Metadata-Version: 2.4
Name: mujax
Version: 0.1.0
Summary: MuZero for Gymnasium environments, built on JAX and mctx.
License-Expression: Apache-2.0
Project-URL: Homepage, https://github.com/okayapinar/mujax
Project-URL: Repository, https://github.com/okayapinar/mujax
Project-URL: Issues, https://github.com/okayapinar/mujax/issues
Keywords: muzero,reinforcement-learning,jax,mcts,gymnasium
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: jax>=0.4.30
Requires-Dist: flax>=0.10.0
Requires-Dist: optax>=0.2.0
Requires-Dist: orbax-checkpoint>=0.11.0
Requires-Dist: mctx>=0.0.6
Requires-Dist: flashbax>=0.1.3
Requires-Dist: gymnasium>=1.1.0
Requires-Dist: numpy>=2.0.0
Requires-Dist: tqdm
Requires-Dist: rlax>=0.1.9
Requires-Dist: distrax>=0.1.9
Requires-Dist: tyro>=1.0.16
Provides-Extra: cuda12
Requires-Dist: jax[cuda12]>=0.4.30; extra == "cuda12"
Provides-Extra: wandb
Requires-Dist: wandb; extra == "wandb"
Dynamic: license-file

# mujax

Gumbel MuZero and Stochastic MuZero for [Gymnasium](https://gymnasium.farama.org/) environments, built on
JAX, [mctx](https://github.com/google-deepmind/mctx) and [flashbax](https://github.com/instadeepai/flashbax).

- **Two algorithms, one pipeline**: `GMZ` (Gumbel MuZero) and `SMZ` (Stochastic MuZero with afterstates and a
  VQ chance codebook) share replay, reanalyze, actor, learner and checkpointing.
- **Single process**: actors and evaluator run in threads on CPU, the learner runs on GPU when available. The CPU
  search releases the GIL, so `num_actors` actor threads (`--num-actors`) scale env throughput with the core count.
- **Reanalyze**: a fraction of every batch is re-searched with the latest parameters.
- **Categorical value/reward** with HL-Gauss targets in symlog space.
- **Symlog observations**: raw observations enter the networks through `signed_logp1` (DreamerV3), so no per-dimension normalization is needed.
- **Checkpoints**: best parameters by evaluation score plus full training state for resuming (Orbax), optional W&B upload.

The design is inspired by DeepMind's [Acme](https://github.com/google-deepmind/acme): the split into actor, learner,
replay and environment loop, the observers and the counter follow Acme's structure, reduced to a small single-process
library.

Requirements: 1-D `Box` observations and `Discrete` actions (use `gym.wrappers.FlattenObservation` otherwise).
Illegal actions can be masked through `info["invalid_actions"]`.

## Installation

```bash
pip install mujax                 # CPU
pip install "mujax[cuda12]"       # NVIDIA GPU
pip install "mujax[wandb]"        # Weights & Biases logging
```

## Command line

```bash
mujax --env CartPole-v1 --num-steps 20000 gmz --num-simulations  16
mujax --env Acrobot-v1 smz --batch-size 1024
mujax --help          # all options; `mujax gmz --help`, `mujax smz --help` for algorithm options
```

Checkpoints are written to `checkpoints/<algo>_<timestamp>/`.

## Python API

```python
import gymnasium as gym
from mujax import CheckpointingConfig, ExperimentConfig, GMZConfig, load_actor, run_experiment


def make_env(num_envs):
    return gym.make_vec(
        "CartPole-v1", num_envs=num_envs, vectorization_mode="sync", vector_kwargs={"autoreset_mode": gym.vector.AutoresetMode.SAME_STEP}
    )


run_experiment(
    ExperimentConfig(
        config=GMZConfig(batch_size=256).with_num_steps(20_000),
        environment_factory=lambda seed: make_env(16),
        eval_environment_factory=lambda seed: make_env(1),
        max_num_learner_steps=20_000,
        checkpointing=CheckpointingConfig(directory="checkpoints/cartpole"),
    )
)

actor, metadata = load_actor("checkpoints/cartpole")
```

Environments must use `AutoresetMode.SAME_STEP`. See [`examples/`](https://github.com/okayapinar/mujax/tree/master/examples) for complete training (`cartpole.py`) and evaluation (`evaluate.py`) scripts.

### Logging

`create_writer` sends metrics to the terminal and, optionally, Weights & Biases; learner, actor and evaluator metrics are written under `learner/`, `actor/` and `evaluator/` with `learner_steps` as the step:

```python
writer = create_writer(console=True, wandb_project="mujax")
ExperimentConfig(..., writer=writer)
```

Any object with `write(step, scalars)`, `write_config(config)` and `close()` works as a writer (`mujax.Writer`). Without a writer only the progress bar is shown.

On the command line metrics are printed to the terminal (`--no-console` turns it off); add `--wandb` (`--wandb-project`) for W&B.

## Layout

| Module | Contents |
| --- | --- |
| `config.py` | `MuZeroConfig`, shared hyperparameters |
| `algorithm.py` | `Algorithm`, `SearchPolicy`: the contract between the infrastructure and an algorithm |
| `algorithms/__init__.py` | `ALGORITHMS` registry, `algorithm_for` |
| `algorithms/muzero.py` | shared math of the MuZero family: `Support`, network blocks, loss skeleton |
| `algorithms/smz.py`, `algorithms/gmz.py` | algorithms: config, networks, search, loss |
| `replay.py`, `reanalyze.py` | flashbax replay and reanalyze iterator |
| `actor.py`, `learner.py`, `loop.py` | environment interaction, gradient step, environment loop |
| `observers.py`, `loggers.py`, `checkpoint.py` | metrics, terminal and W&B writers, Orbax checkpoints |
| `experiment.py` | `ExperimentConfig`, `run_experiment`, `load_actor` |

## Development

```bash
uv sync
uv run pytest
```

## References

- Schrittwieser et al., *Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model* (2020)
- Antonoglou et al., *Planning in Stochastic Environments with a Learned Model* (2022)
- Danihelka et al., *Policy Improvement by Planning with Gumbel* (2022)
- Hoffman et al., *Acme: A Research Framework for Distributed Reinforcement Learning* (2020)

## License

Apache-2.0
