スキナー箱のQ学習 - 化学系エンジニアがAIを学ぶ

はじめに

ネズミがボタンを押して餌を得るという問題について強化学習（Q学習）を適用してみる。本記事の内容に関して下記書籍を参考にした。

牧野浩二、西崎博光, 『Pythonによる深層強化学習入門』, オーム社 (2018)

ネズミの学習問題

対象とする問題は下記の通り。

箱の中にネズミが1匹いる
箱の中にネズミ餌のディスペンサーが1つある
ディスペンサーには「電源ボタン」と「餌ボタン」がある
「電源ボタン」をネズミが押すたびに電源のON/OFFが切り替わる
ディスペンサーの電源がONのときにネズミが「餌ボタン」を押すと餌が出てくる

ネズミは餌を得る手順を学習できるか。

ディスペンサーのモデル

まず、電源のON/OFF状態を保持し、ON状態で「餌ボタン」が押されると餌（報酬）を返すクラスを作成する。

import numpy

class Dispenser(object):
    def __init__(self, init_state):
        """
        初期のON/OFF状態を設定する
        init_state: 0->電源OFF、1->電源ON
        """
        self._state  = init_state
        self._button = self.powerbutton(init_state)

    def powerbutton(self, init_state):
        """
        ON/OFFを切り替えるgenerator
        0: 電源OFF、1: 電源ON を返す。内部的には-1, 1で切り替えている
        """
        state = -1 if init_state==0 else init_state
        while True:
            state = -1 * state
            yield numpy.maximum(0, state)

    def push_powerbutton(self):
        self._state = self._button.__next__()

    def step(self, action):
        """
        action: 0->電源ボタンを押す 1->餌ボタンを押す
        """
        if action == 0:
            self.push_powerbutton()
            reward = 0
        else:
            reward = 1 if self._state==1 else 0
        return self._state, reward

挙動としては下記。電源ボタンを押すとON/OFFが切り替わり、ONのとき餌ボタンを押すと餌(報酬 1)が得られる。具体的には(状態, 報酬)のタプルが返る。これだけ。

>>> env = Dispenser(0) # 電源OFFを初期状態としてディスペンサーオブジェクトを作成
>>> env.step(0)  # 0:電源ボタンを押す
(1, 0) # 状態:電源ON, 報酬:0 が返る

>>> env.step(1) # 1: 餌ボタンを押す
(1, 1) # 状態:電源ON, 報酬:1が返る

Q値

各状態および行動に対するQ値の表を次の通り決める。

Q(0, 0): 状態電源OFF、行動電源ボタンを押す

Q(0, 1): 状態電源OFF、行動餌ボタンを押す

Q(1, 0): 状態電源ON、行動電源ボタンを押す

Q(1, 1): 状態電源ON、行動餌ボタンを押す

各Q値は下記の通り配列で保持することとする。初期値は全て0。

qtable = numpy.zeros((2, 2))

Q値に従って次の行動を決める

Q値および現在の状態から次に取る行動を決める。基本的にはQ値が最大の行動を取るが、時折ランダムな行動を取る $\epsilon$ - グリーディー法を取り入れることとする。

def get_action(qtable, state, episode):
    """
    現在の状態とQ値から次の行動を決める
    episode回数が増えるごとにランダム行動を取らなくする
    """
    ## ε- greedy法 行動を重ねる(episodeが増える)毎にランダム行動を取らなくする
    epsilon = 0.5 * (1 / (episode + 1))
    if epsilon <= numpy.random.uniform(0, 1):
        ## まずQ値が最大の要素を確認(必ずしも1つでないので)
        ## 最大値がTrue、それ以外はFalseとなっているリストを得る
        qmax_bool_list = qtable[state]==qtable[state].max()
        ## Qが最大値の要素のインデックスを得る(複数あるときはその中からランダムに選ぶことになる)
        choices = numpy.where(qmax_bool_list)[0] 
    else:
        ## ランダム行動。全要素のインデックスを得る
        choices = numpy.arange(len(qtable[state]))
    return numpy.random.choice(choices) # 候補の中から1つランダムに選ぶ

Q値を更新する

状態、行動および行動を取った結果の報酬から下記式に従いQ値を更新する。 $s_t$ は状態、 $a_t$ は行動、 $r$ は報酬、 $\text{max}Q$ は次の状態での最大のQ値。 $\alpha$ は学習率、 $\gamma$ は割引率と呼ばれる。

$\displaystyle Q(s_t, a_t) \gets (1 - \alpha) Q(s_t, a_t) + \alpha (r + \gamma \text{max} Q)$

def update_qtable(qtable, state, action, next_state, reward):
    gamma = 0.9
    alpha = 0.5
    next_qmax = max(qtable[next_state])
    qtable[state, action] = (1 - alpha) * qtable[state, action] +\
                            alpha * (reward + gamma * next_qmax)
    return qtable

学習する

1episodeあたり5回行動することとして、10episode学習させる。各episodeの初期のディスペンサー電源ON/OFFはランダムに決める。

MAX_STEP_NUMBER = 5 # 1episodeあたり5回行動することとする
MAX_EPISODE_NUMBER = 10 # episodeは10までとする

for episode in range(MAX_EPISODE_NUMBER):
    total_reward = 0
    state = numpy.random.choice([0, 1]) # ディスペンサー電源初期状態はランダム
    env = Dispenser(state)
    
    for step in range(MAX_STEP_NUMBER):
        action = get_action(qtable, state, episode) # q値、状態から行動を決める
        next_state, reward = env.step(action) # 行動を取った結果を得る
        print(state, action, reward)
        total_reward += reward
        qtable = update_qtable(qtable, state, action, next_state, reward) # Q値の更新
        state = next_state

    print('episode : %d total reward %d' %(episode+1, total_reward))
    print(qtable)

学習結果

初期が電源OFFならtotal rewardは最大4、ONなら最大5となるが、episode 9, 10ではいずれもrewardは最大となっている。初期ON/OFFがランダムでも正しく行動が取られている。

>>>
0 1 0
0 0 0
1 1 1
1 1 1
1 1 1
episode : 1 total reward 3
[[ 0.       0.     ]
 [ 0.       1.42625]]
0 0 0
1 0 0
0 0 0
1 1 1
1 0 0
episode : 2 total reward 1
[[ 0.96271875  0.        ]
 [ 0.57763125  1.8549375 ]]
.
.
.
0 0 0
1 1 1
1 1 1
1 1 1
1 1 1
episode : 9 total reward 4
[[ 6.00931629  1.23142328]
 [ 1.25417229  7.85361236]]
1 1 1
1 1 1
1 1 1
1 1 1
1 1 1
episode : 10 total reward 5
[[ 6.00931629  1.23142328]  # 電源OFF時のQ値、電源ボタンを押すQ値が大きい
 [ 1.25417229  8.33916616]] # 電源ON時のQ値、餌ボタンを押すQ値が大きい

全コード

import numpy

class Dispenser(object):
    def __init__(self, init_state):
        """
        初期のON/OFF状態を設定する
        init_state: 0->電源OFF、1->電源ON
        """
        self._state  = init_state
        self._button = self.powerbutton(init_state)

    def powerbutton(self, init_state):
        """
        ON/OFFを切り替えるgenerator
        0: 電源OFF、1: 電源ON を返す。内部的には-1, 1で切り替えている
        """
        state = -1 if init_state==0 else init_state
        while True:
            state = -1 * state
            yield numpy.maximum(0, state)

    def push_powerbutton(self):
        self._state = self._button.__next__()

    def step(self, action):
        """
        action: 0->電源ボタンを押す 1->餌ボタンを押す
        """
        if action == 0:
            self.push_powerbutton()
            reward = 0
        else:
            reward = 1 if self._state==1 else 0
        return self._state, reward

def get_action(qtable, state, episode):
    ## ε- greedy法 行動を重ねる(episodeが増える)毎にランダム行動を取らなくする
    epsilon = 0.5 * (1 / (episode + 1))
    if epsilon <= numpy.random.uniform(0, 1):
        ## Q値に従った行動。Q値が最大の行動を取る
        ## まずQ値が最大の要素を確認(必ずしも1つでないので)
        qmax_bool_list = qtable[state]==qtable[state].max()
        # 複数あるときはその中からランダムに選ぶ
        choices = numpy.where(qmax_bool_list)[0] 
    else:
        ## ランダム行動
        choices = numpy.arange(len(qtable[state]))
    return numpy.random.choice(choices)

def update_qtable(qtable, state, action, next_state, reward):
    gamma = 0.9
    alpha = 0.5
    next_qmax = max(qtable[next_state])
    qtable[state, action] = (1 - alpha) * qtable[state, action] +\
                            alpha * (reward + gamma * next_qmax)
    return qtable

qtable = numpy.zeros((2, 2))
MAX_STEP_NUMBER = 5 # 1episodeあたり5回行動することとする
MAX_EPISODE_NUMBER = 10 # episodeは10までとする

for episode in range(MAX_EPISODE_NUMBER):
    total_reward = 0
    state = numpy.random.choice([0, 1]) # ディスペンサー電源初期状態はランダム
    env = Dispenser(state)
    
    for step in range(MAX_STEP_NUMBER):
        action = get_action(qtable, state, episode) # q値、状態から行動を決める
        next_state, reward = env.step(action) # 行動を取った結果を得る
        print(state, action, reward)
        total_reward += reward
        qtable = update_qtable(qtable, state, action, next_state, reward) # Q値の更新
        state = next_state

    print('episode : %d total reward %d' %(episode+1, total_reward))
    print(qtable)