報酬がベルヌーイ分布に従うときのThompson Sampling

報酬がベルヌーイ分布に従うときのThompson Samplingの追試をだいぶ前にやったのですが，あまり受けが良くなかったのでブログで供養をしていきます。

問題設定

Thompson Samplingは多腕バンディット問題に対する方策の1つです。多腕バンディット問題については良い解説記事やスライドがたくさんあるので，興味がある人はググってみてください。以下のようなデモもあります*1。

https://arosh.github.io/multi-armed-bandit/

行ったのはChapelleらのNIPS 2011の論文 [1] の中の実験の1つの追試で，報酬が $\lambda=0.5$ のベルヌーイ分布に従うアームが1本， $\lambda=0.4$ のベルヌーイ分布に従うアームが99本あるという設定です。

原理

Thompson Sampling では個々のアーム $i$ を引いたときに当たりが出た回数 $N_\mathrm{S}(i)$ ，外れが出た回数 $N_\mathrm{F}(i)$ を記録しておきます。アームを選ぶときには個々のアームについてベータ分布 $\beta\left(N_\mathrm{S}(i)+1, N_\mathrm{F}(i)+1\right)$ から乱数を発生させ，最も大きい乱数が得られたアームを選択します。

これはベルヌーイ分布のパラメータ $\lambda$ をベイズ推定して $\lambda$ が最大のアームが $i$ である確率をモンテカルロ法で計算するときに乱数サンプルを1個にした場合と考えることができるのですが，詳細は [2] を参照してください。事前分布を一様分布 $\beta(1,1)$ にするかJeffreys分布 $\beta(0.5,0.5)$ にするかは文献によって異なりますが，報酬がベルヌーイ分布に従う場合にはどちらでもあまり関係ないようです。

ソースコード

乱数の種を変更して複数回実験を行いました。以下のプログラムはコマンドラインでシード値を与えるとThompson Samplingに従ってアームを10⁶回引くシミュレーションを行うものです。pure Pythonでは高速化することが難しいのでGoで書きました。乱数ライブラリの leesper/go_rng とプログレスバー表示ライブラリの cheggaaa/pb を使っています。最初に go get してから使ってください。

package main

import (
    "flag"
    "fmt"
    "os"

    "github.com/leesper/go_rng"
    "gopkg.in/cheggaaa/pb.v1"
)

var (
    progress bool
    prior    string
    seed     int64
)

func init() {
    flag.BoolVar(&progress, "progress", false, "")
    flag.StringVar(&prior, "prior", "uniform", "")
    flag.Int64Var(&seed, "seed", 0, "seed value for random generater")
}

type Rng struct {
    bernoulli *rng.BernoulliGenerator
    beta      *rng.BetaGenerator
}

func NewRng(seed int64) *Rng {
    return &Rng{
        bernoulli: rng.NewBernoulliGenerator(seed),
        beta:      rng.NewBetaGenerator(seed),
    }
}

func (rng *Rng) Bernoulli(p float64) bool {
    return rng.bernoulli.Bernoulli_P(p)
}

func (rng *Rng) Beta(alpha, beta float64) float64 {
    return rng.beta.Beta(alpha, beta)
}

func bernoulliBandit(rng *Rng, prior string) ([]float64, error) {
    K := 100
    eps := 0.1
    maxLamb := 0.5
    T := 1000000
    lamb := make([]float64, K)
    lamb[0] = maxLamb
    for i := 1; i < K; i++ {
        lamb[i] = maxLamb - eps
    }
    var alpha, beta float64
    switch prior {
    case "uniform":
        alpha = 1
        beta = 1
    case "jeffreys", "reference":
        alpha = 0.5
        beta = 0.5
    default:
        return nil, fmt.Errorf("Unknown prior: %s", prior)
    }

    var bar *pb.ProgressBar
    if progress {
        bar = pb.New(T)
        bar.Output = os.Stderr
        bar.Start()
    }

    regret := make([]float64, T+1)
    regret[0] = 0
    S := make([]int, K)
    F := make([]int, K)

    for t := 0; t < T; t++ {
        lambEstm := make([]float64, K)
        for i := 0; i < K; i++ {
            lambEstm[i] = rng.Beta(float64(S[i])+alpha, float64(F[i])+beta)
        }
        argMax := 0
        for i := 1; i < K; i++ {
            if lambEstm[i] > lambEstm[argMax] {
                argMax = i
            }
        }
        b := rng.Bernoulli(lamb[argMax])
        if b {
            S[argMax]++
        } else {
            F[argMax]++
        }
        regret[t+1] = regret[t] + lamb[0] - lamb[argMax]
        if progress {
            bar.Increment()
        }
    }
    return regret, nil
}

func main() {
    flag.Parse()
    filename := fmt.Sprintf("%s/%s-seed%03d.txt", prior, prior, seed)
    //計算してからファイル書き込みに失敗したら悲しいので，先にオープンしておく
    w, err := os.Create(filename)
    if err != nil {
        panic(err)
    }
    defer w.Close()
    rng := NewRng(seed)
    regret, err := bernoulliBandit(rng, prior)
    if err != nil {
        panic(err)
    }
    for i, x := range regret[1:] {
        fmt.Fprintln(w, i+1, x)
    }
}

シミュレーションは重たいのでGNU Parallelで並列に実行しました。

kujira16.hateblo.jp

seq 1 200 | parallel --progress "go run chapelle2011.go --seed {} --prior uniform"

リグレットの変化のグラフを描画するためのPythonコードは以下のとおりです。

import seaborn
import scipy
import matplotlib.pyplot as plt
from tqdm import tqdm

seaborn.reset_orig()
seaborn.set_context('talk')
seaborn.set_color_codes()

def load_grid(prefix, st, en, size=1000000):
    A = scipy.empty((size, en - st + 1))
    for i in tqdm(range(st, en + 1)):
        filename = prefix + '{:0>3}'.format(i) + '.txt'
        A[:, i - st] = scipy.loadtxt(filename)[:, 1]
    return A

A = load_grid('uniform/uniform-seed', 1, 200)

label = 'thompson'
ys = scipy.mean(A, axis=1)
del A
xs = scipy.arange(1, ys.shape[0] + 1)
plt.plot(xs, ys, label=label)

ys = scipy.log(xs)
p = 0.5
q = 0.4
coef = (p-q)/(q*scipy.log(q/p)+(1-q)*scipy.log((1-q)/p)) * 99
plt.plot(xs, coef*ys - coef*scipy.log(100), color='k', label='asymptotic bound')

plt.xlabel('plays')
plt.ylabel('regret')
plt.legend(loc='upper left')

plt.xlim(xmin=100)
plt.ylim(0, 10000)
plt.xscale('log')

plt.savefig('chapelle2011.png', dpi=100)