Freeway Environment

class jaxatari.games.jax_freeway.EntityPosition(x, y, width, height)[source]

Bases: NamedTuple

height: Array: Alias for field number 3

width: Array: Alias for field number 2

x: Array: Alias for field number 0

y: Array: Alias for field number 1

class jaxatari.games.jax_freeway.FreewayConstants(screen_width, screen_height, chicken_width, chicken_height, chicken_x, car_width, car_height, num_lanes, lane_spacing, car_speeds, lane_borders, top_border, top_path, bottom_border, throw_back_frames, stun_frames, post_score_stun_frames, post_score_spawn_offset_y, chicken_hit_inset_x, chicken_hit_inset_y_top, chicken_hit_inset_y_bottom, car_hit_inset_x, car_hit_inset_y_top, car_hit_inset_y_bottom, respawn_offset, car_y_offset, cadence_phase_offset, CAR_UPDATES, lane_phase_offset, CAR_COLORS, game_duration_frames, blink_start_frames, score_blink_rate, SCORE_BLINK_COLORS, ASSET_CONFIG)[source]

Bases: NamedTuple

ASSET_CONFIG: tuple: Alias for field number 34

CAR_COLORS: List[Tuple[int, int, int] | None]: Alias for field number 29

CAR_UPDATES: List[int]: Alias for field number 27

SCORE_BLINK_COLORS: List[Tuple[int, int, int]]: Alias for field number 33

blink_start_frames: int: Alias for field number 31

bottom_border: int: Alias for field number 13

cadence_phase_offset: List[int]: Alias for field number 26

car_height: int: Alias for field number 6

car_hit_inset_x: int: Alias for field number 21

car_hit_inset_y_bottom: int: Alias for field number 23

car_hit_inset_y_top: int: Alias for field number 22

car_speeds: List[float]: Alias for field number 9

car_width: int: Alias for field number 5

car_y_offset: int: Alias for field number 25

chicken_height: int: Alias for field number 3

chicken_hit_inset_x: int: Alias for field number 18

chicken_hit_inset_y_bottom: int: Alias for field number 20

chicken_hit_inset_y_top: int: Alias for field number 19

chicken_width: int: Alias for field number 2

chicken_x: int: Alias for field number 4

game_duration_frames: int: Alias for field number 30

lane_borders: List[int]: Alias for field number 10

lane_phase_offset: List[int]: Alias for field number 28

lane_spacing: int: Alias for field number 8

num_lanes: int: Alias for field number 7

post_score_spawn_offset_y: int: Alias for field number 17

post_score_stun_frames: int: Alias for field number 16

respawn_offset: int: Alias for field number 24

score_blink_rate: int: Alias for field number 32

screen_height: int: Alias for field number 1

screen_width: int: Alias for field number 0

stun_frames: int: Alias for field number 15

throw_back_frames: int: Alias for field number 14

top_border: int: Alias for field number 11

top_path: int: Alias for field number 12

class jaxatari.games.jax_freeway.FreewayInfo(time)[source]

Bases: NamedTuple

time: Array: Alias for field number 0

class jaxatari.games.jax_freeway.FreewayObservation(chicken, car)[source]

Bases: NamedTuple

car: Array: Alias for field number 1

chicken: EntityPosition: Alias for field number 0

class jaxatari.games.jax_freeway.FreewayRenderer(consts: FreewayConstants | None = None)[source]

Bases: JAXGameRenderer

render(state)[source]

Bases: NamedTuple

Represents the current state of the game

cars: Array | ndarray | bool_ | number: Alias for field number 1

chicken_y: Array | ndarray | bool_ | number: Alias for field number 0

cooldown: Array | ndarray | bool_ | number: Alias for field number 5

game_over: Array | ndarray | bool_ | number: Alias for field number 8

lane_time: Array | ndarray | bool_ | number: Alias for field number 2

lives_lost: Array | ndarray | bool_ | number: Alias for field number 7

score: Array | ndarray | bool_ | number: Alias for field number 3

time: Array | ndarray | bool_ | number: Alias for field number 4

walking_frames: Array | ndarray | bool_ | number: Alias for field number 6

class jaxatari.games.jax_freeway.JaxFreeway(consts: FreewayConstants | None = None)[source]

Bases: JaxEnvironment[FreewayState, FreewayObservation, FreewayInfo, FreewayConstants]

ACTION_SET: Array = Array([0, 2, 5], dtype=int32)

action_space() → Discrete[source]: Returns the action space for Freeway.

image_space() → Box[source]: Returns the image space for Freeway. The image is a RGB image with shape (210, 160, 3).

obs_to_flat_array(obs: FreewayObservation) → Array[source]: Convert observation to a flat array.

observation_space() → Dict[source]: Returns the observation space for Freeway. The observation contains: - chicken: EntityPosition (x, y, width, height) - car: array of shape (10, 4) with x,y,width,height for each car

render(state: FreewayState) → Array[source]: Render the game state to a raster image.

reset(key: PRNGKey | None = None) → Tuple[FreewayObservation, FreewayState][source]: Initialize a new game state

step(state: FreewayState, action: int) → tuple[FreewayObservation, FreewayState, float, bool, FreewayInfo][source]: Take a step in the game given an action