README.md

metadata

license: mit
pipeline_tag: reinforcement-learning
tags:
  - reinforcement-learning
  - planning
  - discrete-diffusion
  - remdm
  - craftax
  - jax
  - flax
  - orbax

ReMDM Planner — Craftax checkpoints

Trained weights accompanying The Double Intractability of Reinforcement Learning for Discrete Diffusion Planners: a remasking discrete diffusion model (ReMDM) used as an action-sequence planner in Craftax, together with the PPO-RNN experts that supervise it.

Code, configs and evaluation harness: https://github.com/mathisweil/craftax-ReMDM-planner

Contents

Path	Role	Environment	Architecture	Latest step	Training	Size
checkpoints/offline/Craftax-Classic-Symbolic-v1-OfflineDiffusion-BC-100M	Diffusion planner (offline BC)	Craftax-Classic-Symbolic-v1	6L, d_model 384, 8 heads, horizon 32	100000000	97,600 grad steps	97 MB
checkpoints/offline/Craftax-Symbolic-v1-OfflineDiffusion-BC-100M	Diffusion planner (offline BC)	Craftax-Symbolic-v1	6L, d_model 384, 8 heads, horizon 32	100000000	111,552 grad steps	128 MB
checkpoints/online/Craftax-Classic-Symbolic-v1-OnlineDiffusion-DAgger-100M	Diffusion planner (online DAgger)	Craftax-Classic-Symbolic-v1	6L, d_model 384, 8 heads, horizon 32	100000000	97,600 grad steps	33 MB
checkpoints/online/Craftax-Symbolic-v1-OnlineDiffusion-DAgger-100M	Diffusion planner (online DAgger)	Craftax-Symbolic-v1	6L, d_model 384, 8 heads, horizon 32	100000000	111,552 grad steps	128 MB
checkpoints/ppo_agents/Craftax-Classic-Symbolic-v1-PPO_RNN-1000M	PPO-RNN expert	Craftax-Classic-Symbolic-v1	RNN, layer size 512	1000000000	1e+09 frames	35 MB
checkpoints/ppo_agents/Craftax-Symbolic-v1-PPO_RNN-1000M	PPO-RNN expert	Craftax-Symbolic-v1	RNN, layer size 512	1000000000	1e+09 frames	50 MB

Each diffusion checkpoint ships a resume_metadata.json holding the full config snapshot it was trained under; each PPO expert ships config.yaml and wandb-summary.json (final training metrics).

Weights are Orbax checkpoint directories (OCDBT format), not safetensors — the models are Flax modules restored via orbax.checkpoint, and the paths above mirror the source repository so a snapshot can be dropped straight into a working copy.

Download

from huggingface_hub import snapshot_download

# everything (~471 MB)
snapshot_download(repo_id="MathisW78/remdm-craftax-checkpoint", local_dir=".")

# a single model
snapshot_download(
    repo_id="MathisW78/remdm-craftax-checkpoint",
    local_dir=".",
    allow_patterns="checkpoints/online/Craftax-Classic-*/**",
)

Use

From a clone of the code repository, after downloading into it:

uv run python main.py --mode inference \
    --checkpoint_path checkpoints/online/Craftax-Classic-Symbolic-v1-OnlineDiffusion-DAgger-100M

Programmatic loading uses src.planners.model.load_checkpoint for the diffusion planners and src.planners.ppo.load_ppo_agent for the experts; both take the checkpoint directory path and restore the latest step. Architecture arguments should be read from the checkpoint's own resume_metadata.json rather than hardcoded.

Training

The diffusion planners are bidirectional transformers that denoise a masked action plan conditioned on the symbolic observation, trained either by offline behaviour cloning on PPO rollouts or by online DAgger against the PPO expert. Model size and horizon differ per run (see the table); the PPO-RNN experts are the Craftax baselines. Exact hyperparameters for every run, including the remasking strategy, schedule and sampling settings, are in the per-checkpoint metadata files listed above, which are the authoritative record.

Directory names encode the environment and the total environment timesteps the run was trained for. Latest step is whatever each run used as its Orbax step counter: environment frames for most runs, and update steps for the full Craftax DAgger run, whose 1,743 updates cover the same ~100M timesteps.

Limitations

These are research artefacts tied to specific Craftax versions and symbolic observation encodings; they are not general-purpose agents and will not transfer to other environments or to pixel observations. Evaluation results and their variance are reported in the paper.

Citation

@inproceedings{remdm-craftax-planner,
  title  = {The Double Intractability of Reinforcement Learning for Discrete Diffusion Planners},
  author = {Weil, Mathis},
  year   = {2026},
  note   = {NeurIPS 2026 Workshop: Beyond Next-Token Prediction}
}

License

MIT, see LICENSE.