factoriax.engine.envs.wrappers

factoriax.engine.envs.wrappers#

The wrappers that compose over a Factoriax environment.

class factoriax.engine.envs.wrappers.ActionMaskWrapper(inner: Environment[Any, EnvParams], blocked_actions: set[int] | list[int] | tuple[int, ...])#

Replace a blocked action with Action.NOOP.

The wrapper rewrites a blocked action and does not remove it. The action stays in the action space and still costs a timestep. An agent that sends one therefore loses a turn, and nothing stops it from choosing that action.

Parameters:

  • inner – Environment to wrap. Any gymnax-compatible environment with the state type EnvState works.

  • blocked_actions – Action integers to block. The constructor captures them and writes them into the JIT graph of step_env(). A different set therefore needs a new environment and a new compile.

action_space(params: EnvParams) → Discrete#

Return the action space of the inner environment, unchanged.

Parameters:

params – The method does not read this argument. It is present for the gymnax signature.

Returns:

The action space of the inner environment, unchanged. A blocked action stays in the space, and the wrapper rewrites it instead of removing it. An agent can therefore still send it, and the output width of a policy does not depend on the mask.

Return type:

spaces.Discrete

property default_params: EnvParams#

Return the default parameters of the inner environment.

get_obs(state: Any, params: EnvParams) → Array#

Return the observation of the inner environment.

Parameters:

  • state – State to observe.

  • params – Environment parameters.

Returns:

The observation of the inner environment.

Return type:

jax.Array

is_terminal(state: Any, params: EnvParams) → Array#

Ask the inner environment whether the episode is at its end.

Parameters:

  • state – State to test.

  • params – Environment parameters.

Returns:

Scalar bool from the inner environment.

Return type:

jax.Array

property map_height: int#

Return the map height in tiles, from the inner environment.

property map_width: int#

Return the map width in tiles, from the inner environment.

property num_players: int#

Return the number of players, from the inner environment.

observation_space(params: EnvParams) → Box#

Return the observation space of the inner environment, unchanged.

Parameters:

params – The method does not read this argument. It is present for the gymnax signature.

Returns:

The observation space of the inner environment, unchanged.

Return type:

spaces.Box

reset_env(key: Array, params: EnvParams) → tuple[Array, Any]#

Reset the inner environment and return its result.

Parameters:

  • key – PRNG key for the world generation.

  • params – Environment parameters.

Returns:

(obs, state) from the inner environment.

Return type:

tuple

step_env(key: Array, state: Any, action: int | Array, params: EnvParams) → tuple[Array, Any, Array, Array, dict[str, Any]]#

Rewrite a blocked action to NOOP, then step the inner environment.

Parameters:

  • key – PRNG key. The method passes it to the inner environment.

  • state – State to advance.

  • action – Action to take. The method rewrites it to NOOP when it is blocked.

  • params – Environment parameters.

Returns:

(obs, new_state, reward, done, info) from the inner environment, after the rewrite. A blocked action still costs a timestep.

Return type:

tuple

class factoriax.engine.envs.wrappers.AutoResetState(env_state: EnvState, reset_state: EnvState)#

State that holds a copy of the start world, for a cheap auto-reset.

The class keeps a frozen copy of the start state next to the live state. At the end of an episode, lax.select puts the copy back. It does not call reset_env, which runs the whole terrain generation.

The cost is one extra copy of EnvState for each batch element. At 32x32 that is about 12 KB for each element. At 128x128 with max_machines=4096 it grows to about 185 KB for each element, which doubles the total state memory. On a large map, lower batch_size or max_machines.

replace(**updates)#

Returns a new object replacing the specified fields with new values.

class factoriax.engine.envs.wrappers.AutoResetWrapper(inner: FactoriaxEnv, resample: bool = False)#

A gymnax wrapper that auto-resets from a cached state.

Use this wrapper when you need an auto-reset inside a lax.scan training loop, such as a PPO loop in the PureJaxRL style. To manage the episodes yourself, use FactoriaxEnv directly.

Parameters:

  • inner – Core Factoriax environment to wrap.

  • resample – Whether a new episode builds a new world. See __init__().

action_space(params: EnvParams) → Discrete#

Return the action space of the inner environment, unchanged.

Parameters:

params – The method does not read this argument. It is present for the gymnax signature.

Returns:

The action space of the inner environment, unchanged.

Return type:

spaces.Discrete

property default_params: EnvParams#

Return the default parameters of the inner environment.

get_obs(state: AutoResetState, params: EnvParams) → Array#

Return the observation of the inner environment.

Parameters:

  • state – Wrapped state.

  • params – Environment parameters.

Returns:

The observation that the inner environment makes of the live state, and not of the cached state.

Return type:

jax.Array

is_terminal(state: AutoResetState, params: EnvParams) → Array#

Ask the inner environment whether the episode is at its end.

Parameters:

  • state – Wrapped state.

  • params – Environment parameters.

Returns:

Scalar bool from the inner environment.

Return type:

jax.Array

property map_height: int#

Return the map height in tiles, from the inner environment.

property map_width: int#

Return the map width in tiles, from the inner environment.

property num_players: int#

Return the number of players, from the inner environment.

property obs: str#

Return the observation variant key, from the inner environment.

property obs_radius: int#

Return the half-width of the local window, from the inner environment.

observation_space(params: EnvParams) → Box#

Return the observation space of the inner environment, unchanged.

Parameters:

params – The method does not read this argument. It is present for the gymnax signature.

Returns:

The observation space of the inner environment, unchanged.

Return type:

spaces.Box

reset_env(key: Array, params: EnvParams) → tuple[Array, AutoResetState]#

Reset, and cache the start state for the later auto-resets.

Parameters:

  • key – PRNG key for the world generation.

  • params – Environment parameters.

Returns:

(obs, state). The method takes the cached copy here and never replaces it. With resample=False, every episode of the run therefore repeats this exact world.

Return type:

tuple

step_env(key: Array, state: AutoResetState, action: int | Array, params: EnvParams) → tuple[Array, AutoResetState, Array, Array, dict[str, Any]]#

Step the environment, and auto-reset from the cache at the end.

When done is True, lax.select puts the cached reset state in place of the live state. No terrain generation runs.

Parameters:

  • key – PRNG key. The method splits it into a step key and a reset key, so the reset branch draws a new world under resample.

  • state – Wrapped state, which holds both the live world and the cached world.

  • action – Action for the selected player.

  • params – Environment parameters.

Returns:

(obs, new_state, reward, done, info). When done is True, the state in the result is already the first state of the next episode, and reward and done still describe the step that ended. Both branches run in every step, because jax.lax.select() evaluates both.

Return type:

tuple

class factoriax.engine.envs.wrappers.LogEnvState(env_state: Any, episode_returns: Array, episode_lengths: Array, returned_episode_returns: Array, returned_episode_lengths: Array, timestep: Array)#

State that totals the return and the length of each episode, for the log.

Every scalar field is a JAX array, so the state is safe under vmap and under scan. The returned_* fields hold the numbers of the last finished episode. They are non-zero only on the step where done becomes True.

env_state#

State of the inner wrapper. It can be any JAX pytree.

Type:

Any

episode_returns#

Running return of the current episode. It returns to 0 on done.

Type:

jax.jaxlib._jax.Array

episode_lengths#

Running step count of the current episode. It returns to 0 on done.

Type:

jax.jaxlib._jax.Array

returned_episode_returns#

Return of the episode that just finished. It is 0 while an episode runs.

Type:

jax.jaxlib._jax.Array

returned_episode_lengths#

Length of the episode that just finished. It is 0 while an episode runs.

Type:

jax.jaxlib._jax.Array

timestep#

Global step counter. Nothing resets it.

Type:

jax.jaxlib._jax.Array

replace(**updates)#

Returns a new object replacing the specified fields with new values.

class factoriax.engine.envs.wrappers.LogWrapper(inner: Environment[Any, EnvParams])#

Add the numbers of each episode to the info dict in every step.

This wrapper takes any gymnax-compatible environment. Every call to step_env() adds these keys to the info dict:

  • "returned_episode_returns", the return of the episode that finished last. It is non-zero only on the final step of an episode.

  • "returned_episode_lengths", the length of the episode that finished last. It is non-zero only on the final step of an episode.

  • "returned_episode", which is True on the step where an episode ends.

  • "timestep", the global step counter, which rises with every call.

These keys follow the LogWrapper convention of PureJaxRL. A training loop can therefore read the episode numbers directly from traj_batch.info, which jax.lax.scan returns.

Examples

>>> import factoriax
>>> env, params = factoriax.make("EasyRocket-v1", auto_reset=True)
>>> from factoriax.engine.envs.wrappers import LogWrapper
>>> log_env = LogWrapper(env)
>>> # log_env.step_env(...) now returns info["returned_episode_returns"]
action_space(params: EnvParams) → Discrete#

Return the action space of the inner environment, unchanged.

Parameters:

params – The method does not read this argument. It is present for the gymnax signature.

Returns:

The action space of the inner environment, unchanged.

Return type:

spaces.Discrete

property default_params: EnvParams#

Return the default parameters of the inner environment.

get_obs(state: LogEnvState, params: EnvParams) → Array#

Return the observation of the inner environment.

Parameters:

  • state – Wrapped state.

  • params – Environment parameters.

Returns:

The observation of the inner environment.

Return type:

jax.Array

is_terminal(state: LogEnvState, params: EnvParams) → Array#

Ask the inner environment whether the episode is at its end.

Parameters:

  • state – Wrapped state.

  • params – Environment parameters.

Returns:

Scalar bool from the inner environment.

Return type:

jax.Array

property map_height: int#

Return the map height in tiles, from the inner environment.

property map_width: int#

Return the map width in tiles, from the inner environment.

observation_space(params: EnvParams) → Box#

Return the observation space of the inner environment, unchanged.

Parameters:

params – The method does not read this argument. It is present for the gymnax signature.

Returns:

The observation space of the inner environment, unchanged.

Return type:

spaces.Box

reset_env(key: Array, params: EnvParams) → tuple[Array, LogEnvState]#

Reset the inner environment, and set every episode counter to zero.

Parameters:

  • key – PRNG key for the world generation.

  • params – Environment parameters.

Returns:

(obs, state), with the running episode counters at zero.

Return type:

tuple

step_env(key: Array, state: LogEnvState, action: int | Array, params: EnvParams) → tuple[Array, LogEnvState, Array, Array, dict[str, Any]]#

Step the inner environment, and update the episode counters.

On the step where done is True:

  • returned_episode_returns takes the return that just finished.

  • returned_episode_lengths takes the length that just finished.

  • episode_returns and episode_lengths return to zero.

Every info key is always present. Use info["returned_episode"] as a boolean mask to select the valid returned_episode_returns entries.

Parameters:

  • key – PRNG key. The method passes it to the inner environment.

  • state – Wrapped state, which carries the running counters.

  • action – Action for the selected player.

  • params – Environment parameters.

Returns:

(obs, new_state, reward, done, info). info carries the return and the length of the episode. Those two have meaning only on the step where done is True.

Return type:

tuple