POPGym Arcade - GPU-Accelerated POMDPs

POPGym Arcade contains 10 pixel-based tasks in the style of the Arcade Learning Environment. Each environment provides: - 3 Difficulty settings - Common observation and action space shared across all envs - Fully observable and partially observable configurations - Fast and easy GPU vectorization using jax.vmap and jax.jit

Gradient Visualization

We also provide tools to visualize how policies use memory.

See below for further instructions.

Throughput

You can expect millions of frames per second on a consumer-grade GPU. With obs_size=128, most policies converge within 30-60 minutes of training.

Getting Started

Installation

To install the environments, run

bash pip install popgym-arcade If you plan to use our training scripts, install the baselines as well

bash pip install 'popgym-arcade[baselines]'

Human Play

To best understand the environments, you should try and play them yourself. The play script lets you play the games yourself using the arrow keys and spacebar.

bash popgym-arcade-play NoisyCartPoleEasy # play MDP 256 pixel version popgym-arcade-play BattleShipEasy -p -o 128 # play POMDP 128 pixel version

Creating and Stepping Environments

Our envs are gymnax envs, so you can use your wrappers and code designed to work with gymnax. The following example demonstrates how to integrate POPGym Arcade into your code.

```python import popgym_arcade import jax

Create both POMDP and MDP env variants

pomdp, pomdpparams = popgymarcade.make("BattleShipEasy", partialobs=True) mdp, mdpparams = popgymarcade.make("BattleShipEasy", partialobs=False)

Let's vectorize and compile the envs

Note when you are training a policy, it is better to compile your policyupdate rather than the envstep

pomdpreset = jax.jit(jax.vmap(pomdp.reset, inaxes=(0, None))) pomdpstep = jax.jit(jax.vmap(pomdp.step, inaxes=(0, 0, 0, None))) mdpreset = jax.jit(jax.vmap(mdp.reset, inaxes=(0, None))) mdpstep = jax.jit(jax.vmap(mdp.step, inaxes=(0, 0, 0, None)))

Initialize four vectorized environments

n_envs = 4

Initialize PRNG keys

key = jax.random.key(0) resetkeys = jax.random.split(key, nenvs)

Reset environments

observation, envstate = pomdpreset(resetkeys, pomdpparams)

Step the POMDPs

for t in range(10): # Propagate some randomness actionkey, stepkey = jax.random.split(jax.random.key(t)) actionkeys = jax.random.split(actionkey, nenvs) stepkeys = jax.random.split(stepkey, nenvs) # Pick actions at random actions = jax.vmap(pomdp.actionspace(pomdpparams).sample)(actionkeys) # Step the env to the next state # No need to reset, gymnax automatically resets when done observation, envstate, reward, done, info = pomdpstep(stepkeys, envstate, actions, pomdpparams)

POMDP and MDP variants share states

We can plug the POMDP states into the MDP and continue playing

actionkeys = jax.random.split(jax.random.key(t + 1), nenvs) stepkeys = jax.random.split(jax.random.key(t + 2), nenvs) markovstate, envstate, reward, done, info = mdpstep(stepkeys, envstate, actions, mdpparams) ```

Memory Introspection Tools

We implement visualization tools to probe which pixels persist in agent memory, and their impact on Q value predictions. Try code below or vis example to visualize the memory your agent uses

```python from popgymarcade.baselines.model.builder import QNetworkRNN from popgymarcade.baselines.utils import getsaliencymaps, vis_fn import equinox as eqx import jax

config = { # Env string "ENVNAME": "NavigatorEasy", # Whether to use full or partial observability "PARTIAL": True, # Memory model type (see models directory) "MEMORYTYPE": "lru", # Evaluation episode seed "SEED": 0, # Observation size in pixels (128 or 256) "OBS_SIZE": 128, }

Initialize the random key

rng = jax.random.PRNGKey(config["SEED"])

Initialize the model

network = QNetworkRNN(rng, rnntype=config["MEMORYTYPE"], obssize=config["OBSSIZE"])

Load the model

model = eqx.treedeserialiseleaves("PATHTOYOURMODELWEIGHTS.pkl", network)

Compute the saliency maps

grads, obsseq, gradaccumulator = getsaliencymaps(rng, model, config)

Visualize the saliency maps

If you have latex installed, set use_latex=True

visfn(grads, obsseq, config, use_latex=False) ```

Other Useful Libraries

• gymnax - The (deprecated) jax-capable gymnasium API
• stable-gymnax - A maintained and patched version of gymnax
• popgym - The original collection of POMDPs, implemented in numpy
• popjaxrl - A jax version of popgym
• popjym - A more readable version of popjaxrl environments that served as a basis for our work

Citation

@article{wang2025popgym, title={POPGym Arcade: Parallel Pixelated POMDPs}, author={Wang, Zekang and He, Zhe and Zhang, Borong and Toledo, Edan and Morad, Steven}, journal={arXiv preprint arXiv:2503.01450}, year={2025} }