将棋でディープラーニングするその18(報酬に応じた勾配)

前回の日記で、RL policy networkの勾配 $\Delta \rho$ を求める際に、対数尤度の偏微分に報酬に応じた重み（勝敗の報酬 $z_t^i$ から状態価値 $v(s_t^i)$ を引いた値）を掛ける計算の実装が、Chainerでは難しいということを書いた。

Chainerでは損失関数のbackwardを行うと、ミニバッチで1つの勾配が計算されるため、ミニバッチの要素(局面)単位の勾配に $(z_t^i - v(s_t^i))$ を掛ける計算が実装できない。

そこで、ミニバッチの要素1件ずつ、順伝播と逆伝播を行い、計算された勾配を保存しておき、ミニバッチのすべての要素の勾配を計算できたら、保存しておいた勾配に、報酬に応じた重みを掛け合わせて、それらの平均をとることでミニバッチの勾配とする。
最後に、モデルのupdateを行い、パラメータを更新する。

この方法で、報酬に応じた勾配が実装できる。

しかし、ミニバッチの単位で順伝播、逆伝播の計算ができないため、1件ずつ計算すると時間がかかってしまう。
この部分を効率化するにはChainerの内部で行っている勾配計算に手を入れるか、TensorFlowなどでより低レベルな実装を行う必要がある。

また、勾配の計算はGPUからCPUに値を移して計算する必要があるため、メモリ転送で時間のロスが発生する。

実行時間を気にしなければ、以下のように記述することで実装できる。
下記の例は、3層パーセプトロンで、ミニバッチの要素ごとに勾配に重みを掛けている。

import numpy as np
from chainer import cuda, Chain, Function, Variable
from chainer import optimizers
import chainer.functions as F
import chainer.links as L

W1 = [[ 1.21082544, -0.42751756],
      [ 1.35623264, -0.1971387 ],
      [-0.77883673,  0.28367677]]
W2 = [[ 0.08621028, -0.19540818,  0.78203094],
      [ 0.30133799,  1.3698988 , -0.01031571]]

class MyChain(Chain):
    def __init__(self):
        super(MyChain, self).__init__(
            l1=L.Linear(2, 3, initialW=np.array(W1)),
            l2=L.Linear(3, 2, initialW=np.array(W2)),
        )

    def __call__(self, x):
        h = self.l1(x)
        return self.l2(h)

model = MyChain()
model.to_gpu()

optimizer = optimizers.SGD()
optimizer.setup(model)

# print param data
for path, param in model.namedparams():
    print(path)
    print(param.data)
print()

x_data = [[1, 2],
          [3, 4]]
t_data = [0, 1]

# init grads
dic_grads = {}
for path, param in model.namedparams():
    dic_grads[path] = []

for x_elem, t_elem in zip(x_data, t_data):
    x = Variable(cuda.to_gpu(np.array([x_elem], dtype=np.float32)))
    t = Variable(cuda.to_gpu(np.array([t_elem], dtype=np.int32)))

    y = model(x)

    model.cleargrads()
    loss = F.softmax_cross_entropy(y, t)
    loss.backward()

    # save grad
    for path, param in model.namedparams():
        dic_grads[path].append(cuda.to_cpu(param.grad))

# manipulate grad
z_data = [1.0, 0.5]
z = np.array(z_data, dtype=np.float32)
for path, param in model.namedparams():
    grads = np.array(dic_grads[path], dtype=np.float32)
    grad = np.tensordot(z, grads, axes=1)
    grad /= len(z_data)
    param.grad = cuda.to_gpu(grad)

optimizer.update()

# print param data and grad
for path, param in model.namedparams():
    print(path)
    print(cuda.to_cpu(param.data))
    print(cuda.to_cpu(param.grad))