Environment

The Environment class and the Timestep it produces.

An Environment is a frozen JAX pytree: it holds the grid geometry (height, width, max_steps) and a set of pluggable functions - observation_fn, reward_fn, termination_fn, transitions_fn, action_set - that together define the task. reset and step are pure functions of (key) / (timestep, action), so they compose with jax.jit, jax.vmap (a batch of environments) and jax.lax.scan (a whole rollout in one compiled loop).

Concrete environments (navix.environments.Empty, DoorKey, ...) subclass Environment and implement _reset to lay out their grid; everything else is inherited. Build one with navix.make(id) or SomeEnv.create(...).

Environment

Bases: PyTreeNode

A gridworld task as a frozen JAX pytree.

The task is defined by five pluggable pieces, each a plain function you can override per-make/create call:

• observation_fn(state) -> Array - what the agent sees.
• transitions_fn(state, action, action_set) -> State - the world dynamics (applies action_set[action], then any stochastic entity updates).
• reward_fn(prev_state, action, state) -> f32[] - the reward for a transition. The (prev_state, action, state) triple is the convention across navix.rewards / terminations / events: prev_state is $s_t$, state is $s_{t+1}$.
• termination_fn(prev_state, action, state) -> bool[] - whether $s_{t+1}$ is a genuine terminal state (truncation at max_steps is added on top automatically).
• action_set - the tuple of state -> state primitives an integer action indexes into.

Subclasses only implement _reset (the initial grid layout); reset/step are inherited. Instances are immutable - use env.replace(...) (from flax.struct) to get a modified copy.

Attributes:

Name	Type	Description
height	int	grid height in cells, including the surrounding wall.
width	int	grid width in cells, including the surrounding wall.
max_steps	int	truncation horizon - step returns StepType.TRUNCATION once t >= max_steps. Default (via create) is 4 * height * width.
observation_space	Space	Space describing observation_fn's output (shape, dtype, bounds).
action_space	Space	Discrete over len(action_set).
reward_space	Space	Continuous bound on the reward, [-1, 1] by default.
disable_autoreset	bool	if False (default), calling step on a timestep whose episode already ended returns a fresh reset instead of stepping. Set True to handle episode boundaries yourself.
gamma	float	discount factor. Not used by step itself - carried here so agents and reward_fns (e.g. time-discounted goal rewards) can read it off the environment.
penalty_coeff	float	if non-zero, a terminating reward is reduced by penalty_coeff * (t / max_steps), i.e. finishing later is worth less. 0.0 disables it. (The old misspelled name penality_coeff still works as a deprecated read-only alias.)
observation_fn	Callable[[State], Array]	state -> observation. One of the functions in navix.observations.
reward_fn	Callable[[State, Array, State], Array]	(prev_state, action, state) -> f32[].
termination_fn	Callable[[State, Array, State], Array]	(prev_state, action, state) -> bool[].
transitions_fn	Callable[[State, Array, Tuple[Callable[[State], State], ...]], State]	(state, action, action_set) -> state. Usually transitions.deterministic_transition or transitions.stochastic_transition (the default, which also moves balls).
action_set	Tuple[Callable[[State], State], ...]	tuple of state -> state callables; an integer action a applies action_set[a].

penality_coeff property

Deprecated misspelling of penalty_coeff. Reads still work (with a warning); pass penalty_coeff to create going forward.

create(height, width, max_steps=None, observation_fn=observations.symbolic, reward_fn=rewards.DEFAULT_TASK, termination_fn=terminations.DEFAULT_TERMINATION, transitions_fn=transitions.DEFAULT_TRANSITION, action_set=DEFAULT_ACTION_SET, observation_space=None, action_space=None, reward_space=None, disable_autoreset=False, **kwargs) classmethod

Builds an environment, filling in the spaces and max_steps that weren't given.

This is the shared constructor concrete environments call from their own create; navix.make(id, **kwargs) ends up here too.

Parameters:

Name	Type	Description	Default
height	int	grid height in cells (walls included).	required
width	int	grid width in cells (walls included).	required
max_steps	int | None	truncation horizon. None -> 4 * height * width.	None
observation_fn	Callable	state -> observation, from navix.observations. Default observations.symbolic.	symbolic
reward_fn	Callable	(prev_state, action, state) -> f32[].	DEFAULT_TASK
termination_fn	Callable	(prev_state, action, state) -> bool[].	DEFAULT_TERMINATION
transitions_fn	Callable	(state, action, action_set) -> state.	DEFAULT_TRANSITION
action_set	tuple[Callable, ...]	state -> state primitives indexed by the integer action.	DEFAULT_ACTION_SET
observation_space	Space | None	None infers it from observation_fn and the grid size (works for the built-in observation functions; pass one explicitly for a custom observation_fn).	None
action_space	Space | None	None -> Discrete(len(action_set)).	None
reward_space	Space | None	None -> Continuous((), -1, 1).	None
disable_autoreset	bool	see the class attribute.	False
**kwargs		extra fields forwarded to the subclass constructor (e.g. gamma, penalty_coeff, or an environment's own layout options like random_start).	{}

Returns:

Name	Type	Description
Environment	Environment	the constructed environment.

reset(key, cache=None)

Starts a new episode.

Parameters:

Name	Type	Description	Default
key	Array	a jax.random PRNG key. Split it yourself across a batch (jax.vmap(env.reset)(keys)).	required
cache	RenderingCache | None	optional pre-built rendering cache (see _reset).	None

Returns:

Name	Type	Description
Timestep	Timestep	the first timestep - t = 0, is_start() true,
	Timestep	info["return"] = 0.0. state.key is a fresh key derived
	Timestep	from key, so the environment's own stochasticity is
	Timestep	reproducible from the single seed you pass here.

step(timestep, action)

Advances one timestep, or auto-resets at an episode boundary.

If timestep already ended an episode (its step_type > 0) and disable_autoreset is False, this ignores action and returns a fresh reset seeded from timestep.state.key. Otherwise it applies action via transitions_fn, then evaluates reward_fn, termination_fn and observation_fn on the result.

Autoreset is deferred by one call: the terminal timestep is returned as-is (so you can read its final reward), and the reset happens on the next step. Detect a fresh episode with timestep.is_start(), not by looking back at is_done().

Parameters:

Name	Type	Description	Default
timestep	Timestep	the current timestep (from reset or a previous step).	required
action	Array	an integer action, i32[], in [0, len(action_set)). Indexes action_set.	required

Returns:

Name	Type	Description
Timestep	Timestep	the next timestep. info["return"] accumulates the
	Timestep	undiscounted episodic reward.

termination(prev_state, action, state, t)

Combines the task's termination_fn with the max_steps truncation into a single StepType.

Parameters:

Name	Type	Description	Default
prev_state	State	$s_t$.	required
action	Array	$a_t$.	required
state	State	$s_{t+1}$.	required
t	Array	the step count of state (i32[]).	required

Returns:

Name	Type	Description
Array	Array	a StepType value - TERMINATION if termination_fn
	Array	fired, else TRUNCATION if t >= max_steps, else
	Array	TRANSITION. Termination takes precedence over truncation.

StepType

Bases: PyTreeNode

The three kinds of timestep, stored as Timestep.step_type. The distinction between the two "episode over" cases matters for bootstrapping a value function: bootstrap through a TRUNCATION, but not through a TERMINATION.

TERMINATION = jnp.asarray(2) class-attribute instance-attribute

The episode reached a genuine terminal (absorbing) state - the goal, lava, a wrong toggle, etc. Its value is 0; do not bootstrap.

TRANSITION = jnp.asarray(0) class-attribute instance-attribute

The episode continues; step will keep advancing it.

TRUNCATION = jnp.asarray(1) class-attribute instance-attribute

The episode was cut off at max_steps while still ongoing. The final state is not absorbing - its value is non-zero - so a value estimate should bootstrap through it.

Timestep

Bases: PyTreeNode

What Environment.reset and Environment.step return - one moment in a trajectory.

Read it as the result of the transition that produced it: state and t describe the world now, and action / reward / step_type are the action that got here and its consequences. So after ts_next = env.step(ts, a): ts_next.state is $s_{t+1}$, ts_next.action is $a$ (ts.action is discarded), ts_next.reward is $R(s_t, a, s_{t+1})$, and ts_next.observation is the observation of $s_{t+1}$.

Every field is a JAX array (or a pytree of them), so a Timestep vmaps over a batch of environments and scans over time.

Attributes:

Name	Type	Description
t	Array	steps since the last reset, i32[] (or i32[batch] when vmapped). 0 exactly on a reset output.
observation	Array	the agent's view of state, produced by the environment's observation_fn. Shape and dtype are described by env.observation_space; for a pomdp observation function this is a partial (first-person, cropped) view.
action	Array	the action that led to this timestep, i32[]. On a reset output it is a placeholder 0.
reward	Array	the scalar reward for the transition into this timestep, f32[], produced by the environment's reward_fn (bounds given by env.reward_space).
step_type	Array	a StepType value (0/1/2) - see StepType.
state	State	the full State (the true MDP state); the observation is a function of this. Also carries the PRNG key and the rendering cache.
info	Dict[str, Any]	a plain dict for extra per-step quantities. navix populates info["return"] (undiscounted return so far this episode); you may add your own keys.

is_done()

True where the episode has ended for either reason - truncation or termination. This is the mask you use to segment trajectories or to stop bootstrapping a return.

is_start()

True on the first timestep of an episode (t == 0) - both the initial reset and every post-autoreset step. Robust to navix deferring autoreset by one step call, unlike checking a previous step's is_done().

is_termination()

True where step_type == StepType.TERMINATION (genuine absorbing terminal state). Boolean, same batch shape as step_type.

is_transition()

True where the episode is still ongoing (step_type == StepType.TRANSITION). Boolean.

is_truncation()

True where step_type == StepType.TRUNCATION (episode cut off at max_steps). Boolean, same batch shape as step_type.