factoriax.engine.envs.base

factoriax.engine.envs.base#

The environment object a training loop talks to, and the hooks that shape it.

FactoriaxEnv wraps the engine in the gymnax interface. reset_env builds a world, step_env advances it, and the two space methods describe what an agent sends and what it receives.

The engine itself knows nothing about a task. It has no reward, no goal, and no end except the end of the time budget. An author builds a scenario with functions that the constructor takes, and not with a subclass. Those functions turn the same simulation into mining practice or into a rocket launch:

achievement_fn

Reads a state and returns the achievement bits that hold in it.

reward_fn

Scores a step. Without it, every step scores zero.

done_fn

Ends an episode early. The environment ORs it with the timeout and never replaces the timeout.

terrain_fn and level

Decide which world a reset produces.

step_hooks and reset_hooks

Change a state after the engine finishes with it.

The constructor captures each one. Two environments that differ in one of them are therefore two JIT cache entries, and each one pays for its own compile.

class factoriax.engine.envs.base.FactoriaxEnv(achievement_fn: Callable[[EnvState], Array] | None = None, level: Level | None = None, terrain_fn: Callable[[Array, EnvParams], Array] | None = None, step_hooks: tuple[Callable[[Array, EnvState, EnvParams], EnvState], ...] = (), reset_hooks: tuple[Callable[[Array, EnvState, EnvParams], EnvState], ...] = (), reward_fn: Callable[[EnvState, EnvState, EnvParams], Array] | None = None, done_fn: Callable[[EnvState, EnvParams], Array] | None = None, obs: str = 'x_ray_global', obs_radius: int = 7, map_width: int = 32, map_height: int = 32, num_players: int = 1, max_machines: int = 0)#

The JAX grid environment of Factoriax.

The class advances the world state, which holds the terrain, the machines, and the player inventories. It also evaluates an optional achievement-condition function in each step. The reward and the other parts that shape a policy belong in the gymnax wrappers around this environment. factoriax.engine.renderer.JaxRenderer draws the pixels. This class exposes no render method.

The constructor captures achievement_fn, and step_env() folds its result into state.achievements_unlocked. Two different functions give two JIT cache entries, through static_argnames=("self",) on step(). Pass None to skip that pass, which then costs nothing.

terrain_fn and level decide what a reset builds, in that order. With a terrain_fn, reset_env() calls it with the PRNG key, and the scenario owns the world generation. Without one, a level of None means a generated world from the key, and a Level goes through build_state(), which ignores the key for the layout. The reset_hooks then change the new state on every reset path, for example to fill a player inventory. They take the same arguments as step_hooks. reset_env(key, params) is the one reset method, as gymnax requires.

Parameters:

  • achievement_fn – Reads a state and returns bool[MAX_ACHIEVEMENTS]. The environment evaluates it in every step, and a True bit latches into state.achievements_unlocked. None skips the pass.

  • level – Fixed Level for reset_env() to build. None means a world generated from the PRNG key.

  • terrain_fn – Builds a terrain grid from the key. It wins over level, so a scenario that supplies both gets generated terrain and no placed machines.

  • step_hooks – Hooks that run in order after every step. Each one takes (key, state, params) and returns a state.

  • reset_hooks – Hooks that run in order after every reset, with the same arguments. Use them to fill a player inventory or to place the machines of a scenario.

  • reward_fn – Scores each step from (prev_state, new_state, params). None means that every step returns 0.0.

  • done_fn – Ends an episode from (state, params). The environment ORs it with the timeout, so it can end an episode early but can never extend one.

  • obs – Key into OBSERVATIONS that names the profile and the view, for example "x_ray_global". An unknown key raises ValueError.

  • obs_radius – Half-width of the window for a local observation. A global view ignores it.

  • map_width – Map width in tiles. The environment ignores it when level supplies the map.

  • map_height – Map height in tiles. The environment ignores it when level supplies the map.

  • num_players – Number of players in the world. Only selected_player acts.

  • max_machines – Number of entity slots to allocate. A value of 0 asks for a size from the map area. The value is fixed for the episode, and a map with every slot in use refuses a new placement.

action_space(params: EnvParams) → Discrete#

Describe the action space, which holds every action of the engine.

The size is NUM_ACTIONS for every scenario. A scenario whose recipe book cannot make an item still exposes the CRAFT_ action of that item. That action does nothing, and it is not an error. A wrapper masks the actions that a scenario cannot use.

Parameters:

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

Returns:

A space of size NUM_ACTIONS.

Return type:

spaces.Discrete

property default_params: EnvParams#

Return the default parameters of this environment.

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

Build the observation of the selected player.

The method always observes state.selected_player. With more than one player, the result is therefore the view of one agent, and not a stack with one view for each player.

Parameters:

  • state – State to observe.

  • params – Environment parameters.

Returns:

1-D float32. observation_space() gives its length.

Return type:

jax.Array

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

Report whether the episode is at its end.

The timeout always applies. The method ORs a done_fn on top of it, so a scenario can end an episode early. No scenario can make an episode run longer than params.max_timesteps.

Parameters:

  • state – State to test.

  • params – Supplies max_timesteps.

Returns:

Scalar bool.

Return type:

jax.Array

observation_space(params: EnvParams) → Box#

Describe the observation space of the selected variant.

The bounds are 0.0 to 1.0, and every channel scales to stay inside them. This includes the x_ray facing readouts, which divide each count by a value equal to or larger than the MACHINE_MAX_STACK of every machine.

Parameters:

params – The method does not read this argument. The shape comes from the constructor arguments, so it is fixed for the life of the environment.

Returns:

1-D float32, shape (obs_size,).

Return type:

spaces.Box

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

Reset the environment to a start state.

There are three ways to build a world, and the method tries them in this order: terrain_fn, then a bound level, then the terrain generator. With a level, the method calls build_state(), which ignores the key for the layout.

CAUTION: A scenario that supplies both a terrain_fn and a level gets the terrain and loses the machines of the level. Few scenarios want that.

The reset_hooks run after the build, on every path.

Parameters:

  • key – PRNG key for the world generation. A bound level ignores it for the layout, but the method still passes it to the hooks.

  • params – Environment parameters.

Returns:

(obs, state), with timestep at 0.

Return type:

tuple

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

Step the environment, with no auto-reset.

This method overrides the gymnax step(). That method calls reset_env in every tick for its auto-reset. It therefore runs the whole terrain generation in every step, even when the episode is not done, and about triples the cost of a step. This method calls step_env directly instead.

Use AutoResetWrapper when you need auto-reset for a lax.scan training loop.

Parameters:

  • key – PRNG key. The method passes it to the hooks. The step itself is deterministic.

  • state – State to advance.

  • action – Action for the selected player.

  • params – Environment parameters. None takes default_params.

Returns:

(obs, new_state, reward, done, info). There is no auto-reset. When done is True, the state in the result is the final state of the episode, and not the start of a new one.

Return type:

tuple

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

Run one environment step.

The method runs the game rules. It then evaluates the optional achievement condition function and folds the result into state.achievements_unlocked with OR. It skips that pass when the constructor received no achievement_fn. The JIT trace resolves the is None test, so that path costs no branch in a step.

The order matters. The engine runs first, then the step_hooks, then the reward against the state after the hooks, then the end test, then the observation. A hook can therefore change the value of a step, and reward_fn sees the state that the agent observes.

Parameters:

  • key – PRNG key. The method passes it to the hooks.

  • state – State to advance.

  • action – Action for the selected player. The method casts it to int32.

  • params – Environment parameters.

Returns:

(obs, new_state, reward, done, info). reward is 0.0 when the constructor received no reward_fn, and info is always empty.

Return type:

tuple

factoriax.engine.envs.base.achievement_hook(condition_fn: Callable[[EnvState], Array]) → Callable[[Array, EnvState, EnvParams], EnvState]#

Wrap an achievement condition into a step hook that latches its bits.

The latch makes an achievement a milestone and not a state flag. Once a bit is set, it stays set for the rest of the episode, even after the condition stops holding. factoriax.engine.rewards.achievement_reward() therefore pays for it exactly one time.

Parameters:

condition_fn – Reads a state and returns the bits that hold now. Shape (MAX_ACHIEVEMENTS,), bool.

Returns:

A hook that folds those bits into achievements_unlocked with OR. The hook ignores its key argument and its params argument.

Return type:

StepHook