Question

我正在尝试优化策略梯度代理的报酬，即使前向传递对所有4种方法都给出相同的结果，但在反向传播期间，所有@ tf.function变体都出了问题。基本方法可以正常工作，但与其他方法相比太慢了。

这是可行的基准实施，但速度较慢：

head

这些是我的优化版本，可以在前向传递中使用，但是在反向传播期间出了点问题：

main

def reward_to_go(self, rewards, gamma=1.0):
    n = len(rewards)
    ret = np.zeros_like(rewards)
    for i in reversed(range(n)):
        ret[i] += rewards[i] + (gamma*ret[i+1] if i+1 < n else 0)
    return ret

@tf.function
def reward_to_go_array(self, rewards, gamma=1.0):
    n = rewards.shape[0]
    ret = tf.TensorArray(tf.float32, size=n, clear_after_read=False)
    for i in tf.range(n-1, -1, -1):
        if i+1 < n:
            ret = ret.write(i, rewards[i] + gamma*ret.read(i+1))
        else:
            ret = ret.write(i, rewards[i])
    return ret.stack()

调用代码如下：

@tf.function
def reward_to_go_array_cond(self, rewards, gamma=1.0):
    n = rewards.shape[0]
    ret = tf.TensorArray(tf.float32, size=n, clear_after_read=False)
    for i in tf.range(n-1, -1, -1):
        ret = tf.cond(tf.less(i+1, n),
                      lambda: ret.write(i, rewards[i] + gamma*ret.read(i+1)),
                      lambda: ret.write(i, rewards[i]))
    return ret.stack()

其中ep_rews是每一步的奖励列表，而self.gamma是恒定的标量。使用reward_to_go可以正常工作，其余的则不能，但是我不知道为什么。

有人可以说明这里出了什么问题吗？我真的没有主意。

策略梯度奖励去和张量流反向传播

0 个答案: