factoriax.engine.state

State of the environment, and its per-episode parameters.

EnvState is everything that a step reads and writes. EnvParams is the tuning that stays fixed while an episode runs. Both are flax.struct.PyTreeNode, so their array fields are PyTree leaves. Both therefore pass through jax.jit and jax.vmap with no further registration.

The engine holds machine state as entity lists, and not as one grid for each attribute. An entity id indexes every ent_ array. The capacity is fixed when the level is built, so the shapes stay static, whatever a player builds during an episode.

class factoriax.engine.state.EnvParams(max_timesteps: int = 1000, water_probability: float = 0.1, iron_probability: float = 0.12, copper_probability: float = 0.12, coal_probability: float = 0.12, tin_probability: float = 0.1, silicon_probability: float = 0.1, base_resources: int = 1000, miner_mining_rate: int = 3, player_mining_yield: int = 1, recipe_table: RecipeTable = RecipeTable(outputs=Array([7, 8, 9, 10, 11, 12, 13, 14, 15, 26, 24, 27, 25, 28, 30, 16, 17, 23, 29, 31, 32, 33, 18, 19, 20, 21, 22], dtype=int32), output_counts=Array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], dtype=int32), input_items=Array([[2, 1], [3, 1], [4, 1], [5, 1], [7, 9], [8, 10], [8, 9], [11, 13], [12, 13], [7, 8], [7, 13], [11, 12], [9, 13], [8, 13], [7, 18], [14, 9], [15, 10], [17, 15], [19, 22], [12, 14], [9, 1], [8, 1], [6, 1], [11, 7], [14, 13], [12, 15], [20, 21]], dtype=int32), input_counts=Array([[1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [6, 4], [2, 2], [1, 1], [1, 1], [1, 1], [2, 2], [2, 1], [2, 2], [1, 1]], dtype=int32), ticks=Array([2, 2, 2, 2, 4, 4, 4, 6, 6, 4, 6, 8, 4, 4, 6, 8, 8, 8, 300, 8, 4, 4, 4, 8, 8, 10, 10], dtype=int32), machine_type=Array([7, 7, 7, 7, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 7, 4, 4, 4, 4], dtype=int32), output_to_recipe=Array([-1, -1, -1, -1, -1, -1, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 15, 16, 22, 23, 24, 25, 26, 17, 10, 12, 9, 11, 13, 18, 14, 19, 20, 21], dtype=int32), names=('Iron Plate', 'Copper Plate', 'Tin Plate', 'Wafer', 'Frame', 'Circuit', 'Wire', 'Motor', 'Sensor', 'Conveyor Belt', 'Miner', 'Assembler', 'Pallet', 'Arm', 'Furnace', 'Tier 1 Science Pack', 'Tier 2 Science Pack', 'Tier 3 Science Pack', 'Rocket', 'Science Lab', 'Splitter', 'Crossing', 'Refractory', 'Hull', 'Engine Unit', 'Avionics', 'Rocket Core')))

Tuning that stays fixed for the length of an episode.

Every engine function takes this object next to the state, so a scenario can change the balance and needs no change to engine code. The defaults describe the generated sandbox map. A scenario that brings its own level overrides the terrain fields, and nothing reads them.

The six terrain probabilities apply only when the engine generates a level, and not when it loads one. The engine tests each probability against its own smooth-noise field. They are therefore per-tile shares and not bands, and their total does not have to be 1. If two draws claim the same tile, the later draw wins. The order is silicon, tin, coal, copper, iron, water, so water replaces any ore on the same tile. A tile that no draw claims stays plain ground.

max_timesteps

Steps before the episode reports done. The engine compares this value against EnvState.timestep.

Type:

int

water_probability

Share of generated tiles that become water, which blocks movement and placement.

Type:

float

iron_probability

Share of generated tiles that become iron ore.

Type:

float

copper_probability

Share of generated tiles that become copper ore.

Type:

float

coal_probability

Share of generated tiles that become coal.

Type:

float

tin_probability

Share of generated tiles that become tin ore.

Type:

float

silicon_probability

Share of generated tiles that become silicon.

Type:

float

base_resources

Ore units in every mineable tile at the start. The engine writes the value into EnvState.block_resources, and a tile that no player can mine gets 0. The value must fit in int16. With more than BLOCK_MAX_RESOURCES the tile still works, but it normalizes to more than 1.0 in the observation.

Type:

int

miner_mining_rate

Ore that a placed miner extracts in one step. Its output slot holds factoriax.engine.machines.MINER_OUTPUT_CAP, so a miner reaches that limit in one step. It then stays idle until something empties the slot.

Type:

int

player_mining_yield

Ore that one MINE action gives to the player. The ore left in the tile and the stack limit of the player both limit this amount.

Type:

int

recipe_table

Recipes that this episode runs, as a RecipeTable. The default is the table of the engine itself. Engine code takes the size of its recipe loops from this table, and not from the module-level default. A scenario can therefore supply a different number of recipes.

Type:

factoriax.engine.recipes.RecipeTable

NUM_ACTIONS: ClassVar[int] = 87

Size of the action space. This is a ClassVar, so it is not a PyTree leaf and no episode can tune it. The Action enum fixes the action space.

replace(**updates)

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

class factoriax.engine.state.EnvState(map: Array, block_resources: Array, machine_types: Array, tile_entity: Array, ent_y: Array, ent_x: Array, ent_type: Array, ent_direction: Array, ent_power: Array, ent_buf_type: Array, ent_buf_count: Array, ent_asm_in_type: Array, ent_asm_in_count: Array, ent_asm_out_type: Array, ent_asm_out_count: Array, ent_health: Array, player_positions: Array, player_directions: Array, player_inventory: Array, selected_player: int, timestep: int, items_mined: Array, science_consumed_step: Array, achievements_unlocked: Array)

Everything one environment step reads and writes.

Machine state lives in the ent_ arrays. An entity id indexes them, not a tile. A slot is free when ent_y < 0. Placement claims the first free slot, and pickup releases it. The capacity is fixed when the level is built. A map with every slot in use therefore refuses a new placement, and no array grows.

Two fields hold the same fact, and this is deliberate. machine_types answers “is something on this tile” for movement and placement. ent_type carries the kind of each entity for the machine tick. tile_entity joins the two, and placement writes all three together.

The shapes below use H and W for the map, E for the entity capacity, and P for the player count.

map

Terrain, one BlockType per tile. Shape (H, W), int8.

Type:

jax.jaxlib._jax.Array

block_resources

Ore units left in each tile. A mine action lowers the value. Shape (H, W), int16. Zero means that the tile is empty. The value has no meaning on a tile that was never an ore block.

Type:

jax.jaxlib._jax.Array

machine_types

Machine kind on each tile, NONE where the tile is clear. Shape (H, W), int8.

Type:

jax.jaxlib._jax.Array

tile_entity

Entity id on each tile, -1 where the tile holds no entity. Shape (H, W), int16. A neighbour lookup reads this array to go from a tile to entity state.

Type:

jax.jaxlib._jax.Array

ent_y

Tile row of each entity. Shape (E,), int16. A negative value marks a free slot. This is the test that the engine uses for “this entity is not on the map”.

Type:

jax.jaxlib._jax.Array

ent_x

Tile column of each entity. Shape (E,), int16.

Type:

jax.jaxlib._jax.Array

ent_type

Machine kind per entity. Shape (E,), int8.

Type:

jax.jaxlib._jax.Array

ent_direction

Facing of each entity, as a Direction value, 0 when nothing has set it. Shape (E,), int8. A crossing packs both its axes into this one byte, and factoriax.engine.tables.CROSSING_AXIS_DIRS unpacks them.

Type:

jax.jaxlib._jax.Array

ent_power

Steps left in the current craft of the machine. The value counts down. Shape (E,), int16. The craft finishes on the step where the value reaches 1, and 0 means idle. Only assemblers and furnaces use this field. Every other kind holds 0.

Type:

jax.jaxlib._jax.Array

ent_buf_type

Item in the general-purpose buffer slot. Shape (E,), int8. This slot holds the output of a miner, the storage of a pallet, and the contents of a belt.

Type:

jax.jaxlib._jax.Array

ent_buf_count

Items in that buffer. Shape (E,), int16. factoriax.engine.tables.MACHINE_MAX_STACK gives the limit for each kind.

Type:

jax.jaxlib._jax.Array

ent_asm_in_type

Item in each of the two input slots. Shape (E, 2), int8. Assemblers, furnaces, and science labs use these slots. A crossing uses the same pair as one buffer for each axis, at the indexes CROSSING_VERT_SLOT and CROSSING_HORIZ_SLOT.

Type:

jax.jaxlib._jax.Array

ent_asm_in_count

Items in each input slot, in the same order as ent_asm_in_type. Shape (E, 2), int16.

Type:

jax.jaxlib._jax.Array

ent_asm_out_type

Item in the finished-output slot. Shape (E,), int8. A finished craft stays here, and nothing writes over it. The machine therefore stops until something withdraws the item.

Type:

jax.jaxlib._jax.Array

ent_asm_out_count

Items in the output slot. Shape (E,), int16.

Type:

jax.jaxlib._jax.Array

ent_health

Hit points of each entity. Shape (E,), int16. Placement sets MACHINE_MAX_HEALTH. A repair clamps back to that value, and the engine refuses a pickup below it.

Type:

jax.jaxlib._jax.Array

player_positions

(x, y) tile of each player. Shape (P, 2), int16.

Type:

jax.jaxlib._jax.Array

player_directions

Facing of each player, a Direction value. Shape (P,), int8.

Type:

jax.jaxlib._jax.Array

player_inventory

Items that each player carries, by player and item id. Shape (P, NUM_ITEM_TYPES), int16. Column 0 is ItemType.EMPTY and stays zero. factoriax.engine.tables.PLAYER_MAX_STACK gives the limit for each item.

Type:

jax.jaxlib._jax.Array

selected_player

Index of the player that an action applies to. Scalar int32. The single-agent step sets this field.

Type:

int

timestep

Steps since the last reset. Scalar int32. The episode ends when this value reaches EnvParams.max_timesteps.

Type:

int

items_mined

Total mined for each item id since the last reset. Shape (NUM_ITEM_TYPES,), int32. Nothing clears this field during an episode. Reward functions read the difference between two states, not the total.

Type:

jax.jaxlib._jax.Array

science_consumed_step

Science packs that the last step consumed, by position in SCIENCE_PACK_TYPES. Shape (NUM_SCIENCE_PACK_TYPES,), int32. This field is a per-step delta, not a total.

Type:

jax.jaxlib._jax.Array

achievements_unlocked

One latched bit for each achievement. Shape (MAX_ACHIEVEMENTS,), bool. The engine folds a bit in with OR when its condition first holds, and the bit stays set for the rest of the episode. A scenario with fewer achievements leaves the bits at the end False.

Type:

jax.jaxlib._jax.Array

Examples

>>> import jax
>>> from factoriax.engine.constants import NUM_ITEM_TYPES
>>> from factoriax.make import env_from_name
>>> env, params = env_from_name("EasyRocket-v1", auto_reset=False)
>>> _, state = env.reset_env(jax.random.PRNGKey(0), params)
>>> state.player_inventory.shape == (1, NUM_ITEM_TYPES)
True

replace(**updates)

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