Политический градиент вознаграждений до go и обратное распространение тензорного потока - PullRequest
Новая
       20

Политический градиент вознаграждений до go и обратное распространение тензорного потока

0 голосов
/ 09 февраля 2020

Я пытаюсь оптимизировать вознаграждение до go для моего агента градиента политики, и хотя прямой проход дает одинаковые результаты со всеми 4 методами, во время обратного распространения что-то идет не так со всеми вариантами @ tf.function , Метод basi c работает правильно, но он слишком медленный по сравнению с другими.

Вот базовая реализация, которая работает, но медленная: 

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

Это мои оптимизированные версии, которые работают в прямом проходе, но что-то идет не так во время backprop: 

@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()

@tf.function
def reward_to_go_matmul(self, rewards, gamma=1.0):
    # initialize coefficients (gamma cumprod to the number of steps)
    coeffs = tf.math.cumprod(tf.fill([self.input_length], tf.constant(gamma)),
                             exclusive=True)
    # create upper triangular matrix with coefficients
    coeff_m = tf.zeros([self.input_length, self.input_length])
    for i in range(self.input_length):
        coeff_m = matrix_set_diag_v2(coeff_m, tf.fill([self.input_length-i], coeffs[i]), k=i)
    # calculate reward to go
    rtg = tf.matmul(coeff_m, tf.squeeze(rewards))
    return tf.expand_dims(rtg, axis=-1)

Код вызова выглядит так: 

batch_rtg = self.reward_to_go_array_cond(tf.stack(ep_rews), tf.constant(self.gamma))
batch_rtg = self.reward_to_go_array(tf.stack(ep_rews), tf.constant(self.gamma))
batch_rtg = self.reward_to_go_matmul(tf.stack(ep_rews), self.gamma)
batch_rtg = self.reward_to_go(ep_rews, self.gamma)

Где ep_rews список наград за каждый шаг и self.gamma - это постоянный скаляр. Использование reward_to_ go работает должным образом, остальные нет, но я не могу понять, почему ..

Может кто-нибудь, пожалуйста, пролить свет на то, что здесь идет не так? У меня действительно нет идей.

...