ChainerRL を用いた NoisyNet-DQN の実装～深層強化学習における探索手法

近年 AlphaGo の活躍で注目されている深層強化学習ですが、この進歩により人工知能の自動制御が世の中に革命を起こすかもしれないと期待されています (実際にどこまでいけるかは賛否両論ですが、少なくとも僕は期待しています)。

やはり実際に使ってみると、適切な試行錯誤でいかにして最適な動作を学習させるかのプロセスが非常に難しい…。

今回の NoisyNet は、ネットワークそのものにノイズを乗せて、効率的な探索に (一部) 成功したというもの。ノイズは最近の流行りなのだろうか…
https://mochitam.com/2017/06/30/2017-06-30-181850/

今回は DeepMind の NoisyNet-DQN を試してみようと思うのですが、OpenAI も似た論文を出しているということで、こんな世界最強クラスの熾烈な開発競争を見せられては自分の戦う場所は無さそう…

強化学習のおける探索の実現方法として、Q関数や方策のモデルパラメータにノイズを入れる方法がほぼ同時に提案された。エピソード中連続して違う空間を探せるので効率的。従来手法を上回る性能 https://t.co/b83rGNRHfI https://t.co/lfrTasPJ34
— 岡野原大輔 (@hillbig) 2017年7月4日

理論

以下の DeepMind の Fortunato et al. 2017 では、NoisyNet で Atari をやってみたら $\epsilon$ -greedy で探索させるよりも平均パフォーマンスが向上したとのこと。

[1706.10295] Noisy Networks for Exploration

深層強化学習の基本は非常に簡単なので、今さら僕が何か語れることも無いのですが簡単におさらいしておきます。

qiita.com

DQN

ある時刻 $t$ にエージェントは環境 $s$ を観測し、 $u$ という行動をすると、時刻 $t+1$ に再度環境 $s'$ が観測され、報酬 $r$ を獲得します。

時刻 $t =$ 1~T の間に獲得した合計報酬 $R = \sum^{T}_{t=1} r_t$ が最も多くなるような $u$ を選ぶにはどうしたらいいかを考えます。

Q学習では、現在から考えうる最適動作を取り続けた結果、時刻 T までに得られる合計報酬を $Q$ とすると、以下の漸化式が成り立つわけです。

${ \displaystyle Q(s, u) = r + \gamma \max_{u'} Q(s', u') }$

$\gamma$ が無ければ単純に $\sum r_t$ となっていることが分かります。

$\gamma$ は減衰率です。めちゃくちゃ長い時間の動作だと、 $\sum r_t$ が発散してしまうので、未来の影響を適切に減衰していい感じにするのが通例らしいです。

とはいえ $Q$ は当然未知です。ここで $Q$ をディープラーニングで表現される魔法のような関数だとするのが深層強化学習です (これが言いたかっただけなのに長かった)。最初はあてずっぽうで、何回も試行錯誤しながらいい感じの $Q$ を学習できれば人工知能による自動制御の完成というわけです。

$Q$ が $\theta$ でパラメトライズされるとすると、今予想している $Q$ と、実際に観測された報酬との差が損失関数 $L(\theta)$ になる。

${ \displaystyle L(\theta) \sim \left( r+\gamma\max_{u'} Q(s', u'; \theta')-Q(s,u;\theta) \right)^{2} }$

ここで、 $\theta'$ には適切なタイミングで現在の $\theta$ が代入される。報酬は未来から逆伝搬してくるので、一時的に $\theta'$ を止めて擬似的な教師有り学習することで、学習を安定させるのが一般的とされる。

NoisyNet-DQN

ここからが NoisyNet の出番であり、Q 関数のもつ重みとバイアスを $\mathcal{N} (\mu, \Sigma)$ という空間を漂わせる。つまり $\epsilon$ をガウシアンノイズとすると、

$\theta = \mu + \Sigma \odot \epsilon$

となる。こうして $\epsilon$ -greedy のような探索を使わずとも、ネットワークそのものが最適探索をしてしまうらしい。

新たに $\zeta \equiv (\mu, \Sigma)$ でパラメトライズされた損失関数に書き直すと、 $\epsilon = \mathcal{N} (0, 1)$ からサンプリングされた $\xi$ に対して以下のように勾配が計算される。

$\nabla \bar{L} (\zeta) \sim \nabla L(\mu + \Sigma \odot \xi)$

実装

実際に NoisyNet-DQN を実装しようと思うと、Linear 関数の中でノイズを発生させるように書き直すだけで従来の DQN がそのまま使えそう。

Links の Linear を継承して、以下のような NoisyLinear 関数を書いてみた (めちゃくちゃ汚いけど論文を再現するための僕の限界であった)。初期値関係は後述。

class NoisyLinear(L.Linear):
    def _initialize_params(self, in_size):
        super(NoisyLinear, self)._initialize_params(in_size)
        self.in_size = in_size
        factor = 1.0 / np.sqrt(in_size)
        
        ini_w = np.ones_like(self.W) * factor
        ws_initializer = initializers._get_initializer(ini_w)
        self.w_sigma = variable.Parameter(ws_initializer)
        self.w_sigma.initialize((self.out_size, in_size))
        
        ini_b = np.zeros(self.out_size)
        bs_initializer = initializers._get_initializer(ini_b)
        self.b_sigma = variable.Parameter(bs_initializer)
        self.b_sigma.initialize((self.out_size))
        
    def __call__(self, x, noise=True, test=False):
        if self.W.data is None:
            self._initialize_params(x.size // x.shape[0])
        
        if not test and noise:
            xp = self.xp
            e_i = xp.random.normal(size=(self.in_size)).astype(xp.float32)
            e_j = xp.random.normal(size=(self.out_size)).astype(xp.float32)
            e_w = xp.outer(e_j,e_i)
            e_w = variable.Variable(xp.sign(e_w)*xp.sqrt(xp.abs(e_w)))
            e_b = variable.Variable(xp.sign(e_j)*xp.sqrt(xp.abs(e_j)))
            
            W = self.W + self.w_sigma * e_w
            b = self.b + self.b_sigma * e_b
        else:
            W = self.W
            b = self.b
        
        return F.connection.linear.linear(x, W, b)

DQN をスクラッチで書くのはそんなに大変でも無いけど、せっかく PFN さんが公開している Chainer RL の DQN をそのまま使わせてもらうのがパフォーマンスは安定するだろうと期待。

深層強化学習ライブラリChainerRL | Preferred Research

とりあえずのテストのために OpenAI gym の CartPole-v0 あたりを解かせてみる。

gym.openai.com

マニュアルを参照させて頂くと、以下のように gym と Chainer RL を呼べばとりあえず DQN で CartPole-v0 を解かせることが出来るらしい。

import numpy as np

import chainer 
import chainer.functions as F
import chainerrl
import gym

class MTNNet(chainer.Chain):
    def __init__(self, obs_size, n_actions, n_units, noise=True):
        super(MTNNet, self).__init__()
        self.noise = noise
        self.unit = n_units
        self.test = test
        
        with self.init_scope():
            self.lin1 = NoisyLinear(obs_size,self.unit, 
                                    initialW=np.random.uniform(-1.0/np.sqrt(obs_size), 1.0/np.sqrt(obs_size), size=(self.unit,obs_size)))
            self.lin2 = NoisyLinear(self.unit,self.unit,
                                    initialW=np.random.uniform(-1.0/np.sqrt(self.unit), 1.0/np.sqrt(self.unit), size=(self.unit,self.unit)))
            self.lin3 = NoisyLinear(self.unit,n_actions,
                                    initialW=np.random.uniform(-1.0/np.sqrt(self.unit), 1.0/np.sqrt(self.unit), size=(n_actions,self.unit)))
            
    def __call__(self, x):
        h = self.lin1(x, noise=self.noise, test=self.test)
        h = F.leaky_relu(h)
        h = self.lin2(h, noise=self.noise, test=self.test)
        h = F.leaky_relu(h)
        h = self.lin3(h, noise=self.noise, test=self.test)
        a = chainerrl.action_value.DiscreteActionValue(h)
        return a

env = gym.make('CartPole-v0')
obs_size = env.observation_space.shape[0]
n_actions = env.action_space.n

model = MTNNet(obs_size, n_actions, 100, self.noise)
optimizer = chainer.optimizers.Adam()
optimizer.setup(model)

batchsize = 32
gamma = 0.9

explorer = chainerrl.explorers.ConstantEpsilonGreedy(
                epsilon=0.0,
                random_action_func=env.action_space.sample)
        
replay_buffer = chainerrl.replay_buffer.ReplayBuffer(capacity=10**6)

phi = lambda x: x.astype(np.float32, copy=False)
        
agent = chainerrl.agents.DQN(
            model, optimizer, replay_buffer, gamma, explorer, 
            phi=phi,
            update_interval=1,
            replay_start_size=batchsize,
            minibatch_size=batchsize,
            target_update_interval=batchsize*10
            )

n_episodes = 500
max_episode_len = 200

for i in range(1, n_episodes + 1):
    obs = env.reset()
    reward = 0
    done = False
    R = 0
    t = 0
    while not done and t < max_episode_len:
        action = agent.act_and_train(obs, reward)
        obs, reward, done, _ = env.step(action)
        R += reward
        t += 1

初期値系を論文準拠としたので、めんどくさいコーディングになっていますが、ここではそのまま引用するのが分かり良いかと。ノイズについては Factorised Gaussian noise を採用してます。