sagantaf

メモレベルの技術記事を書くブログ。

強化学習のコーディングをする上で便利なdefaultdictの使い方

強化学習の勉強を進める上で、defaultdictをよく使うので、簡単に使い方を残しておきます。

まずはimportします。

In [1]:
from collections import defaultdict

さて、例えば強化学習させる環境の行動(action)が、UP、DOWN、LEFT、RIGHTを取れるとします。

In [3]:
actions = ['UP', 'DOWN', 'LEFT', 'RIGHT']

この時、行動価値Qは、

状態s1: ([UPの価値], [DOWNの価値], [LEFTの価値], [RIGHTの価値]),
状態s2: ([UPの価値], [DOWNの価値], [LEFTの価値], [RIGHTの価値]),
・・・

と設定したいはずです。しかし、強化学習の問題を解くときには、事前に状態の内容や数が分からないことが多々あります。

そこで、defaultdictの出番になります。

下記のように設定します。

In [4]:
Q = defaultdict(lambda: [0] * len(actions))

defaultdictの引数は、関数をとります。

この関数には、辞書のvalue部分をどう初期化するかを設定します。

そうすることで、Qにkeyをセットするたびに、lambda関数で初期化されたvalueが格納されるようになります。

では、実際に試してみます。

Qの中身を見てみると、最初は空っぽであることがわかります。

In [138]:
Q
Out[138]:
defaultdict(<function __main__.<lambda>>, {})


ここに状態となるkeyを設定します。

In [139]:
Q["state1"]
Out[139]:
[0, 0, 0, 0]
In [140]:
Q["state2"]
Out[140]:
[0, 0, 0, 0]


こうすることで、keyと初期化処理されたvalueのセットで格納されます。

In [141]:
Q
Out[141]:
defaultdict(<function __main__.<lambda>>,
            {'state1': [0, 0, 0, 0], 'state2': [0, 0, 0, 0]})


事前にdictの大きさやkeyを設定しておく必要がないため、後から状態や行動の内容、数が分かる環境の場合は、便利に使えます。