Skip to content

States

State - the full, true world state - and the per-step Event / EventsManager record it carries.

State bundles the static grid, the batched entities dict, a rendering cache and a PRNG key; an observation_fn derives what the agent sees from it. EventsManager is a small fixed set of Event slots that the transition pipeline writes ("goal reached", "ball hit", ...) and that navix.events / rewards / terminations read.

Pseudo entity-type key for a wall-hit against the grid boundary or a non-walkable empty cell, as opposed to hitting an actual Wall entity - EventsManager.record_grid_hit/record_wall_hit write to separate (GRID, EventType.HIT)/(Entities.WALL, EventType.HIT) slots so they never fight over the same one; navix.events.on_wall_hit ORs both together, so the distinction is invisible to callers that only care whether a wall was hit at all.

Bases: Positionable, HasColour

A struct representing an event that happened in the environment. It contains the position of the event, the colour of the entity involved in the event, and whether the event happened.

Note

Notice that we need the happened property, which flags if an event has happened or not, because JAX does not support variable size arrays. This means that we cannot add an event to the list in the middle of training. Instead, we initialise all events, and mask them out as not happened.

Every field of an Event stored in EventsManager.events is batched to match the entity type it tracks (see empty_like) - one slot per instance, not one scalar per event type - so e.g. two balls hitting the player the same step are two independent True entries, not collapsed into one (see issue #139).

Attributes:

Name Type Description
position Array

The (row, column) position of the event in the grid.
colour Array

The colour of the entity involved in the event.
happened Array

A boolean flag indicating whether the event happened.

An all-happened=False Event, batched to match entity's own shape (one slot per instance of entity) rather than a single scalar slot.

Parameters:

Name Type Description Default
entity Entity

The entity type this event tracks - e.g. state.entities[Entities.BALL].

 required

Returns:

Name Type Description
Event Event

position/colour/happened, each broadcast to
Event

entity.shape (plus position's own trailing (2,)).

Enumeration of the different types of events that can happen in the environment.

Plain strings, not jax arrays - EventsManager.events keys off (entity_type, event_type) tuples, and a jax pytree's dict keys must be static, hashable Python values, not traced arrays.

Bases: PyTreeNode

A struct that manages the events that happened in the environment this step, such as the goal being reached, the player being hit by a ball, etc.

Keyed by (entity_type, event_type) rather than one named field per event type, and each slot's Event is batched to match that entity type's own instance count (see Event.empty_like) - this is what lets two independent occurrences of the same event type (e.g. two balls both hitting the player) coexist within one step, rather than one replacing the other (see issue #139). record_* merges new hits into a slot rather than replacing it wholesale, so a call earlier in the same step's transition pipeline (e.g. the player walking into a ball) isn't discarded by a later one (e.g. a different ball moving onto the player) writing the same slot.

Attributes:

Name Type Description
events Dict[Tuple[str, str], Event]

One Event slot per (entity_type, event_type) this environment's entities can actually produce - see create. Not meant to be constructed directly; State.__post_init__ calls create automatically from State.entities.

Builds one all-happened=False Event slot per (entity_type, event_type) this environment's entities can actually produce - e.g. (Entities.BALL, EventType.HIT) only exists if Entities.BALL in entities. Called automatically by State.__post_init__; not meant to be called directly.

Parameters:

Name Type Description Default
entities Dict[str, Entity]

State.entities.

 required

Returns:

Name Type Description
EventsManager EventsManager

A fresh, all-happened=False manager.

Whether any instance of key's (entity_type, event_type) slot fired this step - False (not a KeyError) if this environment's entities never included that entity type at all (see create).

Parameters:

Name Type Description Default
key Tuple[str, str]

The (entity_type, event_type) slot to check.

 required

Returns:

Name Type Description
Array Array

A boolean scalar.

Whether key's slot fired this step for the specific instance that was at position when recorded - Event.position keeps the firing instance's own position even after game logic later moves the entity itself (e.g. to the discard pile on pickup), so this can identify which instance fired, not just whether any did (see happened). False (not a KeyError) if this environment's entities never included that entity type at all.

Parameters:

Name Type Description Default
key Tuple[str, str]

The (entity_type, event_type) slot to check.

 required
position Array

The position to match against.

 required

Returns:

Name Type Description
Array Array

A boolean scalar.

Records hit into key's Event slot, OR-merging onto whatever already happened this step rather than replacing it wholesale - so a call earlier in the same step's transition pipeline isn't silently discarded by a later one writing the same slot (see issue #139). position/colour are written only where hit is newly True; already-True entries keep their existing stored position/colour.

Parameters:

Name Type Description Default
key Tuple[str, str]

The (entity_type, event_type) slot to update - must already exist (see create).

 required
hit Array

Boolean - per-instance for a batched slot (e.g. Entities.BALL, EventType.HIT), a bare scalar for a single-instance slot (e.g. GRID, EventType.HIT).

 required
position Array

This instance's own position - written where hit is True.

 required
colour Array

This instance's own colour - written where hit is True.

 required

Returns:

Name Type Description
EventsManager EventsManager

The updated events manager.

Flags an event when the player is hit by a ball as happened and returns the updated events manager.

Parameters:

Name Type Description Default
ball Ball

Every Ball instance in the environment.

 required
hit Array

Boolean, one entry per ball instance.

 required

Returns:

Name Type Description
EventsManager EventsManager

The updated events manager.

Flags an event when the player picks up a ball as happened and returns the updated events manager.

Parameters:

Name Type Description Default
ball Ball

Every Ball instance in the environment.

 required
hit Array

Boolean, one entry per ball instance.

 required

Returns:

Name Type Description
EventsManager EventsManager

The updated events manager.

Flags an event when the player picks up a box as happened and returns the updated events manager.

Parameters:

Name Type Description Default
box Box

Every Box instance in the environment.

 required
hit Array

Boolean, one entry per box instance.

 required

Returns:

Name Type Description
EventsManager EventsManager

The updated events manager.

Flags an event when the player opens a door as happened and returns the updated events manager.

Parameters:

Name Type Description Default
door Door

Every Door instance in the environment.

 required
hit Array

Boolean, one entry per door instance.

 required

Returns:

Name Type Description
EventsManager EventsManager

The updated events manager.

Flags an event when the player unlocks a door as happened and returns the updated events manager.

Parameters:

Name Type Description Default
door Door

Every Door instance in the environment.

 required
hit Array

Boolean, one entry per door instance.

 required

Returns:

Name Type Description
EventsManager EventsManager

The updated events manager.

Flags an event when the player reaches the goal as happened and returns the updated events manager.

Parameters:

Name Type Description Default
goal Goal

Every Goal instance in the environment.

 required
hit Array

Boolean, one entry per goal instance.

 required

Returns:

Name Type Description
EventsManager EventsManager

The updated events manager.

Flags an event when the player hits the grid boundary or a non-walkable empty cell (no Wall entity there) as happened and returns the updated events manager. Kept in a separate GRID slot from record_wall_hit (see GRID's docstring).

Parameters:

Name Type Description Default
position Array

The position hit.

 required

Returns:

Name Type Description
EventsManager EventsManager

The updated events manager.

Flags an event when the player picks up a key as happened and returns the updated events manager.

Parameters:

Name Type Description Default
key Key

Every Key instance in the environment.

 required
hit Array

Boolean, one entry per key instance.

 required

Returns:

Name Type Description
EventsManager EventsManager

The updated events manager.

Flags an event when the lava falls as happened and returns the updated events manager.

Parameters:

Name Type Description Default
lava Lava

Every Lava instance in the environment.

 required
hit Array

Boolean, one entry per lava instance.

 required

Returns:

Name Type Description
EventsManager EventsManager

The updated events manager.

Flags an event when the player picks up an entity as happened and returns the updated events manager.

Parameters:

Name Type Description Default
entity Entity

The entity the player picked up.

 required
position Array

The position of the entity in the grid.

 required

Returns:

Name Type Description
EventsManager EventsManager

The updated events manager.

Flags an event when the player walks into an entity as happened and returns the updated events manager.

Parameters:

Name Type Description Default
entity Entity

The entity the player walked into.

 required
position Array

The position of the entity in the grid.

 required

Returns:

Name Type Description
EventsManager EventsManager

The updated events manager.

Flags an event when the player hits a wall as happened and returns the updated events manager.

Parameters:

Name Type Description Default
wall Wall

Every Wall instance in the environment.

 required
hit Array

Boolean, one entry per wall instance.

 required

Returns:

Name Type Description
EventsManager EventsManager

The updated events manager.

Bases: PyTreeNode

The full, true state of the world - everything step needs to compute the next state, and everything an observation_fn reads. A frozen flax.struct pytree (vmaps over a batch of environments). Timestep.state is one of these.

The get_* / set_* helpers below are the convenient way to reach into entities; the raw entities dict is also fine to use directly.

Pre-rendered tile patches, so rgb observations only re-draw the cells that changed. Carried in the state so it survives jax.jit.

Maps an Entities key ("player", "key", ...) to a single batched Entity holding every instance of that type. A type with no instances in this environment is simply absent from the dict.

A struct indicating which events happened this timestep. For example, the goal is reached, or the player is hit by a ball. Left at its default (empty) here - __post_init__ populates it from entities via EventsManager.create, since a struct.PyTreeNode field's default can't see its sibling fields.

i32[H, W] static base map: 0 is floor, -1 marks a wall cell. Fixed for the lifetime of an episode; moving entities live in entities, not here.

PRNG key for this environment's own stochasticity (ball motion, stochastic goals). Advanced by the functions that consume it.

The environment's mission target(s), if any - e.g. (door,) in GoToDoor, (target,) in GoToObject/Fetch, (carry, drop_near) in PutNear (index 0 is always the "primary"/carry target; index 1, where present, is a second, independently-tracked target - only PutNear needs two today). Empty for an environment with no mission at all. A tuple, not a fixed number of separate fields, so a future environment needing a third simultaneous target is just a longer tuple, not another numbered field.

Gets the ball entity from the state.

Gets the box entity from the state.

Gets the door entity from the state.

The batched Entity for one type. The typed helpers (get_keys, get_doors, ...) are thin wrappers around this.

Parameters:

Name Type Description Default
entity_enum str

an Entities key, e.g. Entities.KEY.

 required

Returns:

Name Type Description
Entity Entity

every instance of that type, batched (leading axis =
Entity

instance count).

Raises:

Type Description
KeyError

if this environment has no entity of that type -

Gets the goal entity from the state.

Gets the key entity from the state.

Gets the lava entity from the state.

The player, unbatched (navix is single-agent). Unlike the other get_* helpers this indexes into the batch and returns one Player.

Parameters:

Name Type Description Default
idx int

which player - only 0 is meaningful today.

 0

Returns:

Name Type Description
Player Player

the player entity, no leading instance axis.

Every entity instance's (row, col), concatenated across all types into one i32[N, 2] (N = total instance count). The ordering matches get_tags / get_sprites / get_transparency, so they can be zipped - this is what the observation functions do to paint entities onto the grid.

Every entity instance's RGB sprite, concatenated into u8[N, TILE_SIZE, TILE_SIZE, 3], in the same order as get_positions. Used to build an rgb observation.

Like get_sprites, but the player's sprite is forced to its north-facing variant. First-person observations rotate the whole view so the player always points up, so its sprite must not also carry a rotation.

Every entity instance's tag, concatenated into i32[N], in the same order as get_positions. Used to build a categorical observation.

Every entity instance's transparent flag, concatenated into bool[N], in the same order as get_positions. Feeds the first-person view cone.

Gets all the WALL entities from the state.

Sets the ball entity in the state.

Sets the box entity in the state.

Sets the door entity in the state.

Replaces one entity type's batch. Mutates self.entities in place and returns self (it does not build a new State), so state.set_entity(...) and state = state.set_entity(...) are equivalent.

Parameters:

Name Type Description Default
entity_enum str

an Entities key.

 required
entity Entity

the new batched entity (its shape may differ from the old one - e.g. after a pickup moves an instance off-grid).

 required

Returns:

Name Type Description
State State

self.

Sets the events in the state.

Sets the goal entity in the state.

Sets the key entity in the state.

Sets the player entity in the state. Notice that we only support one player in the environment for now, but this can easily be extended to multiple players.

Sets the WALL entities in the state.