Skip to content

Transitions

Transition functions: given (state, action, action_set), produce the next State.

An Environment's transitions_fn is one of these. They all start by applying the agent's action (action_set[action], a state -> state primitive - see navix.actions) and then optionally advance any autonomous entities. DEFAULT_TRANSITION is stochastic_transition.

The transitions_fn an Environment uses unless overridden: stochastic_transition (agent action + random ball motion).

Applies only the agent's action - nothing else in the world moves.

Parameters:

Name Type Description Default
state State

the current state, $s_t$.

 required
action Array

a scalar integer action, i32[], in [0, len(actions_set)).

 required
actions_set tuple[Callable, ...]

the environment's action_set; actions_set[action] is applied via jax.lax.switch.

 required

Returns:

Name Type Description
State State

$s_{t+1}$.

Applies the agent's action, then moves every Ball one random step (update_balls). This is DEFAULT_TRANSITION - environments without balls behave identically to deterministic_transition.

Parameters:

Name Type Description Default
state State

the current state, $s_t$.

 required
action Array

a scalar integer action, i32[], in [0, len(actions_set)).

 required
actions_set tuple[Callable, ...]

the environment's action_set.

 required

Returns:

Name Type Description
State State

$s_{t+1}$, with balls advanced.

Moves every Ball one cell in a uniformly random direction, or leaves it in place if that cell is blocked. A ball that would move onto the player instead stays put and records a (BALL, HIT) event (used by rewards.on_ball_hit / terminations.on_ball_hit).

Parameters:

Name Type Description Default
state State

the current state. If it has no Ball entities this is a no-op.

 required

Returns:

Name Type Description
State State

the state with ball positions and ball-hit events
State

updated, and state.key advanced.