cleanup

Author: lucidrains

improving_transformers_world_model/agent.py

@@ -101,13 +101,9 @@ def calc_target_and_gae(
     delta = rewards + gamma * values_next * masks - values
     gates = gamma * lam * masks
 
-    gates, delta = gates[..., :, None], delta[..., :, None]
-
     scan = AssocScan(reverse = True, use_accelerated = use_accelerated)
     gae = scan(gates, delta)
 
-    gae = gae[..., :, 0]
-
     returns = gae + values
 
     gae, returns = tuple(inverse_pack(t) for t in (gae, returns))

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "improving-transformers-world-model"
-version = "0.0.52"
+version = "0.0.53"
 description = "Improving Transformers World Model for RL"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -26,7 +26,7 @@ classifiers=[
 
 dependencies = [
     "adam-atan2-pytorch>=0.1.18",
-    "assoc-scan",
+    "assoc-scan>=0.0.2",
     "einx>=0.3.0",
     "einops>=0.8.0",
     "ema-pytorch>=0.7.7",