mujax

Gumbel MuZero and Stochastic MuZero for Gymnasium environments, built on JAX, mctx and 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: 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

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

Command line

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

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

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

Development

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