はじめに

2017/3/6の強化学習関連論文であるNeural Episodic ControlについてChainerRLのQickStartを参考にしながら、実装に挑戦してみました。

以前より論文内容を実装しての確認を行えるようにならなければと考えており、勉強のためにも挑戦してみることとしました。Qiita初投稿です。
作成したコードはhttps://github.com/ISakony/NEC_chainerrl_CartPole-v0 に置いてあります。
不正確、間違い等のご指摘いただければ幸いです。

環境

python：2.7.13
sklearn: 0.18.1
chainerrl: 0.1.0

概要

NEC(Neural Episodic Control)は深層強化学習（DQNなど)の構造にメモリを持たせたモデルの一種になります。LSTMやDNCのようにエピソードごとの初期化は行いません。モデル構造としては下図のようになっています。

NECのメモリはアクションの数だけ存在し、Key、Q値のセットで保存されています。動作はまず環境からの入力情報をネットワークを通してkey情報(h)に変換します。次に出力されたhをもとにメモリ内の過去hを近傍探索アルゴリズム(論文ではkd-Treeを使用)をもとに探索し、抽出したQ全てを現在のhと過去情報hとの間の距離による重み付けを行った後、総和を取ることでQ値を算出します。この動作は全てのアクションでそれぞれに行われ、最終的に出そろったQ値のうち最大のものを動作するべきアクションとして選択します。

実装

実装はChainerRLのQuickstartガイドをベースにkd-Treeはscikit-learnのものを使用しました。本当はChainerRLのAgentクラスなどを利用しようと思ったのですが、動作を把握しきる時間が取れず、わかるところだけお借りした形になってしまっています。

環境動作の部分

流れとしてはQuickstartのようにQfunctionの定義、CartPole-v0環境の設定、環境のループの順になります。

agent_nec_cp.py

from sklearn.neighbors.kd_tree import KDTree

class QFunction(chainer.Chain):
    def __init__(self, n_actions, n_hidden_channels=50):
        self.q_list = np.ndarray((1,n_actions),dtype=np.float32)#for discreteActionValue
        super(QFunction, self).__init__(
            l0=L.Linear(4, 4),)
            #l1=L.Linear(n_hidden_channels, n_hidden_channels),
            #l2=L.Linear(n_hidden_channels, 2))# 4))

    def __call__(self, x, ma):
        h = F.tanh(self.l0(x))
        #h = F.tanh(self.l1(h))
        #h = F.tanh(self.l2(h))

        #kd_tree
        q_train = [] #for train [variable,variable]
        ind_list = []#for train
        dist_list = [] #for train
        for j in range(len(ma.maq)):#アクションの数だけ実行する
            h_list = ma.mah[j]
            lp = len(h_list)
            leaf_size = lp + (lp/2)

            tree = KDTree(h_list,leaf_size=leaf_size)#各keyメモリごとにデータ構造を形成
            h_ = h.data
            
            if lp < 50:
                k = lp
            else:
                k = 50
            dist, ind = tree.query(h_,k=k)#k=50個のデータ間距離とデータ構造のインデックスを出力

            count = 0
            for ii in ind[0]:
                mahi = np.zeros((1,4),dtype=np.float32)
                mahi[0] = ma.mah[j][ii]
                hi = chainer.Variable(cuda.to_cpu(mahi))
                wi = F.expand_dims(1/(F.batch_l2_norm_squared((h-hi))+0.001),1)

                if count == 0:
                    w = wi
                    maqi = np.zeros((1,1),dtype=np.float32)
                    maqi[0] = ma.maq[j][ii]
                    q = chainer.Variable(cuda.to_cpu(maqi))
                    qq = wi * q
                    count += 1
                else:
                    w += wi
                    maqi = np.zeros((1,1),dtype=np.float32)
                    maqi[0] = ma.maq[j][ii]
                    q = chainer.Variable(cuda.to_cpu(maqi))
                    qq += wi * q
            qq /= w

            q_train.append(qq)
            ind_list.append(ind)
            dist_list.append(dist)
            self.q_list[0][j] = qq.data[0][0]
        qa = chainer.Variable(cuda.to_cpu(self.q_list))
        return chainerrl.action_value.DiscreteActionValue(qa),q_train,ind_list,dist_list,h.data


env = gym.make('CartPole-v0')
print('observation space:', env.observation_space)
print('action space:', env.action_space)
obs = env.reset()

n_actions = env.action_space.n
q_func = QFunction(n_actions)

q_func.to_cpu()

explorer = LinearDecayEpsilonGreedy(1.0,0.1,10**4,env.action_space.sample)

optimizer = chainer.optimizers.RMSpropGraves(lr=0.00025, alpha=0.95, momentum=0.95, eps=0.0001)
optimizer.setup(q_func)

gamma = 0.99

phi = lambda x: x.astype(np.float32, copy=False)

agent = Nec_cp.NEC(q_func, optimizer, gamma, explorer,n_actions,False)


n_episodes = 1000
max_episode_len = 200

train_Flug = True
all_frame = 0

try:
    os.mkdir("./models")
except:
    pass

for i in range(1, n_episodes + 1):
    obs = env.reset()
    reward = 0
    done = False
    R = 0 # return (sum of rewards)
    t = 0 # time step
    while not done and t < max_episode_len:
        # Uncomment to watch the behaviour
        #env.render()
        xx1 = np.zeros((1,4),dtype=np.float32)
        xx1[0] = obs
        xx = chainer.Variable(cuda.to_cpu(xx1))
        action, ind, dist, key = agent.act(xx,i)
        obs, reward, done, _ = env.step(action)

        if train_Flug == True:
            agent.append_memory_and_train(xx1,action,ind,dist,reward,key,done)
        R += reward
        t += 1
        all_frame += 1

まずQfunctionからですが、最初のhを出力する部分はCartPole-v0の入力情報が元々4次元しかないので、簡略化して1回のLinearにtanhの活性化関数を1回通すだけの構造になっています。
次に出力されたkey情報(h)を元にkd-treeによりメモリから近傍のhを探索します。実行部は以下のようになっています。

agent_nec_cp.py

        for j in range(len(ma.maq)):#アクションの数だけ実行する
            h_list = ma.mah[j]
            lp = len(h_list)
            leaf_size = lp + (lp/2)

            tree = KDTree(h_list,leaf_size=leaf_size)#各keyメモリごとにデータ構造を形成
            h_ = h.data
            
            if lp < 50:
                k = lp
            else:
                k = 50
            dist, ind = tree.query(h_,k=k)#k=50個のデータ間距離とデータ構造のインデックスを出力

後述するNECクラスではメモリクラス(Ma)がKeyメモリとQメモリを持っており、型はnumpyarrayにしておくことで、keyメモリ情報を直接kd-Treeの処理に代入します。
kd-Treeは探索元データからデータ構造を形成し、ネットワークからの出力hに近いk=50個のデータのインデックスとデータ間の距離情報を出力します。出力された距離情報、インデックス情報を元に以下の式で重みを計算します。ただし、実装の都合上で分母の重みの総和は最後に割っています。（δは0.001)

重み計算部はchainerのF.batch_l2_norm_squaredを用いることで、逆伝搬可能な二乗ノルムを算出しています。

その他CartPole-v0の動作についてはほぼQuickstartの解説通りです。(act and trainではなく actで次のobsとrewardを得た後、メモリーへの追加とtrainで別れているところだけ異なる。）

NECエージェント部

NECエージェント部は以下のようになっています。Ma(メモリー)、D(リプレイバッファ）、N-stepを管理するクラス、NEC本体のクラスの４つで構成しています。

Nec_cp.py

class Ma_memory(object):
    def __init__(self,n_of_action,limit_n_of_memory):
        self.limit_n_of_memory = limit_n_of_memory
        self.m = 1
        self.maq = self.first_memory_q(n_of_action)
        self.maq_c = self.m
        self.mah = self.first_memory_key(n_of_action)
        self.mah_c = self.m

    def first_memory_q(self,n_of_action):
        mlist = []
        for ii in range(n_of_action):
            ml = np.zeros((self.limit_n_of_memory,1),dtype=np.float32)
            for i in range(self.m):
                mm = np.random.rand(1)
                mm = mm.astype(np.float32)
                ml[0] = mm
            mlist.append(ml)
        return mlist

    def first_memory_key(self,n_of_action):
        mlist = []
        for ii in range(n_of_action):
            ml = np.zeros((self.limit_n_of_memory,4),dtype=np.float32)
            for i in range(self.m):
                mm = np.random.rand(4)
                mm = mm.astype(np.float32)
                ml[0] = mm
            mlist.append(ml)
        return mlist

    def add_memory(self,h,q,greedy_action,dist,ind,alpha):
        if self.mah_c == self.limit_n_of_memory+1:
                print("change")
                self.mah[greedy_action][ind[greedy_action][0][0]] = h[0]
                self.maq[greedy_action][ind[greedy_action][0][0]] = q[0]

        elif dist[greedy_action][0][0] == 0.0:
                print("update",self.maq[greedy_action][ind[greedy_action][0][0]],alpha * (q[0] - self.maq[greedy_action][ind[greedy_action][0][0]]))
                self.maq[greedy_action][ind[greedy_action][0][0]] = self.maq[greedy_action][ind[greedy_action][0][0]] + alpha * (q[0] - self.maq[greedy_action][ind[greedy_action][0][0]])
        else:

            self.mah[greedy_action][self.mah_c-1] = h[0]
            self.mah_c += 1

            self.maq[greedy_action][self.maq_c-1] = q[0]
            self.maq_c += 1

Maクラスはメモリーを保持する役割です。論文の記述とは異なり、keyとQは分けてあります。（kd-Treeにそのまま代入する形式にするため、keyとQは分けました。）ただ追加のタイミングは同じであり、常にセットで操作されます。またnumpyarrayで最初にリストのサイズは固定されているため、中身のデータ量はmah_cおよびmah_q変数が保持しており、限界量に達した場合は最近傍のhと置き換えます(changeの部分)。また既にメモリ内にあるhとまったく同一(distの値が0.0)のものは以下の式で記録されているQ値が更新されます。(updateの部分）

Nec_cp.py

class D_memory(object):
    def __init__(self,limit_n_of_memory):
        self.limit_n_of_memory = limit_n_of_memory
        self.mem = []

    def add_memory(self,state,action,q,ind):
        if len(self.mem) == self.limit_n_of_memory:
            self.mem = self.replace(state,action,q,ind)
        else:
            self.mem.append((state,action,q,ind))

    def replace(self,state,action,q,ind):
        mem_re = []
        for i in range(len(self.mem)):
            if i == len(self.mem)-1:
                mem_re.append((state,action,q,ind))
            elif len(self.mem) != 1 and i == 0:
                pass
            else:
                mem_re.append(self.mem[i+1])
        return mem_re

Dクラスはリプレイバッファーを扱っており、論文の記述通りstateとaction、Nstep-Q値を記憶します。学習時に呼び出され、stateを入力、Nstep-Qをターゲットとして、ネットワークを更新するのに使います。

Nec_cp.py

class N_step_memory(object):
    def __init__(self,gamma,N_horizon,ma_memory,d_memory,alpha):
        self.mem = []
        self.gamma = gamma
        self.N_horizon = N_horizon
        self.ma_memory = ma_memory
        self.d_memory = d_memory
        self.alpha = alpha

    def add_replace_memory(self,state,action,key,reward,t_step,done,q_ontime,dist,ind):

        poplist = []
        for i in range(len(self.mem)):
            self.mem[i][4] = self.mem[i][4] + 1 #n_step +1
            if done == False and self.mem[i][4] != self.N_horizon:
                self.mem[i][3] += (self.gamma ** (self.mem[i][4]-1)) * reward
            elif done == True:
                self.mem[i][3] += (self.gamma ** (self.mem[i][4]-1)) * reward

                t_qn = np.ndarray((1,1),dtype = np.float32)
                t_qn[0] = self.mem[i][3] #qn

                self.ma_memory.add_memory(self.mem[i][2],t_qn,self.mem[i][1],self.mem[i][5],self.mem[i][6],self.alpha)
                self.d_memory.add_memory(self.mem[i][0],self.mem[i][1],t_qn,self.mem[i][6])

                poplist.append(i)

            elif self.mem[i][4] == self.N_horizon:
                self.mem[i][3] += (self.gamma ** (self.mem[i][4]-1)) * q_ontime[0]

                t_qn = np.ndarray((1,1),dtype = np.float32)
                t_qn[0] = self.mem[i][3] #qn

                self.ma_memory.add_memory(self.mem[i][2],t_qn,self.mem[i][1],self.mem[i][5],self.mem[i][6],self.alpha)
                self.d_memory.add_memory(self.mem[i][0],self.mem[i][1],t_qn,self.mem[i][6])

                poplist.append(i)

        for ii in range(len(poplist)):
            self.mem.pop(poplist[ii]-ii)

        if done == True:
            self.mem = []

        if t_step == 1:
            qn = reward
            self.mem.append([state,action,key,qn,t_step,dist,ind])

N-step-memoryクラスについては、勉強不足もあり、N-step学習の扱いはこれでいいのかという部分もありますが、ひとまずエピソードの始めから1stepずつ記憶しておき、それぞれにN-1stepまで割引率込みの報酬を足していき、N-step目に到達したものから、その時のQ値を加算してDクラスに送ります。Nstepに満たないものは割引率込みの報酬和だけとなります。

例：200ステップでゲームが終了した場合。N=100の時、Nstep到達済みstate,action,N-stepQは100組生成され、残り100組はstate,action,割引率込み報酬和のデータになる。

Nec_cp.py

class NEC(object):
    def __init__(self,q_function, optimizer, gamma,explorer,n_actions,retrain):
        self.model = q_function
        self.n_actions = n_actions
        self.N_horizon = 100       
        self.phi=lambda x: x
        self.update_frequency=1
        self.minibatch_size = 16#32
        self.size_of_memory = 100000 #5*10^5

        if retrain == False:
            self.ma_memory = Ma_memory(n_actions,self.size_of_memory)
        else:
            self.ma_memory = Ma_memory(n_actions,self.size_of_memory)
            self.ma_memory.maq = pickle.load(open("Ma_memory_maq_.pickle","rb"))
            self.ma_memory.mah = pickle.load(open("Ma_memory_mah_.pickle","rb"))

        self.d_memory = D_memory(100000)
        self.alpha = 0.1 #learning late
        self.n_step_memory = N_step_memory(gamma,self.N_horizon,self.ma_memory,self.d_memory,self.alpha)
        self.optimizer = optimizer
        self.gamma = gamma #0.99
        self.explorer = explorer
        self.n_step = 0
        self.q_ontime = None #
        self.q_target = None
        self.t = 0
        self.average_loss_all = None
        self.average_loss_c =0
        self.average_loss = None

    def act(self, state, episode):#
        self.n_step += 1
        if self.n_step == self.N_horizon + 1:
            self.n_step = 1

        #action
        with chainer.no_backprop_mode():
            action_value = self.model(state,self.ma_memory)


        greedy_action = cuda.to_cpu(action_value[0].greedy_actions.data)[0]

        if episode < 1:
            action = np.random.randint(0,self.n_actions)
        else:
            action = self.explorer.select_action(
                self.t, lambda: greedy_action, action_value=action_value[0])
        self.t += 1

        #print(action_value[1])
        self.q_ontime = cuda.to_cpu(action_value[1][greedy_action].data)

        self.last_action = action
        ind = action_value[2]
        dist = action_value[3]
        key = cuda.to_cpu(action_value[4])

        return self.last_action,ind,dist,key

    def append_memory_and_train(self,state,action,ind,dist,reward,key,done):
        t_step = 1
        self.n_step_memory.add_replace_memory(state,action,key,reward,t_step,done,self.q_ontime,dist,ind)

        if len(self.d_memory.mem) < self.minibatch_size:
            n_loop = len(self.d_memory.mem)
        else:
            n_loop = self.minibatch_size

        if self.t % 5 == 0:#16
            loss = 0.0
            for ii in range(n_loop):
                rnd = np.random.randint(len(self.d_memory.mem))
                obs, action, t_q, t_ind = self.n_step_memory.d_memory.mem[rnd]
                obs = chainer.Variable(cuda.to_cpu(obs))
                t_q = chainer.Variable(cuda.to_cpu(t_q))
                tav = self.model(obs,self.n_step_memory.ma_memory)
                greedy_action = cuda.to_cpu(tav[0].greedy_actions.data)[0]
                train_q = tav[1][greedy_action]

                loss += mean_squared_error.mean_squared_error(train_q, t_q)
                self.average_loss_all += loss
                self.average_loss_c += 1.0
                self.average_loss = self.average_loss_all / self.average_loss_c


            if n_loop != 0:
                loss /= n_loop
                self.optimizer.zero_grads()
                loss.backward()
                self.optimizer.update()

NECクラスはMaメモリーを100000、Dメモリーを100000に設定。（論文条件通り）atariゲームではないので、初期の動作は1epoch分をランダム動作にする以外はε-greedyにしています。ネットワークの学習は論文記述では16フレーム置きに、32バッチで実行されますが、今回のQFunctionは1バッチの処理しかできないので、5フレーム置きに16バッチではなく16回メモリから呼び出して学習を行う処理にしています。

学習

当初100～200epochではスコアが上がらなかったため、1000epoch程動かしました。

学習結果

学習曲線は下図のようになりました。

400epochを過ぎる当たりまでは停滞するような挙動をしめしており、そこからスコアが伸びていく挙動になっています。（200スコアが最大）

感想

論文内では、AtariゲームにおいてDQNやA3C等よりも早く学習できる（メモリが更新されるほうが効率が良く学習できる。）とされていますが、CartPole-v0を学習させた限りでは環境の状態がある程度メモリに出揃わないと正しいQ値を算出できないように見えます。
chainerRLのquickstartを動かした限りではDQNはCartPole-v0を100epoch～200epochも学習すればスコア200まで到達していた・・・。
（このことがあるので、実装がそもそも致命的に間違えているのではとも考えています。ご指摘いただければ幸いです。）
Atariゲームなら早く学習するのが確認できるかと思いKey(h)を出力するネットワーク構造をCNNに変更し、MsPackman-v0を実行してみましたが、14万フレーム動かすのにGPU利用でも2日間かかってしまい断念しました。その間の挙動は時々高いスコアを出しつつも安定しない挙動となっていました。論文のグラフではミリオンフレームの軸となっていたので100万フレームまでいけばスコアが伸びてくるのかもしれませんが、今回のひとまず動けば実装では確認は難しそうです。

参考

日本語訳
http://qiita.com/miyamotok0105/items/91fd680a4ed322b9a00a
Quickstart
https://github.com/pfnet/chainerrl/blob/master/examples/quickstart/quickstart.ipynb
解説スライド
https://www.slideshare.net/DeepLearningJP2016/dlneural-episodic-controlmodelfree-episodic-control-75250168