Metadata-Version: 2.4
Name: alenax
Version: 0.1.3
Summary: Add your description here
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: gymnax>=0.0.9
Requires-Dist: ale-py
Provides-Extra: dev
Requires-Dist: mypy>=1.19.1; extra == "dev"
Requires-Dist: pytest>=8; extra == "dev"
Requires-Dist: black; extra == "dev"
Requires-Dist: isort; extra == "dev"

# alenax

A JAX wrapper for the [Arcade Learning Environment](https://github.com/Farama-Foundation/Arcade-Learning-Environment), compatible with [Gymnax](https://github.com/RobertTLange/gymnax).

## Installation

```bash
pip install alenax
```

Requires Python >= 3.10.

## Usage

```python
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

```python
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
```
