Fix current shared next states usage

Toni-SM · Aug 14, 2024 · a243cae · a243cae
1 parent f6198e6
commit a243cae
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/skrl/multi_agents/jax/mappo/mappo.py b/skrl/multi_agents/jax/mappo/mappo.py
@@ -497,7 +497,7 @@ def _update(self, timestep: int, timesteps: int) -> None:
 
             # compute returns and advantages
             value.training = False
-            last_values, _, _ = value.act({"states": self._shared_state_preprocessor[uid](self._current_shared_next_states[uid])}, role="value")  # TODO: .float()
+            last_values, _, _ = value.act({"states": self._shared_state_preprocessor[uid](self._current_shared_next_states)}, role="value")  # TODO: .float()
             value.training = True
             if not self._jax:  # numpy backend
                 last_values = jax.device_get(last_values)