pobax 0.0.1

Reinforcement learning benchmark that tests all forms of partial observability in JAX.

POBAX: Partially Observable Benchmarks in JAX

---

POBAX is a reinforcement learning benchmark that tests all forms of partial observability.

POBAX has been accepted to RLC 2025. Check out our paper!

The benchmark is entirely written in JAX, allowing for fast, GPU-scalable experimentation.

Environments

---

POBAX includes environments (as well as recommended hyperparameter settings) across diverse forms of partial observability. We list our environments from smallest to largest (in terms of neural network size requirements for PPO RNN):

Environment	Category	IDs	Description
Simple Chain		simple_chain	Diagnostic POMDP for testing algorithms.
T-Maze		tmaze_10	Bakker's classic memory testing environment (hallway disambiguation).
RockSample		rocksample_11_11 and rocksample_15_15	The classic rock collecting POMDP, where an agent needs to uncover and collect rocks.
Battleship		battleship_10	Single-player battleship (10x10).
Masked Mujoco		{env_name}-{F/P/V}-v0	Mujoco with state features masked out. env_name can be Walker, Ant, Hopper, or HalfCheetah. F/P/V stands for fully observable, position only, or velocity only versions of environments, respectively.
DMLab Minigrid		Navix-DMLab-Maze-{01/02/03}-v0	MiniGrid versions of the DeepMind Lab mazes. 01/02/03 refer to the DeepMind Lab Minigrid mazes in ascending difficulty.
Visual Continuous Control		{env_name}-pixels	Pixel-based versions of Mujoco control. Requires the Madrona_MJX package. env_name can be ant, halfcheetah, hopper, or walker2d.
No-Inventory Crafter		craftax-pixels	Crafter without the inventory. Requires the Craftax package.

Basic Usage

import jax
from pobax.envs import get_env

rand_key = jax.random.PRNGKey(2025)
env_key, rand_key = jax.random.split(rand_key)

# Creates a vectorized environment
env, env_params = get_env("rocksample_11_11", env_key)

# Reset 10 environments
reset_key, rand_key = jax.random.split(rand_key)
reset_keys = jax.random.split(rand_key, 10)

obs, env_state = env.reset(reset_keys, env_params)

# Take steps in all environments
step_key, action_key, rand_key = jax.random.split(rand_key, 3)
step_keys = jax.random.split(step_key, 10)
action_keys = jax.random.split(action_key, 10)

actions = jax.vmap(env.action_space(env_params).sample)(action_keys)

obs, env_state, reward, done, info = env.step(step_keys, env_state, actions, env_params)

Installation

Agents

POBAX includes algorithms loosely based on the PureJAXRL framework, with algorithms based on proximal policy optimization (PPO). These include:

• Recurrent PPO,
• λ-discrepancy,
• GTrXL.

Memoryless versions of the recurrent PPO algorithm is also included with the --memoryless flag.

Citation