license: mit
library_name: pytorch
pipeline_tag: reinforcement-learning
tags:
- reinforcement-learning
- planning
- discrete-diffusion
- remdm
- minihack
- nethack
- pytorch
ReMDM Planner β MiniHack artefacts
Anonymous release accompanying Return-Weighted ELBO Fine-Tuning Degrades Masked Diffusion Planners (under double-blind review). It contains the code, the trained checkpoints and the pre-computed results behind every MiniHack number in the paper.
Anonymity notice. This repository is an anonymised mirror prepared for
double-blind review. Author names, institutional affiliations, experiment-tracking
identifiers and absolute filesystem paths have been removed or replaced with
placeholders. Cluster hardware is referred to only as GPU-24GB and GPU-H200.
What is here
src/ configs/ environments/ experiments/ scripts/ tests/
main.py pyproject.toml uv.lock LICENSE demo_minihack.ipynb
checkpoint_inference.pth Stripped EMA weights only (~20 MB), what the notebook loads
checkpoints/
online/Minihack-Online-Diffusion-DAgger-100M/ DAgger planner (the checkpoint every
ablation fine-tunes)
offline/Minihack-Offline-Diffusion-BC-100M/ Offline BC planner
ablation_assets/ Flat copy of the assets the notebook reads:
9 PNGs + 4 CSVs + results.json
experiments/rl_finetuning/outputs/minihack_ablations/ The published 25-condition suite
figures/ 113 pre-computed PNGs
tables/ main_results.csv, group_summary.csv, hypothesis_verdict.csv,
gradient_analysis.csv, significance_test.txt, gdelta.csv, results.tex, ...
gdelta/ gdelta_{seed0,seed1,seed2,aggregate}.json <- the return-term measurement
results.json, diagnosis.md
results/*.csv Per-analysis summaries (main results, per-env win rates,
forgetting, gradient diagnostics, representation drift, verdicts)
results/inference/ Single-seed evaluations of both checkpoints
results/paper_figures/ The 11 manuscript figures as vector PDF
Headline numbers (MiniHack, in-distribution win rate)
| Win rate | |
|---|---|
| Pretrained (DAgger) checkpoint | 47.5% |
| Baseline return-weighted ELBO fine-tuning, 500 iters | 43.8% Β± 6.1 |
| Worst of 25 conditions (normalised advantages) | 12.1% |
Three conditions finish nominally above the checkpoint, but the best-versus-baseline difference has a bootstrap interval of [β2.1, +13.7] points at p = 0.40, so the paper does not read them as improvements.
MiniHack is not a second confirmation, and the paper says so. The effect here is small relative to a seed standard deviation of 6.1 points. It is reported because it is the same suite run identically and its ordering matches Craftax Classic, where the effect is large and unambiguous. The central claim rests on Craftax Classic.
The return term
experiments/rl_finetuning/outputs/minihack_ablations/gdelta/ holds the direct
measurement of g_Ξ΄ at the pretrained checkpoint, over 3 rollout seeds Γ 8 noise draws
on 4,608-window batches. Reproduce it on a CPU with run_ablations.py --measure-gdelta.
MiniHack does not look like Craftax Classic here: βg_Ξ΄β/ββL_BCβ = 1.956 Β± 0.275 at
cosine 0.803 Β± 0.039 (imitation against itself gives 0.986 Β± 0.006 on this
benchmark), against 0.485 Β± 0.007 at cosine 0.02 Β± 0.05 on Craftax Classic. A
shuffled-Ξ΄ null costs about five sixths of the ratio here and about a quarter there,
so the ratio is largely weight dispersion on Craftax Classic and largely not here.
β One caveat the paper records. The Eq.-4 correctness residual is 4.8 Γ 10β»β΅ here against 4.8 Γ 10β»β· on Craftax Classic β two orders looser, and it did not tighten when the measurement went from 3 draws to 8.
The review arms are not in this release
Three questions outside the 25-condition suite were run on both benchmarks β the
unweighted control (bc_all), a learning-rate sweep, and advantage_clip at a matched
effective step. The Craftax Classic run directories are published, in
AnonMLuser/remdm-planner-craftax
under experiments/rl_finetuning/outputs/review_*/.
The MiniHack run directories are not part of this release. Their numbers are
reported in the paper and quoted in demo_minihack.ipynb, but no artefact in this
bundle backs them:
| Arm | lr | Win rate | vs published baseline 0.4375 |
|---|---|---|---|
bc_all β uniform weights, all windows |
3e-4 | 0.3375 Β± 0.0270 | β0.1000 |
advantage_clip β matched effective step |
1.636e-4 | 0.3708 Β± 0.0425 | β0.0667 |
baseline_rl |
1e-4 | 0.4417 Β± 0.0156 | +0.0042 |
baseline_rl |
1e-5 | 0.4667 Β± 0.0059 | +0.0292 |
Unlike Craftax Classic, no same-stack anchor was run here, so these are quoted against the published suite rather than a contemporaneous rerun.
Quick start
from huggingface_hub import snapshot_download
path = snapshot_download(repo_id="AnonMLuser/remdm-planner-minihack", local_dir="remdm-planner-minihack")
Then open demo_minihack.ipynb, or see the project documentation below.
Citation
@inproceedings{anonymous2026returnweighted,
title = {Return-Weighted ELBO Fine-Tuning Degrades Masked Diffusion Planners},
author = {Anonymous},
booktitle = {Submitted to NeurIPS 2026},
year = {2026},
note = {Under double-blind review}
}
ReMDM Planner for MiniHack
PyTorch implementation of ReMDM (Remasking Discrete Diffusion Model) for action-sequence planning in MiniHack navigation environments. A dual-stream transformer generates 64-step action plans by iteratively denoising masked token sequences, conditioned on a 9x9 local crop and the full 21x79 dungeon map. Trained with DAgger under BFS oracle supervision, from scratch; generalises zero-shot from 4 in-distribution to 3 out-of-distribution environments.
The sibling repository remdm-planner-craftax implements the same method in JAX on Craftax. Both repos share the same CLI, config layout and README structure; commands transfer between them by swapping the repo name and benchmark-specific values.
Method
The planner starts from a fully-masked action sequence and iteratively unmasks tokens over K denoising steps via the ReMDM Algorithm 1 posterior (per-token Bernoulli unmasking), while ReMDM remasking lets committed tokens be re-predicted for plan refinement.
Two independent training pipelines are compared head-to-head, both supervised by the built-in BFS oracle: --mode online runs DAgger from scratch (primary), --mode offline behaviour-clones a pre-collected oracle dataset. Either output is scored with --mode inference. See Architecture and Diffusion for details.
Setup
Prerequisites: Python 3.12+, uv. nle compiles from source on macOS.
Linux GPU use needs NVIDIA driver >= 580 for CUDA 13, or >= 525 with --extra cuda12.
# macOS (arm64)
brew install cmake
# Linux (x86_64), if building from source
sudo apt-get install build-essential cmake bison flex libbz2-dev
git clone https://github.com/ANONYMOUS/remdm-planners.git
cd remdm-planner-minihack
# Default. macOS gets the PyPI build (MPS); Linux gets PyPI's CUDA 13.0 build.
# Installs the dev group (pytest) too.
uv sync
# Linux, explicit CUDA 13.2 (driver >= 580)
uv sync --extra cuda13
# Linux, CUDA 12.6 fallback (driver >= 525, or Maxwell/Pascal cards)
uv sync --extra cuda12
Extras: cuda13 and cuda12 are mutually exclusive and Linux-only. Neither is needed on a
modern driver β plain uv sync already yields a CUDA 13.0 build on Linux; use cuda12 only
if nvidia-smi reports a driver older than 580.
Install path must not contain spaces. MiniHack's
mh_patch_nhdat.shinterpolates paths unquoted and fails silently on whitespace, leaving every environment as the same default level with no goal staircase.src/envs/minihack_env.pydetects this and substitutes a Python implementation, but a space-free path avoids it entirely.
Repo layout
remdm-planner-minihack/
βββ configs/ Experiment configs (defaults.yaml + presets, see Configuration)
βββ environments/ Custom .des scenario files (empty; user-supplied)
βββ src/ Model, diffusion, envs, planner pipelines
βββ experiments/
β βββ rl_finetuning/ RL fine-tuning ablation suite (run_ablations.py)
βββ scripts/ HF upload utilities, DAgger and ablation profilers
βββ tests/ Smoke suite β uv run pytest
βββ checkpoints/ Gitignored β offline/, online/ (see Checkpoints)
βββ results/ Gitignored, created on demand β inference/ eval JSONs and
β paper_figures/ manuscript PDFs, both published (see Checkpoints)
βββ demo_minihack.ipynb Demo notebook
βββ main.py CLI entry point
βββ pyproject.toml uv project β deps, cuda extra, dev group
Quickstart
Collects a few oracle trajectories, trains under a 5k env-step budget, prints ID evaluation. A few minutes on CPU.
python main.py --mode smoke
Training
Two independent training methods; neither depends on the other. An offline BC checkpoint can warm-start DAgger via --checkpoint, but this was not used for the paper results.
Online DAgger (primary)
python main.py --mode online # full paper recipe (defaults.yaml)
python main.py --mode online --config configs/final_minihack_gpu_24gb.yaml
python main.py --mode online --override total_timesteps=1000000 --override dagger_lr=0.0001
python main.py --mode online --checkpoint checkpoints/iter600.pth # resume
python main.py --mode online --checkpoint checkpoints/iter600.pth --no-warm-start
Per iteration: curriculum-sampled model rollouts, BFS oracle labelling on the same seeds, efficiency filtering into the replay buffer, then grad_steps_per_iteration gradient steps β see DAgger training loop.
Offline BC
First collect a dataset, then train on it:
python main.py --mode collect # 5000 eps/env -> data/dataset.pt
python main.py --mode collect --data data/small.pt --override collect_episodes_per_env=2000
python main.py --mode offline --data data/dataset.pt
python main.py --mode offline --data data/dataset.pt --override total_timesteps=500000
# Resume (restores optimizer, scheduler, step counter, W&B run)
python main.py --mode offline --data data/dataset.pt --checkpoint checkpoints/offline_step40000.pth
Gradient steps default to total_timesteps // offline_batch_size; ID + OOD eval runs on the id_eval_every_timesteps / ood_eval_every_timesteps cadence. The offline_*_grad_steps keys override that in grad-step units β see the hazard note under Configuration.
Evaluation from a checkpoint
python main.py --mode inference --checkpoint checkpoints/iter600.pth # all ID + OOD
python main.py --mode inference --checkpoint wandb:entity/project/checkpoint-iter600:latest
# Specific environments, save JSON
python main.py --mode inference --checkpoint checkpoints/iter600.pth \
--envs MiniHack-Room-Random-5x5-v0 MiniHack-MazeWalk-45x19-v0 \
--episodes 100 --output results/inference/eval.json
python main.py --mode inference --checkpoint checkpoints/iter600.pth \
--des environments/<your_level>.des # custom .des scenarios (dir ships empty)
python main.py --mode inference --checkpoint checkpoints/iter600.pth --no-ema
--checkpoint accepts a local .pth path or a wandb: artifact reference (wandb:entity/project/name:version). Inference uses EMA weights unless --no-ema is given.
Write eval JSONs into results/inference/ (created for you): scripts/hf_upload.py publishes every JSON it finds there.
Match the config to the checkpoint. The model is built from the config, not the checkpoint, and a mismatch raises at load. Always evaluate with the checkpoint's own config snapshot:
DIR=checkpoints/online/Minihack-Online-Diffusion-DAgger-100M
python main.py --mode inference --config $DIR/config.yaml --checkpoint $DIR/iter563.pth
Baselines and ablations
RL and imitation baselines
Six algorithms: SB3 discrete-action RL (ppo, a2c, dqn, ppo-rnn), Behavioural Cloning (bc) on oracle demos, and a causal Decision Transformer (dt). All share total_timesteps, so numbers are comparable to DAgger and offline BC. Hyperparameters live under the baselines_* config namespace; outputs go to baselines_output_dir.
python main.py --mode baselines --algo ppo # any of the six
python main.py --mode baselines --algo dqn --seeds 0 1 2 # explicit seeds
python main.py --mode baselines --algo bc --num-seeds 3 # or a seed count
python main.py --mode baselines --algo ppo --output results/ppo.json
Architecture ablations
# Local-only planner (no global stream, no goal head), trained from scratch
python main.py --mode online --config configs/ablation_local_only.yaml
# Blind-global: zero the global observation of a trained dual-stream model at eval
python main.py --mode inference --checkpoint checkpoints/iter600.pth --blind-global
RL fine-tuning ablation suite
26 registered ablations (same names as in the craftax repo). See experiments/README.md.
python experiments/rl_finetuning/run_ablations.py --list
python experiments/rl_finetuning/run_ablations.py --checkpoint path/to/ckpt.pth --all
python experiments/rl_finetuning/run_ablations.py \
--checkpoint wandb:entity/project/checkpoint-iter600:latest \
--ablations baseline_rl kl_penalty --fast
Pass --emit-tex-macros to also write tables/results.tex, one \newcommand per
headline quantity, so the manuscript cites generated numbers instead of retyping them.
Macros from this repository are prefixed mh and the sibling suite's rw, so both
files can be \input together; the name-mangling rule is shared between the two.
--measure-gdelta splits the return-weighted ELBO gradient into its imitation and return
terms at one parameter point. No training, no GPU needed; results land in gdelta/ beside
the run's results.json β see experiments/README.md.
python experiments/rl_finetuning/run_ablations.py --measure-gdelta --gdelta-seeds 0 1 2 \
--checkpoint path/to/ckpt.pth \
--results-path outputs/minihack_ablations/results.json \
--output-dir outputs/minihack_ablations
Configuration
One YAML config holds the experiment; the CLI holds the run.
Precedence, lowest to highest: configs/defaults.yaml < --config preset < --override and run flags. Exactly two config layers β a preset never inherits from another preset.
configs/defaults.yaml: the shared final paper recipe, not a cheap baseline. Both machines train exactly this; running with no--configtrains it too.- Config files (
configs/*.yaml): deep-merged ontodefaults.yaml, so presets contain only their deltas β never re-state a default value. - Run flags:
--seed,--checkpoint,--data,--output,--episodes,--envs, mode switches. --override KEY=VALUE(repeatable): keys are validated againstdefaults.yamland cast to the key's type, so a typo is an error, not a silent no-op.
Hazard when writing a preset. Four keys silently override an env-step-derived value when non-null, and
defaults.yamlnow sets all four as part of the recipe:offline_total_grad_steps,offline_eval_every_grad_steps,offline_checkpoint_every_grad_steps,offline_buffer_capacity. A preset that wants its owntotal_timestepsto govern the offline budget must pin them back to explicitnullβ omitting them inherits the pins. Left unpinned,smoke.yamlwould train 60,000 offline gradient steps instead of 19.tests/test_config.pyenforces the pins for every preset that derives its own budget.
| Preset | Purpose |
|---|---|
configs/defaults.yaml |
Shared final paper recipe β the full run both clusters train |
configs/smoke.yaml |
Smoke test (total_timesteps=5000, small buffer, W&B off) |
configs/ablation_local_only.yaml |
Local-only planner ablation (use_global_stream: false) |
configs/gpu_24gb_bigger_model.yaml |
GPU-24GB, larger model (384D, 6 heads) |
configs/gpu_24gb_learning_behaviour.yaml |
GPU-24GB learning-behaviour study (eta=0.18, B=6144) |
configs/final_minihack_gpu_h200.yaml |
Paper run, H200. Machine values only: worker counts (32) and dataset path |
configs/final_minihack_gpu_24gb.yaml |
Paper run, RTX 3090 Ti. Machine values only: dataset path (workers stay at the default 8) |
Key hyperparameters are documented inline in configs/defaults.yaml; the appendix tabulates them.
Checkpoints
Training writes to a unique run directory under checkpoint_dir (default checkpoints/), named {tag}_{YYYYMMDD}_{HHMMSS}_{hex4}. DAgger saves iter{N}.pth on the checkpoint_every_timesteps cadence; offline BC saves offline_step{N}.pth and offline_final.pth. With use_wandb on they also upload as versioned W&B artifacts (type model). Every checkpoint stores its W&B run ID, so passing it back via --checkpoint resumes the same curve.
checkpoints/ is gitignored. Released weights live on the Hugging Face Hub: AnonMLuser/remdm-planner-minihack
| Directory | Method | Selected at | Sample-equivalents |
|---|---|---|---|
checkpoints/online/Minihack-Online-Diffusion-DAgger-100M |
DAgger (main result) | iter563 |
100M |
checkpoints/offline/Minihack-Offline-Diffusion-BC-100M |
Offline BC baseline | offline_step50000 |
100M |
# All checkpoints
uv run hf download AnonMLuser/remdm-planner-minihack --include "checkpoints/**" --local-dir .
# One checkpoint
uv run hf download AnonMLuser/remdm-planner-minihack \
--include "checkpoints/online/Minihack-*/**" --local-dir .
Keep the --include. The Hub repo carries its own README.md (the generated model card), LICENSE and .gitattributes; dropping the glob and pulling into --local-dir . overwrites this repository's copies of all three. To fetch everything, add --exclude "README.md" "LICENSE" ".gitattributes", or use a separate --local-dir. Publishing is safe either way β hf_upload.py stages LICENSE and the demo README.md from git, not the working tree.
Each released directory ships <step>.pth (full training state), model.safetensors (EMA weights only, no pickle), config.yaml and selection.json. The -100M suffix counts sample-equivalents, not env steps β the runs behind these train 5,650,000 env steps. See Checkpoint format for the .pth schema and programmatic loading.
Historical note: the released DAgger selection.json records "every": null, "configured_max": null and "unit": "dagger_iterations", written by a selection() that read two since-renamed config keys. It is historical and noncanonical and stays as published (author decision 2026-08-17); the checkpoint's own config_<step>.yaml carries the real cadence and budget. Current code records the candidate set in env steps β "every": 940000, "configured_max": 5650000 for the shipped recipe β and raises rather than writing a null.
Publishing to the Hub
scripts/hf_upload.py rediscovers and uploads four things, each keeping its repo-relative path: checkpoints/ (adding a model.safetensors EMA export and selection.json per directory), every experiments/rl_finetuning/outputs/<run>/ holding a results.json (with diagnosis.md, tables/, figures/, gdelta/), the eval JSONs in results/inference/, and the manuscript figure PDFs in results/paper_figures/. It drops W&B and hub config keys, shortens absolute paths and regenerates the model card.
HF_TOKEN=hf_xxx uv run python scripts/hf_upload.py --repo-id <ANON_HF_REPO_ID> \
--selection-metric "mean ID+OOD win rate" --dry-run
--dry-run prints the staged tree and card without uploading; drop it to upload. --selection-metric records what the best-of-N checkpoints were chosen on. Also --inference-results <FILE|DIR> ... (eval JSONs kept elsewhere), --private, --yes. Publish one model per directory, with a single .pth and config.
The manuscript figures are built by the sibling repo. Each one puts Craftax Classic and MiniHack side by side, so ../remdm-planner-craftax/scripts/paper_figures.py reads both repositories' ablation results.json and neither can build them alone. Copy the PDFs it emits into results/paper_figures/ here; both Hub repos publish the same set, and the upload warns when they are absent rather than passing over them silently.
A hf download --local-dir . overwrites README.md and LICENSE in the working tree. Publishing is unaffected β hf_upload.py stages LICENSE from git cat-file blob HEAD:LICENSE, and hf_upload_demo.py its bundle's README.md, warning if git cannot be consulted β but restore your own files with git checkout -- README.md LICENSE, or avoid the clobber with the download flags above.
Checkpoint discovery expects the released layout, checkpoints/<role>/<name>/*.pth. A training run writes to its own checkpoints/dagger_<timestamp>/, so copy what you mean to release into checkpoints/{offline,online}/<name>/ first, or nothing is staged. checkpoints/hf/ is skipped β that is where a Hub download lands, and publishing from it would nest already-published artefacts under checkpoints/hf/checkpoints/....
Results, citation, licence
Results tables and the full method description are in Return-Weighted ELBO Fine-Tuning Degrades Masked Diffusion Planners (under submission); demo_minihack.ipynb reproduces the headline comparison. Citation to be added on publication. Licence: MIT, see LICENSE.
Appendix: benchmark-specific detail
Environments
| In-distribution (training) | Out-of-distribution (zero-shot eval) |
|---|---|
MiniHack-Room-Random-5x5-v0 (small random room) |
MiniHack-Room-Dark-15x15-v0 (dark room) |
MiniHack-Room-Random-15x15-v0 (large random room) |
MiniHack-Corridor-R5-v0 (five-room corridor) |
MiniHack-Corridor-R2-v0 (two-room corridor) |
MiniHack-MazeWalk-45x19-v0 (large maze) |
MiniHack-MazeWalk-9x9-v0 (small maze) |
Architecture
LocalDiffusionPlannerWithGlobal (~5.2M parameters):
Local stream: 9x9 glyphs -> Embedding(6000,64) -> CNN(64->32->64) -> Linear -> 1 token
Global stream: 21x79 glyphs -> Embedding(6000,32) -> CNN(32->32->64) -> Pool(2,4) -> 8 tokens
Goal head: mean(global) -> MLP -> [B,2] staircase coords (aux loss)
Gate: sigmoid(learnable scalar, init=-3.0) * global_tokens
Action stream: Embedding(14, 256) + timestep_emb(100, 256) + position_emb(64, 256)
Transformer: concat [1 + 8 + 64 = 73 tokens] -> 4-layer encoder (256D, 4 heads, pre-norm)
Output head: last 64 tokens -> Linear(256, 12) -> action logits
Signature: (local_obs, global_obs, noisy_action_seq, t_discrete) -> {"actions": [B,64,12], "goal_pred": [B,2]}.
LocalDiffusionPlanner (no global stream, no goal head) is the ablation_local_only variant. Supported by --mode offline and --mode online; the experiments/ ablation suite assumes the goal head is present.
Diffusion
- Forward process (MDLM): each action token is independently replaced with
MASK(12) with probability1 - alpha(t),alpha(t)linear or cosine. PAD (13) is never masked. - Loss: continuous-time MDLM NELBO: per sample
w(t) * sum_masked(CE) / Lwithw(t) = -alpha'(t) / (1 - alpha(t))clipped to[0, 1000]; optionallabel_smoothing. - Greedy sampling: used for DAgger collection. Same MaskGIT loop, argmax decoding, no temperature/top-K/remasking,
diffusion_steps_collectsteps.
Reverse sampling (ReMDM Algorithm 1), over K steps (default 10). Per step: predict logits, apply temperature and top-p filtering, sample, and record each committed token's decode probability psi; unmask each masked position independently with posterior probability (alpha_s - (1 - sigma) alpha_t) / (1 - alpha_t); remask each committed position with probability sigma from the configured Section-4.1 schedule. A final greedy cleanup commits anything still masked.
| Strategy | Formula | Description |
|---|---|---|
rescale |
p = eta * sigma_max |
Proportional to noise level |
cap |
p = min(eta, sigma_max) |
Fixed upper bound |
conf |
p = softmax(-confidence) * eta * sigma_max |
Low-confidence tokens remasked more |
Key hyperparameters
configs/defaults.yaml is authoritative and commented inline. Tabulated here are the
keys that change a result, carry a hazard, or are named elsewhere in this README.
Model. n_embd 256, n_head 4, n_layer 4, n_global_tokens 8, seq_len 64,
dropout 0.0, global_gate_init -3.0 β the shape every released checkpoint carries
(see Architecture); a checkpoint restores only against a matching
config. Two model keys are result-affecting in their own right:
| Parameter | Default | Description |
|---|---|---|
ema_decay |
0.999 | EMA smoothing for inference weights; an absolute update count, see the budget caveat below |
use_global_stream |
true | false builds the local-only ablation variant |
Diffusion
| Parameter | Default | Description |
|---|---|---|
noise_schedule |
linear |
linear, cosine, or cosine_sq (MDLM App E.1 naming) |
num_diffusion_steps |
100 | Discrete timestep resolution |
diffusion_steps_eval |
10 | Denoising iterations at inference |
diffusion_steps_collect |
5 | Denoising iterations during collection |
remask_strategy |
conf |
rescale, cap, or conf |
eta |
0.15 | Remasking strength |
temperature |
0.5 | Sampling temperature |
top_p |
0.9 | Nucleus threshold (ReMDM Sec 5) |
replan_every |
16 | Env steps before replanning; the actions already executed in the current plan window are locked into the new plan (inpainting) |
loss_weight_clip |
1000.0 | NELBO weight clip bound |
label_smoothing |
0.0 | Cross-entropy label smoothing (0 = exact ELBO) |
Training budget (unified). Offline BC, DAgger and the SB3 baselines share one env-step budget. This is the only knob that should change to scale a run.
| Parameter | Default | Description |
|---|---|---|
total_timesteps |
5,650,000 | Shared env-step budget |
id_eval_every_timesteps |
470,000 | ID eval cadence |
ood_eval_every_timesteps |
470,000 | OOD eval cadence |
checkpoint_every_timesteps |
940,000 | Checkpoint cadence |
- Offline BC: gradient steps =
total_timesteps // offline_batch_size, and the cosine LRT_maxderives from the same quantity, so any run length decays to the 10% floor at its end. - DAgger: tracks cumulative
env.step()calls (model + oracle) and halts attotal_timesteps.episodes_per_iterationandgrad_steps_per_iterationset the collect/train ratio and must not scale with the budget. - Caveat,
ema_decay: an absolute-update-count constant (half-life ~1 / (1 - decay)steps). Shiftingtotal_timestepsby more than ~2x changes the fraction of training the EMA window covers; set a matching decay by hand for very short or long runs.
Offline grad-step pins. These four override the env-step-derived budget whenever
non-null, and defaults.yaml sets all four. A preset whose own total_timesteps should
govern must pin them back to explicit null β see the hazard note under
Configuration.
| Parameter | Default | Description |
|---|---|---|
offline_total_grad_steps |
60000 | Total gradient steps, overriding total_timesteps // offline_batch_size |
offline_eval_every_grad_steps |
5000 | Eval cadence in grad steps |
offline_checkpoint_every_grad_steps |
10000 | Checkpoint cadence in grad steps |
offline_buffer_capacity |
1500000 | Offline replay capacity |
Training
| Parameter | Default | Description |
|---|---|---|
offline_lr / dagger_lr |
0.0003 / 0.00003 | BC LR (cosine-decayed to 10%) and DAgger LR (constant) |
offline_batch_size / dagger_batch_size |
2048 / 2048 | Batch size per pipeline |
offline_grad_clip / dagger_grad_clip |
1.0 / 1.0 | Gradient norm clip per pipeline |
weight_decay |
0.0 | AdamW weight decay (core training; the ablation suite keeps 1e-4) |
grad_steps_per_iteration |
100 | Gradient steps per DAgger iteration |
episodes_per_iteration |
30 | Episodes per DAgger iteration |
aux_loss_weight |
0.5 | Auxiliary goal loss weight |
buffer_capacity |
10000 | Replay buffer size (windows) |
efficiency_multiplier |
1.5 | DAgger efficiency filter threshold |
Collection, evaluation, performance
| Parameter | Default | Description |
|---|---|---|
collect_episodes_per_env |
5000 | Oracle episodes per ID environment |
eval_episodes_per_env |
50 | Episodes per env at eval (per-run: --episodes) |
use_amp |
true | Mixed precision via torch.amp; see Performance tuning |
torch_compile |
true | torch.compile the model |
checkpoint_dir |
checkpoints |
Root for per-run checkpoint directories |
seed |
null | RNG seed (null = random; per-run: --seed) |
Worker counts (collect_num_workers, num_collection_workers, both 8) are machine values.
The collect_output, use_wandb, wandb_* and offline_log_every keys mirror the run
flags under Configuration, the curriculum_* keys the behaviour under
DAgger training loop, and the 21 baselines_* keys hold the
SB3/BC/DT hyperparameters; all are commented where they are declared.
DAgger training loop
- Curriculum sampling: pick an environment weighted by difficulty (low win-rate sampled more).
- Model rollout: EMA model, greedy sampling, replanning every 16 steps with the executed prefix locked,
episodes_per_iterationepisodes. - Oracle rollout: BFS oracle on the same seed.
- Efficiency filter: add the oracle trajectory if the model failed or took >1.5x the oracle's steps.
- Budget accounting:
env_steps_total += model_steps + oracle_steps; halt attotal_timesteps. - Training: sample the buffer, run
grad_steps_per_iterationsteps, update EMA after each.
BFS oracle priority: (1) kick adjacent doors, (2) BFS to staircase, (3) BFS to frontier, (4) BFS to farthest tile, (5) random cardinal.
Reward shaping
| Component | Value | Condition |
|---|---|---|
| Win bonus | +20.0 | Episode won |
| BFS progress | +0.5 * (prev_dist - curr_dist) | Closer to staircase |
| Exploration | +0.05 | New tile visited |
| Step penalty | -0.01 | Every step |
Checkpoint format
# DAgger (iter{N}.pth)
{
"model_state_dict": ..., "ema_state_dict": ...,
"optimizer_state_dict": ..., "scheduler_state_dict": ...,
"curriculum_state": {...},
"iteration": int,
"env_steps": int, # cumulative env.step() calls
"wandb_run_id": str | None,
"rng_states": {"torch", "numpy", "python"},
}
Offline BC step-level (offline_step{N}.pth, when checkpoint_every_timesteps > 0) is the
same minus curriculum_state, with step for iteration and env_steps = step * offline_batch_size. Its rng_states is required: resume raises without it. The final
offline_final.pth carries only model_state_dict, ema_state_dict and wandb_run_id.
Load programmatically
# Inference, from safetensors (already EMA weights)
from safetensors.torch import load_file
from src.config import load_config
from src.models.denoiser import make_model
DIR = "checkpoints/online/Minihack-Online-Diffusion-DAgger-100M"
cfg = load_config(f"{DIR}/config.yaml")
model = make_model(cfg)
model.load_state_dict(load_file(f"{DIR}/model.safetensors"))
model.eval()
From the full .pth instead, to resume or to pick training over EMA weights: torch.load(..., weights_only=False), then model.load_state_dict(ckpt["model_state_dict"]) and, for what evaluation uses by default, ModelEMA(model, decay=cfg.ema_decay) with load_state_dict(ckpt["ema_state_dict"]) and apply_to(model).
W&B artifacts and run resumption
W&B model artifacts contain the .pth and a config.yaml snapshot; the reference format is wandb:entity/project/artifact-name:version, version being latest, v0, v1. Resuming reads the run ID out of the checkpoint and passes it to wandb.init(resume="must"), so curves continue with no gaps.
# Automatic. A checkpoint predating the feature needs the ID passing by hand:
python main.py --mode online --checkpoint old.pth --override wandb_resume_id=abc123xyz
W&B metric namespaces
Declared in src/planners/logging.py; the key lists there are authoritative.
| Namespace | Contents |
|---|---|
diffusion/ |
loss, loss_diff, loss_aux |
train/ |
Buffer, collection and optimiser state β 13 keys including model_steps, oracle_steps, efficiency_ratio, global_gate, env_steps |
speed/ |
Per-iteration timings, throughput and gpu_memory_mb |
model/ |
param_norm, param_drift_from_init, ema_gate_value (every 10 iters) |
eval_id/{env}/, eval_ood/{env}/ |
Per-env win_rate, wins, avg_reward, avg_steps, n_episodes |
eval_id/, eval_ood/ |
mean_win_rate |
ckpt_eval_id/, ckpt_eval_ood/, ckpt_eval/ |
The same, at checkpoint time |
curriculum/{env}/ |
win_rate per training environment |
offline/ |
final_loss, total_steps, total_timesteps (summary only) |
inference/{env}/ |
Per-env metrics from --mode inference |
DAgger and offline BC both emit to eval_id/ and eval_ood/, through the same Evaluator and EMA-weight path.
Performance tuning
use_amp (default true) puts torch.amp.autocast("cuda") + GradScaler in both trainers β
roughly 2x on gradient steps, with loss and win rates statistically equivalent to FP32, and a
no-op on CPU. torch_compile (default true) shows no measured gain beyond AMP.
num_collection_workers affects only the threaded CPU fallback: collection auto-selects
GPU-batched (CUDA, episodes_per_iteration > 1) > threaded CPU > sequential.
Profile with python scripts/profile_dagger.py [--override key=value ...].
Testing
uv run pytest # the default suite
uv run pytest -m slow # slow entry points only (BC + PPO baselines)
A CPU-only suite, 17 modules. Tiny synthetic data and a shrunken model throughout β no real checkpoints, datasets or network calls, and nothing written outside tmp_path. conftest.py forces CPU and disables W&B; slow marks the multi-second CLI smokes and is deselected by default. For a quality signal, use --mode smoke.
| File | Covers |
|---|---|
test_smoke_src.py, test_smoke_experiments.py |
that things run: imports, model from the real config, a forward pass of the expected shape and dtype with no NaNs, a finite training step, save/reload identity, every CLI entry point, and all 26 registry ablations |
test_spec_*.py, test_method_spec*.py |
that things are correct: each canonical statement of the parent workspace's the spec *.md pinned against the implementation |
test_config.py, test_recipe_values.py |
the preset, delta-only and poolability rules, and the shipped recipe values |
test_gdelta.py, test_tex_macros.py |
the --measure-gdelta decomposition, and the --emit-tex-macros output: definitions only, uniquely named, letters only |
test_ablation_perf.py, test_gpu_step_perf.py |
measured throughput expectations |
test_env_reuse.py, test_failure_behaviour.py |
MiniHack env pooling, and failures that must raise rather than be swallowed |
test_gpu_agreement.py |
CPU/GPU agreement, skipped without a device |
Implementation notes
- MDLM loss returns
0.0(not NaN) when no masked positions exist. NELBO-weighted per MDLM eq (10). - PAD tokens are never masked and are excluded from the loss.
- EMA updates after every gradient step, not per iteration.
DataCollectorsyncs EMA weights before each rollout. - Curriculum starts from a 50/50 prior per environment and buckets the rolling win-rate:
[0, 0.15)-> 0.2,[0.15, 0.85)-> 1.0,[0.85, 1.0]-> 0.1. - Replay buffer pins offline data at the front; only online samples are FIFO-evicted. Returns
Nonewhen empty. - Global gate starts at
sigmoid(-3.0) ~ 0.047, nearly closed, so the global stream cannot destabilise early training. - DAgger warm-start: iteration 0 seeds the buffer with 3 oracle trajectories per ID environment (12 total).