README.md

---
license: mit
library_name: pytorch
pipeline_tag: reinforcement-learning
tags:
- reinforcement-learning
- planning
- discrete-diffusion
- remdm
- minihack
- nethack
- pytorch
---

# ReMDM Planner: MiniHack 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
[MiniHack](https://github.com/facebookresearch/minihack), together with the BFS
oracle rollouts that supervise it, and the results reported in the paper.

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

## Contents

| Path | Role | Environment | Architecture | Selected at | Training | Size |
|---|---|---|---|---|---|---|
| `checkpoints/offline/Minihack-Offline-Diffusion-BC-100M` | Diffusion planner (offline BC) | `MiniHack` | 4L, d_model 256, 4 heads, horizon 64, 5M params | gradient step 50,000 | 102,400,000 sample-equivalents | 100 MB |
| `checkpoints/online/Minihack-Online-Diffusion-DAgger-100M` | Diffusion planner (online DAgger) | `MiniHack` | 4L, d_model 256, 4 heads, horizon 64, 5M params | iteration 563 | 5,654,965 env steps | 100 MB |

Each checkpoint ships the `.pth` training state it was published from (weights,
EMA shadow, optimiser, scheduler, and for the DAgger run the curriculum and RNG
state, so training resumes exactly), a `model.safetensors` export of the EMA
weights for inference, the YAML config snapshot it was trained under, and a
`selection.json` recording how it was chosen.

Weights are PyTorch training states with a `safetensors` export of the EMA
weights alongside, 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/minihack_ablations` | `results.json`, `diagnosis.md`, 17 tables, 113 figures | 27 MB |

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

| File | Environment | Evaluation | Headline metric | Size |
|---|---|---|---|---|
| `eval_offline_s0.json` | `7 envs` | 100 episodes per env | mean win rate 0.46 | 1 KB |
| `eval_online_s0.json` | `7 envs` | 100 episodes per env | mean win rate 0.41 | 1 KB |

## Download

```python
from huggingface_hub import snapshot_download

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

# a single model
snapshot_download(
    repo_id="mathisweil/remdm-minihack-checkpoints",
    local_dir=".",
    allow_patterns="checkpoints/offline/Minihack-Offline-Diffusion-BC-100M/**",
)
```

## Use

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

```bash
DIR=checkpoints/offline/Minihack-Offline-Diffusion-BC-100M
uv run python main.py --mode inference \
    --config $DIR/config.yaml --checkpoint $DIR/offline_step50000.pth \
    --output results/inference/eval.json
```

Programmatic loading uses `src.models.denoiser.make_model` with the checkpoint's
own config, then the safetensors export:

```python
from safetensors.torch import load_file
from src.config import load_config
from src.models.denoiser import make_model

cfg = load_config("checkpoints/offline/Minihack-Offline-Diffusion-BC-100M/config.yaml")
model = make_model(cfg)
model.load_state_dict(load_file("checkpoints/offline/Minihack-Offline-Diffusion-BC-100M/model.safetensors"))
model.eval()
```

Architecture arguments should be read from the checkpoint's own config snapshot
rather than from `configs/defaults.yaml`, which tracks the current code.

## Training

The planners are bidirectional transformers that denoise a masked action plan
conditioned on a cropped MiniHack glyph observation, trained either by offline
behaviour cloning on oracle rollouts or by online DAgger against the BFS oracle
under a dynamic environment curriculum. Model size and horizon differ per run
(see the table). Exact hyperparameters for every run, including the
in-distribution and out-of-distribution environment sets, the remasking
strategy, schedule and sampling settings, are in the per-checkpoint config
snapshots listed above, which are the authoritative record.

Both models are best-checkpoint selections rather than final-step dumps: each
trainer evaluates every periodic checkpoint on its configured number of
episodes per environment using EMA weights, and the highest-scoring one is
published (the metric behind that selection is not recorded in this release). Directory names encode the sample-equivalents the
published model consumed (gradient steps x batch size, rounded); file names
carry each trainer's own counter, DAgger iterations online and gradient steps
offline. Each checkpoint's `selection.json` records the configured budget it
was drawn from and the step it was selected at.

## Limitations

These are research artefacts tied to specific MiniHack environment versions and
to the cropped-glyph observation encoding; 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-minihack-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`.