mathisweil/remdm-minihack-checkpoints

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

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:

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:

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

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