Available Environments

Jaxnasium doesn't bundle a large number of environments directly. Instead, it relies on existing work and wraps these environments in various wrappers to conform to the Jaxnasium API. This approach allows users to leverage a wide array of established environments while maintaining the performance and flexibility offered by JAX.

Native Environments

For convenience, the five classic control environments are natively implemented and bundled in Jaxnasium, and require no external dependencies:

• CartPole-v1
• MountainCar-v0
• Acrobot-v1
• Pendulum-v1
• MountainCarContinuous-v0

External Environments using the Jaxnasium API

These environments run without wrappers.

TODO

External Environment Libraries

Jaxnasium integrates with the following external environment libraries through wrapper adapters. See the end of this page for a full list of available environments. These are not bundled as dependencies and need to be installed manually (e.g. via pip) before use.

Gymnax

JAX implementations of OpenAI's Gym environments, offering accelerated and parallelized rollouts. Includes classic control, bsuite, and MinAtar environments.

Jumanji

A suite of diverse, scalable reinforcement learning environments implemented in JAX by DeepMind. Focuses on combinatorial problems and general decision-making tasks.

Brax

A fast and flexible physics simulation engine for training and evaluating rigid body environments in JAX by Google.

Crax

A GPU-accelerated benchmark and algorithm suite for Safe RL, built as a fork of Brax. Alongside the reward, every environment reports a cost signal in the timestep info.

The safety suites are parameterised by a difficulty level; jaxnasium defaults to the easiest, so pass level for the harder variants:

env = jym.make("safe_lift_ant", level=3)

Playground

MuJoCo Playground by DeepMind: MJX-accelerated robotics environments, spanning the DM Control Suite, locomotion tasks for a range of quadrupeds and humanoids, and dexterous manipulation.

Pgx

JAX implementations of various board games and classic environments, including chess, Go, shogi, and more.

JaxMARL

Multi-agent reinforcement learning environments implemented in JAX, including MPE (Multi-Particle Environment) scenarios and other multi-agent tasks.

xMinigrid

JAX implementation of MiniGrid environments, including XLand variants for procedural generation research.

Navix

JAX implementation of navigation environments, providing various gridworld navigation tasks.

Octax

JAX CHIP-8 emulator environments (Brix, Pong, Tetris, and other classic games).

Craftax

JAX implementation of Craftax environments, inspired by Minecraft-like crafting and survival tasks.

Factoriax

A JAX grid world of Factorio-style crafting and logistics: mine ore, place machines, and work up a tech tree towards launching a rocket. Uses the Gymnax API.

PopJym

JAX implementation of POPGym, a suite of partially observable environments. Generally requires a recurrent or memory-based policy.

Usage

To create any of these environments, simply pass its id to jaxnasium.make:

import jaxnasium as jym

# Native environments
env = jym.make("CartPole-v1")

# External environments (requires installing the respective library)
env = jym.make("Breakout-MinAtar")  # Gymnax
env = jym.make("Game2048-v1")  # Jumanji
env = jym.make("ant")  # Brax
env = jym.make("chess")  # Pgx
env = jym.make("octax:brix")  # Octax

Every id can also be given with an explicit provider prefix (e.g. gymnax:Breakout-MinAtar), which is useful when the same name exists in more than one library.

By default, make applies the wrapper for the environment's library (translating it to the Jaxnasium API) and a LogWrapper. Pass wrappers=[...] to control this.

jaxnasium._registry.Registry dataclass

__init__(_environments: dict[str, type[Environment]] = dict(), _aliases: dict[str, str] = dict()) -> None

register(id: str, **kwargs)

Register an environment with the registry.

Arguments: id: The environment ID (e.g., "CartPole-v1") entry_point: Optional entry point string for lazy loading **kwargs: currently unused

register_alias(alias: str, target: str)

Register an alias for an environment.

Arguments: alias: The alias to register target: The target environment ID

make(id: str, wrappers: list[type[Wrapper] | Literal['external_lib_wrapper']] | None = None, **env_kwargs) -> Environment

Create an environment instance.

Arguments: id: The environment ID wrappers: List of wrappers to apply to the environment. can also be "external_lib_wrapper" (string): Wrapper for external libraries (e.g. Gymnax, Jumanji, Brax). only used if a environment is loaded from a supported external library. If None provided, we wrap the external_lib_wrapper (if applicable) and LogWrapper. **env_kwargs: Environment constructor arguments

get_env_class(id: str) -> type[Environment]

Get the environment class for an environment ID.

Arguments: id: The environment ID

registered_envs: list[str] property

List all environments in the registry as a flat list.

print_envs() -> None

Pretty prints the available environments in the registry.

options: heading_level: 3 members: - make - register - register_alias - registered_envs - print_envs

Complete List of Registered Environments

Below is the complete list of all registered environment ids available in Jaxnasium.

Auto-generated List

This list is automatically generated from the Jaxnasium registry.

• Acrobot-v1
• CartPole-v1
• MountainCarContinuous-v0
• MountainCar-v0
• Pendulum-v1
• gymnax:CartPole-v1
• gymnax:Acrobot-v1
• gymnax:Pendulum-v1
• gymnax:MountainCar-v0
• gymnax:MountainCarContinuous-v0
• gymnax:Asterix-MinAtar
• gymnax:Breakout-MinAtar
• gymnax:Freeway-MinAtar
• gymnax:SpaceInvaders-MinAtar
• gymnax:DeepSea-bsuite
• gymnax:Catch-bsuite
• gymnax:MemoryChain-bsuite
• gymnax:UmbrellaChain-bsuite
• gymnax:DiscountingChain-bsuite
• gymnax:MNISTBandit-bsuite
• gymnax:SimpleBandit-bsuite
• gymnax:FourRooms-misc
• gymnax:MetaMaze-misc
• gymnax:PointRobot-misc
• gymnax:BernoulliBandit-misc
• gymnax:GaussianBandit-misc
• gymnax:Reacher-misc
• gymnax:Swimmer-misc
• gymnax:Pong-misc
• jumanji:Game2048-v1
• jumanji:GraphColoring-v1
• jumanji:Minesweeper-v0
• jumanji:RubiksCube-v0
• jumanji:RubiksCube-partly-scrambled-v0
• jumanji:SlidingTilePuzzle-v0
• jumanji:Sudoku-v0
• jumanji:Sudoku-very-easy-v0
• jumanji:BinPack-v2
• jumanji:FlatPack-v0
• jumanji:JobShop-v0
• jumanji:Knapsack-v1
• jumanji:Tetris-v0
• jumanji:Cleaner-v0
• jumanji:Connector-v3
• jumanji:CVRP-v1
• jumanji:MultiCVRP-v0
• jumanji:Maze-v0
• jumanji:RobotWarehouse-v0
• jumanji:Snake-v1
• jumanji:TSP-v1
• jumanji:MMST-v0
• jumanji:PacMan-v1
• jumanji:Sokoban-v0
• jumanji:LevelBasedForaging-v0
• jumanji:SearchAndRescue-v0
• brax:ant
• brax:halfcheetah
• brax:hopper
• brax:humanoid
• brax:humanoidstandup
• brax:inverted_pendulum
• brax:inverted_double_pendulum
• brax:pusher
• brax:reacher
• brax:walker2d
• crax:safe_button_point
• crax:safe_circle_point
• crax:safe_goal_point
• crax:safe_height_humanoid
• crax:safe_lift_ant
• crax:safe_lift_humanoid
• crax:safe_lift_spider
• crax:safe_pathway_walker2d
• crax:safe_push_point
• crax:safe_reacher
• crax:safe_velocity_ant
• crax:safe_velocity_halfcheetah
• crax:safe_velocity_humanoid
• crax:safe_velocity_hopper
• crax:safe_velocity_swimmer
• crax:safe_velocity_walker2d
• pgx:2048
• pgx:animal_shogi
• pgx:backgammon
• pgx:chess
• pgx:connect_four
• pgx:gardner_chess
• pgx:go_9x9
• pgx:go_19x19
• pgx:hex
• pgx:kuhn_poker
• pgx:leduc_holdem
• pgx:minatar-asterix
• pgx:minatar-breakout
• pgx:minatar-freeway
• pgx:minatar-seaquest
• pgx:minatar-space_invaders
• pgx:othello
• pgx:shogi
• pgx:sparrow_mahjong
• pgx:tic_tac_toe
• jaxmarl:MPE_simple_v3
• jaxmarl:MPE_simple_tag_v3
• jaxmarl:MPE_simple_world_comm_v3
• jaxmarl:MPE_simple_spread_v3
• jaxmarl:MPE_simple_crypto_v3
• jaxmarl:MPE_simple_speaker_listener_v4
• jaxmarl:MPE_simple_push_v3
• jaxmarl:MPE_simple_adversary_v3
• jaxmarl:MPE_simple_reference_v3
• jaxmarl:MPE_simple_facmac_v1
• jaxmarl:MPE_simple_facmac_3a_v1
• jaxmarl:MPE_simple_facmac_6a_v1
• jaxmarl:MPE_simple_facmac_9a_v1
• jaxmarl:switch_riddle
• jaxmarl:SMAX
• jaxmarl:HeuristicEnemySMAX
• jaxmarl:ant_4x2
• jaxmarl:halfcheetah_6x1
• jaxmarl:hopper_3x1
• jaxmarl:humanoid_9|8
• jaxmarl:walker2d_2x3
• jaxmarl:storm
• jaxmarl:storm_2p
• jaxmarl:storm_np
• jaxmarl:hanabi
• jaxmarl:overcooked
• jaxmarl:overcooked_v2
• jaxmarl:coin_game
• jaxmarl:jaxnav
• xminigrid:XLand-MiniGrid-R1-9x9
• xminigrid:XLand-MiniGrid-R1-11x11
• xminigrid:XLand-MiniGrid-R1-13x13
• xminigrid:XLand-MiniGrid-R1-15x15
• xminigrid:XLand-MiniGrid-R1-17x17
• xminigrid:XLand-MiniGrid-R2-9x9
• xminigrid:XLand-MiniGrid-R2-11x11
• xminigrid:XLand-MiniGrid-R2-13x13
• xminigrid:XLand-MiniGrid-R2-15x15
• xminigrid:XLand-MiniGrid-R2-17x17
• xminigrid:XLand-MiniGrid-R4-9x9
• xminigrid:XLand-MiniGrid-R4-11x11
• xminigrid:XLand-MiniGrid-R4-13x13
• xminigrid:XLand-MiniGrid-R4-15x15
• xminigrid:XLand-MiniGrid-R4-17x17
• xminigrid:XLand-MiniGrid-R6-13x13
• xminigrid:XLand-MiniGrid-R6-17x17
• xminigrid:XLand-MiniGrid-R6-19x19
• xminigrid:XLand-MiniGrid-R9-16x16
• xminigrid:XLand-MiniGrid-R9-19x19
• xminigrid:XLand-MiniGrid-R9-25x25
• xminigrid:MiniGrid-BlockedUnlockPickUp
• xminigrid:MiniGrid-DoorKey-5x5
• xminigrid:MiniGrid-DoorKey-6x6
• xminigrid:MiniGrid-DoorKey-8x8
• xminigrid:MiniGrid-DoorKey-16x16
• xminigrid:MiniGrid-Empty-5x5
• xminigrid:MiniGrid-Empty-6x6
• xminigrid:MiniGrid-Empty-8x8
• xminigrid:MiniGrid-Empty-16x16
• xminigrid:MiniGrid-EmptyRandom-5x5
• xminigrid:MiniGrid-EmptyRandom-6x6
• xminigrid:MiniGrid-EmptyRandom-8x8
• xminigrid:MiniGrid-EmptyRandom-16x16
• xminigrid:MiniGrid-FourRooms
• xminigrid:MiniGrid-LockedRoom
• xminigrid:MiniGrid-MemoryS8
• xminigrid:MiniGrid-MemoryS16
• xminigrid:MiniGrid-MemoryS32
• xminigrid:MiniGrid-MemoryS64
• xminigrid:MiniGrid-MemoryS128
• xminigrid:MiniGrid-Playground
• xminigrid:MiniGrid-Unlock
• xminigrid:MiniGrid-UnlockPickUp
• navix:Navix-Empty-5x5-v0
• navix:Navix-Empty-6x6-v0
• navix:Navix-Empty-8x8-v0
• navix:Navix-Empty-16x16-v0
• navix:Navix-Empty-Random-5x5-v0
• navix:Navix-Empty-Random-6x6-v0
• navix:Navix-Empty-Random-8x8-v0
• navix:Navix-Empty-Random-16x16-v0
• navix:Navix-DoorKey-5x5-v0
• navix:Navix-DoorKey-6x6-v0
• navix:Navix-DoorKey-8x8-v0
• navix:Navix-DoorKey-16x16-v0
• navix:Navix-DoorKey-Random-5x5-v0
• navix:Navix-DoorKey-Random-6x6-v0
• navix:Navix-DoorKey-Random-8x8-v0
• navix:Navix-DoorKey-Random-16x16-v0
• navix:Navix-FourRooms-v0
• navix:Navix-KeyCorridorS3R1-v0
• navix:Navix-KeyCorridorS3R2-v0
• navix:Navix-KeyCorridorS3R3-v0
• navix:Navix-KeyCorridorS4R3-v0
• navix:Navix-KeyCorridorS5R3-v0
• navix:Navix-KeyCorridorS6R3-v0
• navix:Navix-LavaGapS5-v0
• navix:Navix-LavaGapS6-v0
• navix:Navix-LavaGapS7-v0
• navix:Navix-SimpleCrossingS9N1-v0
• navix:Navix-SimpleCrossingS9N2-v0
• navix:Navix-SimpleCrossingS9N3-v0
• navix:Navix-SimpleCrossingS11N5-v0
• navix:Navix-Dynamic-Obstacles-5x5-v0
• navix:Navix-Dynamic-Obstacles-5x5-Random-v0
• navix:Navix-Dynamic-Obstacles-6x6-v0
• navix:Navix-Dynamic-Obstacles-6x6-Random-v0
• navix:Navix-Dynamic-Obstacles-8x8-v0
• navix:Navix-Dynamic-Obstacles-16x16-v0
• navix:Navix-DistShift1-v0
• navix:Navix-DistShift2-v0
• navix:Navix-GoToDoor-5x5-v0
• navix:Navix-GoToDoor-6x6-v0
• navix:Navix-GoToDoor-8x8-v0
• octax:airplane
• octax:blinky
• octax:brix
• octax:cavern1
• octax:cavern2
• octax:cavern3
• octax:cavern5
• octax:cavern6
• octax:deep
• octax:filter
• octax:flight_runner
• octax:missile
• octax:pong
• octax:rocket
• octax:shooting_stars
• octax:space_flight1
• octax:space_flight2
• octax:space_flight3
• octax:space_flight4
• octax:space_flight5
• octax:space_flight6
• octax:space_flight7
• octax:space_flight8
• octax:space_flight9
• octax:space_flight10
• octax:spacejam
• octax:squash
• octax:submarine
• octax:tank
• octax:target_shooter1
• octax:target_shooter2
• octax:target_shooter3
• octax:tetris
• octax:ufo
• octax:vertical_brix
• octax:wipe_off
• octax:worm
• octax:cavern1
• octax:space_flight1
• octax:target_shooter1
• craftax:Craftax-Classic-Symbolic-v1
• craftax:Craftax-Classic-Pixels-v1
• craftax:Craftax-Symbolic-v1
• craftax:Craftax-Pixels-v1
• factoriax:MinerBootstrap-v1
• factoriax:EasyRocket-v1
• factoriax:Rocket-v1
• popjym:NoisyStatelessMetaCartPole
• popjym:AutoencodeEasy
• popjym:AutoencodeMedium
• popjym:AutoencodeHard
• popjym:BattleshipEasy
• popjym:BattleshipMedium
• popjym:BattleshipHard
• popjym:ConcentrationEasy
• popjym:ConcentrationMedium
• popjym:ConcentrationHard
• popjym:CountRecallEasy
• popjym:CountRecallMedium
• popjym:CountRecallHard
• popjym:HigherLowerEasy
• popjym:HigherLowerMedium
• popjym:HigherLowerHard
• popjym:MinesweeperEasy
• popjym:MinesweeperMedium
• popjym:MinesweeperHard
• popjym:MultiArmedBanditEasy
• popjym:MultiArmedBanditMedium
• popjym:MultiArmedBanditHard
• popjym:StatelessCartPoleEasy
• popjym:StatelessCartPoleMedium
• popjym:StatelessCartPoleHard
• popjym:NoisyStatelessCartPoleEasy
• popjym:NoisyStatelessCartPoleMedium
• popjym:NoisyStatelessCartPoleHard
• popjym:StatelessPendulumEasy
• popjym:StatelessPendulumMedium
• popjym:StatelessPendulumHard
• popjym:NoisyStatelessPendulumEasy
• popjym:NoisyStatelessPendulumMedium
• popjym:NoisyStatelessPendulumHard
• popjym:RepeatFirstEasy
• popjym:RepeatFirstMedium
• popjym:RepeatFirstHard
• popjym:RepeatPreviousEasy
• popjym:RepeatPreviousMedium
• popjym:RepeatPreviousHard
• playground:AcrobotSwingup
• playground:AcrobotSwingupSparse
• playground:BallInCup
• playground:CartpoleBalance
• playground:CartpoleBalanceSparse
• playground:CartpoleSwingup
• playground:CartpoleSwingupSparse
• playground:CheetahRun
• playground:FingerSpin
• playground:FingerTurnEasy
• playground:FingerTurnHard
• playground:FishSwim
• playground:HopperHop
• playground:HopperStand
• playground:HumanoidStand
• playground:HumanoidWalk
• playground:HumanoidRun
• playground:PendulumSwingup
• playground:PointMass
• playground:SwimmerSwimmer6
• playground:WalkerRun
• playground:WalkerStand
• playground:WalkerWalk
• playground:ApolloJoystickFlatTerrain
• playground:BarkourJoystick
• playground:BerkeleyHumanoidJoystickFlatTerrain
• playground:BerkeleyHumanoidJoystickRoughTerrain
• playground:G1JoystickFlatTerrain
• playground:G1JoystickRoughTerrain
• playground:Go1JoystickFlatTerrain
• playground:Go1JoystickRoughTerrain
• playground:Go1Getup
• playground:Go1Handstand
• playground:Go1Footstand
• playground:H1InplaceGaitTracking
• playground:H1JoystickGaitTracking
• playground:Op3Joystick
• playground:SpotFlatTerrainJoystick
• playground:SpotGetup
• playground:SpotJoystickGaitTracking
• playground:T1JoystickFlatTerrain
• playground:T1JoystickRoughTerrain
• playground:AlohaHandOver
• playground:AlohaSinglePegInsertion
• playground:PandaPickCube
• playground:PandaPickCubeOrientation
• playground:PandaOpenCabinet
• playground:LeapCubeReorient
• playground:LeapCubeRotateZAxis
• playground:AeroCubeRotateZAxis
• kinetix:s
• kinetix:random-s
• kinetix:m
• kinetix:random-m
• kinetix:l
• kinetix:random-l