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):

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