mathisweil/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

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,657,661 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, 20 tables, 113 figures, 4 gdelta	29 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	50 episodes per env	mean win rate 0.44	1 KB
eval_online_s0.json	7 envs	50 episodes per env	mean win rate 0.30	1 KB

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.

Figure	Size
fig10_timestep_conditioning.pdf	31 KB
fig11_train_vs_eval.pdf	34 KB
fig1_finetuning_trajectories.pdf	21 KB
fig2_repr_drift.pdf	24 KB
fig3_cka.pdf	19 KB
fig4_grad_alignment.pdf	23 KB
fig5_minihack_per_env.pdf	18 KB
fig6_achievements.pdf	24 KB
fig7_tbin_gradients.pdf	20 KB
fig8_score_vs_kl.pdf	23 KB
fig9_weight_dispersion.pdf	45 KB

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.