feat: pass shape to reset, transition and termination functions

instadeepai · Jan 16, 2024 · a372c91 · a372c91
1 parent 7b0acb6
commit a372c91
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/matrax/env.py b/matrax/env.py
@@ -99,7 +99,7 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
             agent_obs=agent_obs,
             step_count=state.step_count,
         )
-        timestep = restart(observation=observation)
+        timestep = restart(observation=observation, shape=self.num_agents)
         return state, timestep
 
     def step(
@@ -122,7 +122,7 @@ def compute_reward(
             actions: chex.Array, payoff_matrix_per_agent: chex.Array
         ) -> chex.Array:
             reward_idx = tuple(actions)
-            return payoff_matrix_per_agent[reward_idx]
+            return payoff_matrix_per_agent[reward_idx].astype(float)
 
         rewards = jax.vmap(functools.partial(compute_reward, actions))(
             self.payoff_matrix
@@ -143,10 +143,12 @@ def compute_reward(
 
         timestep = jax.lax.cond(
             done,
-            termination,
-            transition,
-            rewards,
-            next_observation,
+            lambda: termination(
+                reward=rewards, observation=next_observation, shape=self.num_agents
+            ),
+            lambda: transition(
+                reward=rewards, observation=next_observation, shape=self.num_agents
+            ),
         )
 
         # create environment state