Pong Environment

class jaxatari.games.jax_pong.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_pong.JaxPong(consts: PongConstants | None = None)[source]

Bases: JaxEnvironment[PongState, PongObservation, PongInfo, PongConstants]

ACTION_SET: Array = Array([ 0, 1, 3, 4, 11, 12], dtype=int32)

action_space() → Discrete[source]: Returns the action space of the environment as an array containing the actions that can be taken. Returns: The action space of the environment as an array.

image_space() → Box[source]: Returns the image space of the environment. Returns: The image space of the environment.

obs_to_flat_array(obs: PongObservation) → Array[source]

observation_space() → <module 'jaxatari.spaces' from '/home/docs/checkouts/readthedocs.org/user_builds/jaxatari/envs/stable/lib/python3.10/site-packages/jaxatari/spaces.py'>[source]: Returns the observation space of the environment. Returns: The observation space of the environment.

render(state: PongState) → Array[source]

Renders the environment state to a single image. :param state: The environment state.

Returns: A single image of the environment state.

reset(key: Array = Array([0, 42], dtype=uint32)) → Tuple[PongObservation, PongState][source]: Resets the environment to the initial state. Returns: The initial observation and the initial environment state.

step(state: PongState, action: Array | ndarray | bool_ | number) → Tuple[PongObservation, PongState, float, bool, PongInfo][source]

class jaxatari.games.jax_pong.PongConstants(MAX_SPEED, BALL_SPEED, ENEMY_STEP_SIZE, WIDTH, HEIGHT, BASE_BALL_SPEED, BALL_MAX_SPEED, MIN_BALL_SPEED, PLAYER_ACCELERATION, BALL_START_X, BALL_START_Y, BACKGROUND_COLOR, PLAYER_COLOR, ENEMY_COLOR, BALL_COLOR, WALL_COLOR, SCORE_COLOR, PLAYER_X, ENEMY_X, PLAYER_SIZE, BALL_SIZE, ENEMY_SIZE, WALL_TOP_Y, WALL_TOP_HEIGHT, WALL_BOTTOM_Y, WALL_BOTTOM_HEIGHT, ASSET_CONFIG)[source]

Bases: NamedTuple

ASSET_CONFIG: tuple: Alias for field number 26

BACKGROUND_COLOR: Tuple[int, int, int]: Alias for field number 11

BALL_COLOR: Tuple[int, int, int]: Alias for field number 14

BALL_MAX_SPEED: int: Alias for field number 6

BALL_SIZE: Tuple[int, int]: Alias for field number 20

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

BALL_START_X: Array | ndarray | bool_ | number: Alias for field number 9

BALL_START_Y: Array | ndarray | bool_ | number: Alias for field number 10

BASE_BALL_SPEED: int: Alias for field number 5

ENEMY_COLOR: Tuple[int, int, int]: Alias for field number 13

ENEMY_SIZE: Tuple[int, int]: Alias for field number 21

ENEMY_STEP_SIZE: int: Alias for field number 2

ENEMY_X: int: Alias for field number 18

HEIGHT: int: Alias for field number 4

MAX_SPEED: int: Alias for field number 0

MIN_BALL_SPEED: int: Alias for field number 7

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

PLAYER_COLOR: Tuple[int, int, int]: Alias for field number 12

PLAYER_SIZE: Tuple[int, int]: Alias for field number 19

PLAYER_X: int: Alias for field number 17

SCORE_COLOR: Tuple[int, int, int]: Alias for field number 16

WALL_BOTTOM_HEIGHT: int: Alias for field number 25

WALL_BOTTOM_Y: int: Alias for field number 24

WALL_COLOR: Tuple[int, int, int]: Alias for field number 15

WALL_TOP_HEIGHT: int: Alias for field number 23

WALL_TOP_Y: int: Alias for field number 22

WIDTH: int: Alias for field number 3

class jaxatari.games.jax_pong.PongInfo(time)[source]

Bases: NamedTuple

time: Array: Alias for field number 0

class jaxatari.games.jax_pong.PongObservation(player, enemy, ball, score_player, score_enemy)[source]

Bases: NamedTuple

ball: EntityPosition: Alias for field number 2

enemy: EntityPosition: Alias for field number 1

player: EntityPosition: Alias for field number 0

score_enemy: Array: Alias for field number 4

score_player: Array: Alias for field number 3

class jaxatari.games.jax_pong.PongRenderer(consts: PongConstants | None = None)[source]

Bases: JAXGameRenderer

render(state)[source]

class jaxatari.games.jax_pong.PongState(player_y, player_speed, ball_x, ball_y, enemy_y, enemy_speed, ball_vel_x, ball_vel_y, player_score, enemy_score, step_counter, acceleration_counter, buffer, key)[source]

Bases: NamedTuple

acceleration_counter: Array | ndarray | bool_ | number: Alias for field number 11

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

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

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

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

buffer: Array | ndarray | bool_ | number: Alias for field number 12

enemy_score: Array | ndarray | bool_ | number: Alias for field number 9

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

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

key: Array: Alias for field number 13

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

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

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

step_counter: Array | ndarray | bool_ | number: Alias for field number 10