alenax

A JAX wrapper for the Arcade Learning Environment, compatible with Gymnax.

Installation

pip install alenax

Requires Python >= 3.10.

Usage

import jax
from alenax import AtariEnv

env = AtariEnv("pong")
key = jax.random.PRNGKey(0)

obs, state = env.reset(key)
obs, state, reward, done, info = env.step(key, state, action)

All environment methods are JIT-compiled and support jax.vmap for batched execution.

Wrappers

from alenax import ClipReward, EpisodicLife, RecordEpisodeStatistics

env = AtariEnv("pong")
env = ClipReward(env)              # Clip rewards to [-1, 1]
env = EpisodicLife(env)            # Treat life loss as episode end
env = RecordEpisodeStatistics(env) # Track episode returns and lengths