States

Event

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, whether the event happened, and the type of event that 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.

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.
event_type	Array	The type of event that happened.

EventType

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

EventsManager

Bases: PyTreeNode

A struct that manages the events. It contains the different events that can happen in the environment, such as the goal being reached, the player being hit by a ball, etc.

Attributes:

Name	Type	Description
goal_reached	Event	An event indicating that the goal has been reached.
ball_hit	Event	An event indicating that the player has been hit by a ball.
wall_hit	Event	An event indicating that the player has hit a wall.
lava_fall	Event	An event indicating that the lava has fallen.
key_pickup	Event	An event indicating that the player has picked up a key.
door_opening	Event	An event indicating that the player has opened a door.
door_unlock	Event	An event indicating that the player has unlocked a door.
ball_pickup	Event	An event indicating that the player has picked up a ball.

record_ball_hit(ball)

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	The ball that hit the player.	required

Returns:

Name	Type	Description
EventsManager	EventsManager	The updated events manager.

record_ball_pickup(ball, position)

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	The ball the player picked up.	required
position	Array	The position of the ball in the grid.	required

Returns:

Name	Type	Description
EventsManager	EventsManager	The updated events manager.

record_door_opening(door, position)

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

Parameters:

Name	Type	Description	Default
door	Door	The door the player opened.	required
position	Array	The position of the door in the grid.	required

Returns:

Name	Type	Description
EventsManager	EventsManager	The updated events manager.

record_door_unlock(door, position)

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

Parameters:

Name	Type	Description	Default
door	Door	The door the player unlocked.	required
position	Array	The position of the door in the grid.	required

Returns:

Name	Type	Description
EventsManager	EventsManager	The updated events manager.

record_goal_reached(goal, position)

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

Parameters:

Name	Type	Description	Default
goal	Goal	The goal the player reached.	required
position	Array	The position of the goal in the grid.	required

Returns:

Name	Type	Description
EventsManager	EventsManager	The updated events manager.

record_grid_hit(position)

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

Parameters:

Name	Type	Description	Default
position	Array	The position of the wall in the grid.	required

Returns:

Name	Type	Description
EventsManager	EventsManager	The updated events manager.

record_key_pickup(key, position)

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	The key the player picked up.	required
position	Array	The position of the key in the grid.	required

Returns:

Name	Type	Description
EventsManager	EventsManager	The updated events manager.

record_lava_fall(lava, position)

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

Parameters:

Name	Type	Description	Default
lava	Lava	The lava that fell.	required
position	Array	The position of the lava in the grid.	required

Returns:

Name	Type	Description
EventsManager	EventsManager	The updated events manager.

record_pickup(entity, position)

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.

record_walk_into(entity, position)

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.

record_wall_hit(wall, position)

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

Parameters:

Name	Type	Description	Default
wall	Wall	The wall the player hit.	required
position	Array	The position of the wall in the grid.	required

Returns:

Name	Type	Description
EventsManager	EventsManager	The updated events manager.

State

Bases: PyTreeNode

The Markovian state of the environment

cache: RenderingCache instance-attribute

The rendering cache to speed up rendering

entities: Dict[str, Entity] = struct.field(default_factory=dict) class-attribute instance-attribute

The entities in the environment, indexed via entity type string representation. Batched over the number of entities for each type

events: EventsManager = EventsManager() class-attribute instance-attribute

A struct indicating which events happened this timestep. For example, the goal is reached, or the player is hit by a ball.

grid: Array instance-attribute

The base map of the environment that remains constant throughout the training

key: Array instance-attribute

The random number generator state

get_balls()

Gets the ball entity from the state.

get_boxes()

Gets the box entity from the state.

get_doors()

Gets the door entity from the state.

get_entity(entity_enum)

Get an entity from the state by its enum.

Parameters:

Name	Type	Description	Default
entity_enum	str	The enum of the entity to get.	required

Returns:

Name	Type	Description
Entity	Entity	The entity from the state.

get_goals()

Gets the goal entity from the state.

get_keys()

Gets the key entity from the state.

get_lavas()

Gets the lava entity from the state.

get_player(idx=0)

Gets the player entity from the state.

get_positions()

Get the positions of all the entities in the state.

get_sprites()

Get the sprites of all the entities in the state.

get_tags()

Get the tags of all the entities in the state.

get_transparency()

Get the transparency of all the entities in the state.

get_walls()

Gets all the WALL entities from the state.

set_balls(balls)

Sets the ball entity in the state.

set_boxes(boxes)

Sets the box entity in the state.

set_doors(doors)

Sets the door entity in the state.

set_entity(entity_enum, entity)

Set an entity in the state by its enum.

Parameters:

Name	Type	Description	Default
entity_enum	str	The enum of the entity to set.	required
entity	Entity	The entity to set.	required

Returns:

Name	Type	Description
State	State	The updated state.

set_events(events)

Sets the events in the state.

set_goals(goals)

Sets the goal entity in the state.

set_keys(keys)

Sets the key entity in the state.

set_player(player, idx=0)

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.

set_walls(walls)

Sets the WALL entities in the state.