Update RoboCup to work with the latest Mava version.

ch33nchan · Nov 24, 2021 · aca4ce8 · aca4ce8
1 parent fa3a5ea
commit aca4ce8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mava/wrappers/robocup.py b/mava/wrappers/robocup.py
@@ -43,7 +43,7 @@ def reset(self) -> Tuple[dm_env.TimeStep, np.array]:
         proc_obs = self._proc_robocup_obs(observations=raw_obs, done=False)
         proccessed_state = self._proc_robocup_state(state, proc_obs)
         timestep = dm_env.restart(proc_obs)
-        return timestep, {"env_state": proccessed_state}
+        return timestep, {"s_t": proccessed_state}
 
     def step(self, nn_actions: types.NestedArray) -> Tuple[dm_env.TimeStep, np.array]:
         """Steps the environment."""
@@ -71,7 +71,7 @@ def step(self, nn_actions: types.NestedArray) -> Tuple[dm_env.TimeStep, np.array
                 discount=self._discount,
                 step_type=self._step_type,
             ),
-            {"env_state": proccessed_state},
+            {"s_t": proccessed_state},
         )
 
     @property