Add test cases (#339)

* Add test cases

* fix test cases

Author: vwxyzjn

.github/workflows/tests.yaml

@@ -37,6 +37,9 @@ jobs:
       - name: Run core tests with jax
         if: runner.os == 'Linux' || runner.os == 'macOS'
         run: poetry run pytest tests/test_classic_control_jax.py
+      - name: Run gae tests with jax
+        if: runner.os == 'Linux' || runner.os == 'macOS'
+        run: poetry run pytest tests/test_jax_compute_gae.py
       - name: Install tuner dependencies
         run: poetry install --with optuna
       - name: Run tuner tests

tests/test_jax_compute_gae.py

@@ -41,7 +41,7 @@ def compute_gae_scan(
         )
         storage = storage.replace(
             advantages=advantages,
-            returns=storage.advantages + storage.values,
+            returns=advantages + storage.values,
         )
         return storage