remdm-minihack-checkpoints

ReMDM Planner: MiniHack checkpoints

Trained weights accompanying Return-Weighted ELBO Fine-Tuning Degrades Masked 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

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.

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

Manuscript figures as vector PDF, under results/paper_figures/. Each puts Craftax Classic and MiniHack side by side, so they are built from both repositories' ablation results.json by the sibling repo's scripts/paper_figures.py and published identically in both Hub repos.

Download

This repo mirrors the code repository's layout, so a snapshot drops straight into a working copy -- but it also carries its own README.md (this card), LICENSE and .gitattributes, and local_dir="." would overwrite the code repository's copies of all three. Exclude them, or download into a directory of its own.

from huggingface_hub import snapshot_download

# everything (~229 MB), into a clone of the code repository
snapshot_download(
    repo_id="mathisweil/remdm-minihack-checkpoints",
    local_dir=".",
    ignore_patterns=["README.md", "LICENSE", ".gitattributes"],
)

# or somewhere of its own, leaving any working copy untouched
snapshot_download(repo_id="mathisweil/remdm-minihack-checkpoints", local_dir="remdm-minihack")

# 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 (selected on mean ID+OOD win rate). 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  = {Return-Weighted ELBO Fine-Tuning Degrades Masked Diffusion Planners},
  author = {Weil, Mathis},
  year   = {2026},
  note   = {NeurIPS 2026 Workshop: Beyond Next-Token Prediction}
}

License

MIT, see LICENSE.