chore: simply ppo gae code (#334)

vwxyzjn · Dec 28, 2022 · d67ae0c · d67ae0c · vercel · Dec 28, 2022
1 parent ad19f23
commit d67ae0c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/cleanrl/ppo_atari_envpool_xla_jax_scan.py b/cleanrl/ppo_atari_envpool_xla_jax_scan.py
@@ -323,8 +323,10 @@ def compute_gae(
         _, advantages = jax.lax.scan(
             compute_gae_once, advantages, (dones[1:], values[1:], values[:-1], storage.rewards), reverse=True
         )
-        storage = storage.replace(advantages=advantages)
-        storage = storage.replace(returns=storage.advantages + storage.values)
+        storage = storage.replace(
+            advantages=advantages,
+            returns=storage.advantages + storage.values,
+        )
         return storage
 
     def ppo_loss(params, x, a, logp, mb_advantages, mb_returns):

diff --git a/tests/test_jax_compute_gae.py b/tests/test_jax_compute_gae.py
@@ -39,8 +39,10 @@ def compute_gae_scan(
         _, advantages = jax.lax.scan(
             compute_gae_once_fn, advantages, (dones[1:], values[1:], values[:-1], storage.rewards), reverse=True
         )
-        storage = storage.replace(advantages=advantages)
-        storage = storage.replace(returns=storage.advantages + storage.values)
+        storage = storage.replace(
+            advantages=advantages,
+            returns=storage.advantages + storage.values,
+        )
         return storage
 
     def compute_gae_python_loop(