vectorised gae calc

vectorised advantage calculation turns out to be 40x faster, modified from here https://github.com/pmlg/deep-rl-bootcamp/blob/master/lab4/alg_utils.py#L35

async_worker/ppo.py

@@ -94,17 +94,16 @@ class ProximalPolicyOptimization:
 
 
             values.append(actor_net.to_torch_variable(R))
-            A = actor_net.to_torch_variable(torch.zeros((1, 1)))
-            discount = actor_net.to_torch_variable([self.config.discount])
-            gae_tau = actor_net.to_torch_variable([self.config.gae_tau])
-            for i in reversed(range(len(rewards))):
-                R = actor_net.to_torch_variable([[rewards[i]]])
-                ret = R + discount * values[i + 1]
-                A = ret - values[i] + discount * gae_tau * A
-                advantages.append(A.detach())
-                returns.append(ret.detach())
-            advantages = list(reversed(advantages))
-            returns = list(reversed(returns))
+            returns = rewards + self.config.discount * values[1:]
+            deltas = returns - values[:-1]
+            advs = []
+            cum_adv = 0
+            multiplier = self.config.discount * self.config.gae_tau
+            for delta in flip(deltas, 0):
+                cum_adv = cum_adv * multiplier + delta
+                advs.append(cum_adv)
+            advantages = advs[::-1]
+            returns = list(returns)
             replay.feed([states, actions, returns, advantages])
 
         batched_rewards /= batched_episode