---
license: mit
library_name: jax
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](https://github.com/MichaelTMatthews/Craftax), together with the
PPO-RNN experts that supervise it, and the results reported in the paper.

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

## Contents

| Path | Role | Environment | Architecture | Selected at | Training | Size |
|---|---|---|---|---|---|---|
| `checkpoints/offline/Craftax-Classic-Symbolic-v1-Offline-Diffusion-BC-100M` | Diffusion planner (offline BC) | `Craftax-Classic-Symbolic-v1` | 6L, d_model 384, 8 heads, horizon 32 | 1,000,000,000 | 97,600 grad steps | 97 MB |
| `checkpoints/online/Craftax-Classic-Symbolic-v1-Online-Diffusion-DAgger-100M` | Diffusion planner (online DAgger) | `Craftax-Classic-Symbolic-v1` | 6L, d_model 384, 8 heads, horizon 32 | 100,000,000 | 97,600 grad steps | 97 MB |
| `checkpoints/ppo_agents/Craftax-Classic-Symbolic-v1-PPO_RNN-1000M` | PPO-RNN expert | `Craftax-Classic-Symbolic-v1` | RNN, layer size 512 | 1,000,000,000 | 1e+09 frames | 35 MB |
| `checkpoints/ppo_agents/Craftax-Symbolic-v1-PPO_RNN-1000M` | PPO-RNN expert | `Craftax-Symbolic-v1` | RNN, layer size 512 | 1,000,000,000 | 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](https://orbax.readthedocs.io) 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.

## Results

RL fine-tuning ablation runs, as produced by `experiments/rl_finetuning/run_ablations.py`. Each run ships its `results.json` summary, the `diagnosis.md` write-up, and the tables (`.csv` and `.tex`) and figures generated from it.

| Run | Contents | Size |
|---|---|---|
| `experiments/rl_finetuning/outputs/craftax_classic_ablations` | `results.json`, `diagnosis.md`, 19 tables, 114 figures | 41 MB |

Evaluation results produced by `main.py --mode inference` on the checkpoints above, under `results/inference/`.

| File | Environment | Evaluation | Headline metric | Size |
|---|---|---|---|---|
| `eval_classic_bc_s0.json` | `Craftax-Classic-Symbolic-v1` | 256 envs x 10000 steps | mean score 4.15 | 1 KB |
| `eval_classic_dagger_s0.json` | `Craftax-Classic-Symbolic-v1` | 256 envs x 10000 steps | mean score 2.72 | 1 KB |

## Download

```python
from huggingface_hub import snapshot_download

# everything (~320 MB)
snapshot_download(repo_id="mathisweil/remdm-craftax-checkpoints", local_dir=".")

# a single model
snapshot_download(
    repo_id="mathisweil/remdm-craftax-checkpoints",
    local_dir=".",
    allow_patterns="checkpoints/offline/Craftax-Classic-Symbolic-v1-Offline-Diffusion-BC-100M/**",
)
```

## Use

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

```bash
uv run python main.py --mode inference \
    --checkpoint checkpoints/offline/Craftax-Classic-Symbolic-v1-Offline-Diffusion-BC-100M \
    --output results/inference/eval.json
```

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. `Selected at` is whatever each run used as its Orbax step
counter, which is environment frames for the runs published here.

## 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

```bibtex
@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`.