Wrappers

Adapt a navix Environment to the gymnax API, so navix environments can be dropped into agents and tooling written against gymnax ((obs, state) from reset, (obs, state, reward, done, info) from step, with an explicit params).

GymnaxState

Bases: EnvState

A gymnax EnvState that simply carries the wrapped navix Timestep. ToGymnax.step unwraps timestep, advances the navix environment, and rewraps.

Attributes:

Name	Type	Description
timestep	Timestep	the underlying navix Timestep.
time	Array	timestep.t, exposed at the top level because gymnax reads state.time for its own truncation bookkeeping.

ToGymnax

Bases: Environment

Wraps a navix Environment as a gymnax.environments.Environment.

Prefer ToGymnax.wrap(env), which returns the (gymnax_env, params) pair gymnax callers expect.

import navix as nx
from navix.environments.wrappers import ToGymnax

gymnax_env, params = ToGymnax.wrap(nx.make("Navix-Empty-5x5-v0"))
obs, state = gymnax_env.reset(key, params)
obs, state, reward, done, info = gymnax_env.step(key, state, action, params)

done merges navix's truncation and termination into gymnax's single boolean; StepType is still available via state.timestep.step_type. Autoreset follows navix's rules (see Environment.step), not gymnax's.

default_params property

A gymnax EnvParams carrying the wrapped env's max_steps as max_steps_in_episode.

__init__(env)

Args: env (Environment): the navix environment to wrap.

action_space(params)

The gymnax Discrete action space, len(env.action_set) values. params is ignored (kept for API compatibility).

observation_space(params)

The wrapped env's observation Space as a gymnax Box (low/high/shape/dtype copied across). params is ignored.

reset(key, params=None)

Resets the wrapped environment.

Parameters:

Name	Type	Description	Default
key	Array	PRNG key.	required
params	EnvParams | None	ignored; max_steps comes from the wrapped env.	None

Returns:

Name	Type	Description
tuple	Tuple[Array, EnvState]	(observation, GymnaxState).

step(key, state, action, params)

Advances the wrapped environment one step.

Parameters:

Name	Type	Description	Default
key	Array	PRNG key (unused - navix carries its own key in state.timestep.state.key).	required
state	GymnaxState	the state returned by the previous reset/step.	required
action	Array	integer action, [0, len(action_set)).	required
params	EnvParams	ignored.	required

Returns:

Name	Type	Description
tuple	Array	(observation, GymnaxState, reward, done, info), where
	EnvState	done = timestep.is_done() (truncation or termination).

wrap(env) classmethod

Builds the wrapper and its params in one call.

Parameters:

Name	Type	Description	Default
env	Environment	the navix environment to wrap.	required

Returns:

Name	Type	Description
tuple	Tuple[Environment, EnvParams]	(ToGymnax(env), EnvParams(max_steps_in_episode=env.max_steps)).