feat: pass randomness keys to environment steps

docs/guides/advanced_usage.md

@@ -42,3 +42,8 @@ rollout = jax.vmap(run_n_steps, in_axes=(0, 0, None))(state, keys, rollout_lengt
 # Shape and type of given rollout:
 # TimeStep(step_type=(7, 5), reward=(7, 5), discount=(7, 5), observation=(7, 5, 6, 6, 5), extras=None)
 ```
+
+For stochastic environments, `env.step(state, action, key)` uses the supplied key instead of
+`state.key`, with the same splitting and returned-state key behavior. Omitting `key` preserves
+the existing behavior. Deterministic environments accept and may ignore the key.
+`AutoResetWrapper` continues to derive reset randomness from the returned state key.

jumanji/env.py

@@ -67,12 +67,19 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
         """
 
     @abc.abstractmethod
-    def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self,
+        state: State,
+        action: chex.Array,
+        key: chex.PRNGKey | None = None,
+    ) -> Tuple[State, TimeStep[Observation]]:
         """Run one timestep of the environment's dynamics.
 
         Args:
             state: State object containing the dynamics of the environment.
             action: Array containing the action to take.
+            key: optional random key used in place of state.key for stochastic steps,
+                including the returned state key. Deterministic environments may ignore it.
 
         Returns:
             state: State object corresponding to the next state of the environment,

jumanji/environments/logic/game_2048/env.py

@@ -163,12 +163,15 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
 
         return state, timestep
 
-    def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Array, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """Updates the environment state after the agent takes an action.
 
         Args:
             state: the current state of the environment.
             action: the action taken by the agent.
+            key: random key used to sample the next tile.
 
         Returns:
             state: the new state of the environment.
@@ -177,8 +180,8 @@ def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observ
         # Take the action in the environment: Up, Right, Down, Left.
         updated_board, reward = move(state.board, action)
 
-        # Generate new key.
-        random_cell_key, new_state_key = jax.random.split(state.key)
+        key = state.key if key is None else key
+        random_cell_key, new_state_key = jax.random.split(key)
 
         # Update the state of the board by adding a new random cell.
         updated_board = jax.lax.cond(

jumanji/environments/logic/game_2048/env_test.py

@@ -103,6 +103,27 @@ def test_game_2048__step_invalid(game_2048: Game2048) -> None:
     assert jnp.array_equal(state.score, new_state.score)
 
 
+def test_game_2048__step_uses_supplied_key(game_2048: Game2048, board: Board) -> None:
+    """The supplied key controls the stochastic successor state."""
+    state = State(
+        board=board,
+        step_count=jnp.array(0),
+        action_mask=game_2048._get_action_mask(board),
+        score=jnp.array(0),
+        key=jax.random.PRNGKey(42),
+    )
+    action = jnp.array(1)
+    step_key = jax.random.PRNGKey(0)
+
+    next_state, _ = game_2048.step(state, action, step_key)
+    state_with_different_key = state.replace(key=jax.random.PRNGKey(43))
+    independent_state, _ = game_2048.step(state_with_different_key, action, step_key)
+    resampled_state, _ = game_2048.step(state, action, jax.random.PRNGKey(1))
+
+    chex.assert_trees_all_equal(next_state, independent_state)
+    assert not jnp.array_equal(next_state.board, resampled_state.board)
+
+
 def test_game_2048__step_action_mask(game_2048: Game2048) -> None:
     """Verify that the action mask returned from `step` is correct."""
     state = State(

jumanji/environments/logic/graph_coloring/env.py

@@ -135,7 +135,9 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
 
         return state, timestep
 
-    def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Array, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """Updates the environment state after the agent takes an action.
 
         Specifically, this function allows the agent to choose

jumanji/environments/logic/minesweeper/env.py

@@ -147,7 +147,9 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
         timestep = restart(observation=observation)
         return state, timestep
 
-    def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Array, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """Run one timestep of the environment's dynamics.
 
         Args:

jumanji/environments/logic/rubiks_cube/env.py

@@ -135,7 +135,9 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
         timestep = restart(observation=observation)
         return state, timestep
 
-    def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Array, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """Run one timestep of the environment's dynamics.
 
         Args:

jumanji/environments/logic/sliding_tile_puzzle/env.py

@@ -109,7 +109,9 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
         timestep = restart(observation=obs, extras=self._get_extras(state))
         return state, timestep
 
-    def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Array, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """Updates the environment state after the agent takes an action."""
         (updated_puzzle, updated_empty_tile_position) = self._move_empty_tile(
             state.puzzle, state.empty_tile_position, action

jumanji/environments/logic/sudoku/env.py

@@ -99,7 +99,9 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
         timestep = restart(observation=obs)
         return state, timestep
 
-    def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Array, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         # check if action is valid
         invalid = ~state.action_mask[tuple(action)]
         updated_board = apply_action(action=action, board=state.board)

jumanji/environments/packing/bin_pack/env.py

@@ -285,7 +285,9 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
 
         return state, timestep
 
-    def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Array, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """Run one timestep of the environment's dynamics. If the action is invalid, the state
         is not updated, i.e. the action is not taken, and the episode terminates.
 

jumanji/environments/packing/flat_pack/env.py

@@ -157,7 +157,9 @@ def reset(
 
         return grid_state, timestep
 
-    def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Array, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """Steps the environment.
 
         Args:

jumanji/environments/packing/job_shop/env.py

@@ -169,7 +169,9 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
 
         return state, timestep
 
-    def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Array, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """Updates the status of all machines, the status of the operations, and increments the
         time step. It updates the environment state and the timestep (which contains the new
         observation). It calculates the reward based on the three terminal conditions:

jumanji/environments/packing/knapsack/env.py

@@ -138,7 +138,9 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
         timestep = restart(observation=self.observe(state))
         return state, timestep
 
-    def step(self, state: State, action: chex.Numeric) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Numeric, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """Run one timestep of the environment's dynamics.
 
         Args:

jumanji/environments/packing/tetris/env.py

@@ -156,20 +156,24 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
         timestep = restart(observation=observation)
         return state, timestep
 
-    def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Array, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """Run one timestep of the environment's dynamics.
 
         Args:
             state: `State` object containing the dynamics of the environment.
             action: `chex.Array` containing the rotation_index and x_position of the tetromino.
+            key: random key used to sample the next tetromino.
 
         Returns:
             next_state: `State` corresponding to the next state of the environment,
             next_timestep: `TimeStep` corresponding to the timestep returned by the environment.
         """
+        key = state.key if key is None else key
         rotation_index, x_position = action
         tetromino_index = state.tetromino_index
-        key, sample_key = jax.random.split(state.key)
+        key, sample_key = jax.random.split(key)
         tetromino = self._rotate(rotation_index, tetromino_index)
         # Place the tetromino in the selected place
         grid_padded, y_position = utils.place_tetromino(state.grid_padded, tetromino, x_position)

jumanji/environments/routing/cleaner/env.py

@@ -193,7 +193,9 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
 
         return state, timestep
 
-    def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Array, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """Run one timestep of the environment's dynamics.
 
         If an action is invalid, the corresponding agent does not move and

jumanji/environments/routing/connector/env.py

@@ -154,7 +154,9 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
         timestep = restart(observation=observation, extras=extras, shape=(self.num_agents,))
         return state, timestep
 
-    def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Array, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """Perform an environment step.
 
         Args:

jumanji/environments/routing/cvrp/env.py

@@ -159,7 +159,9 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
         timestep = restart(observation=self.observe(state))
         return state, timestep
 
-    def step(self, state: State, action: chex.Numeric) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Numeric, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """Run one timestep of the environment's dynamics.
 
         Args:

jumanji/environments/routing/lbf/env.py

@@ -184,7 +184,9 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep]:
 
         return state, timestep
 
-    def step(self, state: State, actions: chex.Array) -> Tuple[State, TimeStep]:
+    def step(
+        self, state: State, actions: chex.Array, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep]:
         """Simulate one step of the environment.
 
         Args:

jumanji/environments/routing/maze/env.py

@@ -192,7 +192,9 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
 
         return state, timestep
 
-    def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Array, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """
         Run one timestep of the environment's dynamics.
 

jumanji/environments/routing/mmst/env.py

@@ -190,12 +190,15 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
         timestep = restart(observation=self.observe(state), extras=extras)
         return state, timestep
 
-    def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Array, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """Run one timestep of the environment's dynamics.
 
         Args:
             state: State object containing the dynamics of the environment.
             action: Array containing the index of the next node to visit.
+            key: random key used to resolve simultaneous actions.
 
         Returns:
             state, timestep: Tuple[State, TimeStep] containing the next state of the
@@ -231,7 +234,8 @@ def step_agent_fn(
 
             return connected_nodes, conn_index, new_node, indices
 
-        key, step_key = jax.random.split(state.key)
+        key = state.key if key is None else key
+        key, step_key = jax.random.split(key)
         action, next_nodes = self._trim_duplicated_invalid_actions(state, action, step_key)
 
         connected_nodes = jnp.zeros_like(state.connected_nodes)

jumanji/environments/routing/multi_cvrp/env.py

@@ -163,7 +163,9 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
 
         return state, timestep
 
-    def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Array, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """
         Run one timestep of the environment's dynamics.
 

jumanji/environments/routing/pac_man/env.py

@@ -245,7 +245,9 @@ def reset(self, key: PRNGKey) -> Tuple[State, TimeStep[Observation]]:
 
         return state, timestep
 
-    def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Array, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """Run one timestep of the environment's dynamics.
 
         If an action is invalid, the agent does not move, i.e. the episode does not
@@ -256,12 +258,16 @@ def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observ
             action: (int32) specifying which action to take: [0,1,2,3,4] correspond to
                 [Up, Right, Down, Left, No-op]. If an invalid action is taken, i.e. there is a wall
                 blocking the action, then no action (no-op) is taken.
+            key: random key used to sample ghost actions. Defaults to state.key.
 
         Returns:
             state: the new state of the environment.
             the next timestep to be observed.
         """
 
+        if key is not None:
+            state = state.replace(key=key)
+
         # Collect updated state based on environment dynamics
         updated_state, collision_rewards = self._update_state(state, action)
 

jumanji/environments/routing/robot_warehouse/env.py

@@ -218,6 +218,7 @@ def step(
         self,
         state: State,
         action: chex.Array,
+        key: chex.PRNGKey | None = None,
     ) -> Tuple[State, TimeStep[Observation]]:
         """Perform an environment step.
 
@@ -229,14 +230,15 @@ def step(
                 - 2 turn left
                 - 3 turn right
                 - 4 toggle load
+            key: random key used to sample new shelf requests.
 
         Returns:
             state: State object corresponding to the next state of the environment.
             timestep: TimeStep object corresponding the timestep returned by the environment.
         """
 
         # unpack state
-        key = state.key
+        key = state.key if key is None else key
         grid = state.grid
         agents = state.agents
         shelves = state.shelves

jumanji/environments/routing/snake/env.py

@@ -165,7 +165,9 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
         timestep = restart(observation=self.observe(state))
         return state, timestep
 
-    def step(self, state: State, action: chex.Numeric) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Numeric, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """Run one timestep of the environment's dynamics.
 
         Args:
@@ -175,12 +177,14 @@ def step(self, state: State, action: chex.Numeric) -> Tuple[State, TimeStep[Obse
                 - 1: move to the right.
                 - 2: move down.
                 - 3: move to the left.
+            key: random key used to sample the next fruit position.
 
         Returns:
             state, timestep: next state of the environment and timestep to be observed.
         """
+        key = state.key if key is None else key
         is_valid = state.action_mask[action]
-        key, fruit_key = jax.random.split(state.key)
+        key, fruit_key = jax.random.split(key)
 
         head_position = self._update_head_position(state.head_position, action)
 

jumanji/environments/routing/sokoban/env.py

@@ -193,7 +193,9 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observation]]:
 
         return state, timestep
 
-    def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
+    def step(
+        self, state: State, action: chex.Array, key: chex.PRNGKey | None = None
+    ) -> Tuple[State, TimeStep[Observation]]:
         """
         Executes one timestep of the environment's dynamics.