fix: matrax env metrics when using global state (#1174)

Author: sash-a

mava/wrappers/matrax.py

@@ -43,6 +43,8 @@ def modify_timestep(
         self, timestep: TimeStep
     ) -> TimeStep[Union[Observation, ObservationGlobalState]]:
         """Modify the timestep for `step` and `reset`."""
+        metrics: Dict[str, Any] = {"env_metrics": {}}
+
         obs_data = {
             "agents_view": timestep.observation.agent_obs,
             "action_mask": self.action_mask,
@@ -52,9 +54,8 @@ def modify_timestep(
             global_state = jnp.concatenate(timestep.observation.agent_obs, axis=0)
             global_state = jnp.tile(global_state, (self.num_agents, 1))
             obs_data["global_state"] = global_state
-            return timestep.replace(observation=ObservationGlobalState(**obs_data))
+            return timestep.replace(observation=ObservationGlobalState(**obs_data), extras=metrics)
 
-        metrics: Dict[str, Any] = {"env_metrics": {}}
         return timestep.replace(observation=Observation(**obs_data), extras=metrics)
 
     def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep]: