Refresh artefacts; align notebook numbers with the updated paper

.gitattributes

@@ -120,3 +120,7 @@ experiments/rl_finetuning/outputs/minihack_ablations/figures/t_bin_grad_norms_tr
 experiments/rl_finetuning/outputs/minihack_ablations/figures/t_bin_norms_heatmap.png filter=lfs diff=lfs merge=lfs -text
 experiments/rl_finetuning/outputs/minihack_ablations/figures/t_distribution_analysis.png filter=lfs diff=lfs merge=lfs -text
 experiments/rl_finetuning/outputs/minihack_ablations/figures/win_rate_and_effective_batch_size.png filter=lfs diff=lfs merge=lfs -text
+ablation_assets/eval_scores_over_training.png filter=lfs diff=lfs merge=lfs -text
+ablation_assets/score_delta_over_baseline_rl.png filter=lfs diff=lfs merge=lfs -text
+experiments/rl_finetuning/outputs/minihack_ablations/figures/per_layer_grad_heatmap_layer_ablation_top1.png filter=lfs diff=lfs merge=lfs -text
+experiments/rl_finetuning/outputs/minihack_ablations/figures/score_delta_over_baseline_rl.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,272 @@
+# ============================================================
+# Project-specific
+# ============================================================
+tmp/
+wandb/
+res/
+runs/
+data/
+checkpoints/
+!checkpoints/*.pth
+play_data/
+texture_cache.pbz2
+texture_cache*.pbz2
+checkpoints*/
+artifacts/
+*.pth
+*.pt
+
+# ============================================================
+# Byte-compiled / optimized / DLL files
+# ============================================================
+**/__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# ============================================================
+# Distribution / packaging
+# ============================================================
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# ============================================================
+# PyInstaller
+# ============================================================
+*.manifest
+*.spec
+
+# ============================================================
+# Installer logs
+# ============================================================
+pip-log.txt
+pip-delete-this-directory.txt
+
+# ============================================================
+# Unit test / coverage reports
+# ============================================================
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# ============================================================
+# Translations
+# ============================================================
+*.mo
+*.pot
+
+# ============================================================
+# Django
+# ============================================================
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# ============================================================
+# Flask
+# ============================================================
+instance/
+.webassets-cache
+
+# ============================================================
+# Scrapy
+# ============================================================
+.scrapy
+
+# ============================================================
+# Sphinx documentation
+# ============================================================
+docs/_build/
+
+# ============================================================
+# PyBuilder
+# ============================================================
+.pybuilder/
+target/
+
+# ============================================================
+# Jupyter / IPython
+# ============================================================
+.ipynb_checkpoints
+profile_default/
+ipython_config.py
+
+# ============================================================
+# Package managers / lock files
+# (lock files are generally recommended for version control;
+#  uncomment to ignore them if your workflow requires it)
+# ============================================================
+
+# pyenv
+# .python-version
+
+# pipenv
+# Pipfile.lock
+
+# uv
+# uv.lock
+
+# poetry
+# poetry.lock
+# poetry.toml
+
+# pdm
+# pdm.lock
+# pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+# pixi.lock
+.pixi/
+
+# PEP 582
+__pypackages__/
+
+# ============================================================
+# Task queues / message brokers
+# ============================================================
+
+# Celery
+celerybeat-schedule
+celerybeat.pid
+
+# Redis
+*.rdb
+*.aof
+*.pid
+
+# RabbitMQ
+mnesia/
+rabbitmq/
+rabbitmq-data/
+
+# ActiveMQ
+activemq-data/
+
+# ============================================================
+# Environments
+# ============================================================
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# ============================================================
+# Type checkers / linters
+# ============================================================
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre
+.pyre/
+
+# pytype
+.pytype/
+
+# Ruff
+.ruff_cache/
+
+# ============================================================
+# Misc Python tools
+# ============================================================
+
+# SageMath
+*.sage.py
+
+# Cython debug symbols
+cython_debug/
+
+# Spyder
+.spyderproject
+.spyproject
+
+# Rope
+.ropeproject
+
+# Abstra
+.abstra/
+
+# mkdocs
+/site
+
+# ============================================================
+# Notebooks / interactive tools
+# ============================================================
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/
+
+# Streamlit
+.streamlit/secrets.toml
+
+# ============================================================
+# IDEs
+# (kept as comments — commit these if your whole team uses
+#  the same editor, otherwise add to your global gitignore)
+# ============================================================
+
+# PyCharm / JetBrains
+.idea/
+
+# VS Code
+# .vscode/
+
+# ============================================================
+# PyPI
+# ============================================================
+.pypirc
+
+# ============================================================
+# OS-specific
+# ============================================================
+
+# macOS
+.DS_Store
+
+__pycache__/
+*.pyc
+
+CHANGES.md
+CLI_CHANGES.md
+
+# Claude Code: per-machine settings stay local; the rest of .claude/ is shared
+.claude/settings.local.json

README.md

@@ -1,63 +1,600 @@
 ---
-license: mit
+license: apache-2.0
+library_name: pytorch
+pipeline_tag: reinforcement-learning
 tags:
-  - reinforcement-learning
-  - discrete-diffusion
-  - planning
-  - minihack
-  - pytorch
+- reinforcement-learning
+- planning
+- discrete-diffusion
+- remdm
+- minihack
+- nethack
+- pytorch
 ---
 
-# ReMDM Planner for MiniHack — anonymous artefact release
+# ReMDM Planner — MiniHack artefacts
 
-Anonymous companion artefacts for *Return-Weighted ELBO Fine-Tuning of Discrete Diffusion Planners Reduces to Self-Imitation*, submitted for double-blind review.
-Code: <https://github.com/vtp5tycwm2-art/remdm-planners>
+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.
 
-This repo is a self-contained bundle: the full project source, the released
-checkpoints, a stripped inference-only checkpoint, and the pre-computed RL
-fine-tuning ablation assets. `demo_minihack.ipynb` in the code repo pulls this
-whole repo with one `snapshot_download` and needs no authentication.
+**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`.
 
-## Contents
+## What is here
 
 ```
-src/ configs/ environments/ experiments/ scripts/ tests/   project source
-checkpoint_inference.pth                                   EMA weights only (~21 MB), what the notebook loads
-checkpoints/online/Minihack-Online-Diffusion-DAgger-100M   DAgger planner (iter563) + EMA safetensors + config
-checkpoints/offline/Minihack-Offline-Diffusion-BC-100M     offline BC baseline (step 50000)
-ablation_assets/                                           flat PNGs + CSVs the notebook displays
-experiments/rl_finetuning/outputs/minihack_ablations/      full ablation figures + tables
+src/  configs/  environments/  experiments/  scripts/  tests/
+main.py  pyproject.toml  uv.lock  LICENSE  demo_minihack.ipynb
+
+checkpoint_inference.pth   Stripped EMA weights only (~21 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:
+                           8 PNGs + 4 CSVs + results.json
+experiments/rl_finetuning/outputs/minihack_ablations/
+  figures/  tables/  results.json
+results/                   Aggregated CSVs and single-seed inference evaluations
 ```
 
-`checkpoint_inference.pth` holds only `ema_state_dict` — no optimiser,
-scheduler, RNG or curriculum state. The `-100M` suffix counts
-**sample-equivalents, not env steps**; the run behind it trained 5,650,000 env
-steps.
+## Headline numbers (MiniHack, ID win rate)
+
+| | Win rate |
+|---|---|
+| Pretrained (DAgger) checkpoint, ablation-harness protocol | **47.5%** |
+| Baseline return-weighted ELBO fine-tuning, 500 iters | 43.8% ± 6.1 |
+| Best of 25 conditions (head-only) | 49.6% ± 3.9 |
+| Worst of 25 conditions (normalised advantages) | 12.1% ± 4.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 reported as a matching ordering, not
+as a second confirmation; the claim rests on Craftax Classic, where the pretrained
+checkpoint scores 11.81 and no condition of the 25 recovers it.
+
+Direct evaluation of the same checkpoint at 50 episodes per environment gives 48.5%
+in distribution and 4.7% on the three held-out layouts (paper Table 6).
 
-## Anonymity
+## Evaluation path
 
-W&B run ids embedded in the `.pth` checkpoints have been cleared, and entity and
-project names removed from the result metadata. Machine names appear as
-`gpu_24gb` and `gpu_h200`. Weight tensors are byte-identical to the originals.
+MiniHack has a single evaluation path. `demo_minihack.ipynb` calls
+`Evaluator._run_episodes_batched` from `src/planners/inference.py`, and the ablation
+harness (`experiments/rl_finetuning/ablations/training.py`) imports the same
+`Evaluator` and calls `evaluator.evaluate(...)`. Notebook and paper numbers therefore
+come from one sampler and one replanning rule, differing only in episode count and
+seed offset.
 
-## Limitations
+## Quick start
 
-Research artefacts tied to specific MiniHack environment versions and to the
-cropped-glyph observation encoding. Not general-purpose agents; they will not
-transfer to other environments or to pixel observations. Evaluation results and
-their variance are reported in the paper.
+```python
+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
 
 ```bibtex
-@inproceedings{remdm-planner-minihack,
-  title  = {Return-Weighted ELBO Fine-Tuning of Discrete Diffusion Planners Reduces to Self-Imitation},
-  author = {Anonymous},
-  year   = {2026},
-  note   = {Under review}
+@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}
 }
 ```
 
-## License
+---
+
+# ReMDM Planner for MiniHack
+
+PyTorch implementation of **ReMDM** (Remasking Discrete Diffusion Model) for action-sequence planning in [MiniHack](https://github.com/facebookresearch/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 [`craftax/`](../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 in the accompanying paper (under submission; citation to follow): **online DAgger** under a BFS oracle (primary) and **offline behavioural cloning** on pre-collected oracle datasets. See [Architecture](#architecture) and [Diffusion](#diffusion) for details.
+
+## Setup
+
+Prerequisites: Python 3.12+, [uv](https://docs.astral.sh/uv/). `nle` compiles from source on macOS.
+Linux GPU use needs NVIDIA driver >= 580 for CUDA 13, or >= 525 with `--extra cuda12`.
+
+```bash
+# 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-planners/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.sh` interpolates paths unquoted and fails silently on whitespace, leaving every environment as the same default level with no goal staircase. `src/envs/minihack_env.py` detects this and substitutes a Python implementation, but a space-free path avoids the issue entirely.
+
+## Repo layout
+
+```
+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/inference/      Eval JSONs from --mode inference (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.
+
+```bash
+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)
+
+```bash
+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, `grad_steps_per_iteration` gradient steps. Halts when cumulative env steps reach `total_timesteps`. See [DAgger training loop](#dagger-training-loop).
+
+### Offline BC
+
+First collect a dataset, then train on it:
+
+```bash
+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](#configuration). See the hazard note under [Configuration](#configuration).
+
+## Evaluation from a checkpoint
+
+```bash
+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:
+
+```bash
+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`.
+
+```bash
+python main.py --mode baselines --algo ppo
+python main.py --mode baselines --algo a2c
+python main.py --mode baselines --algo dqn --seeds 0 1 2
+python main.py --mode baselines --algo ppo-rnn
+python main.py --mode baselines --algo bc --num-seeds 3
+python main.py --mode baselines --algo dt --seeds 0 1 2
+python main.py --mode baselines --algo ppo --output results/ppo.json
+python main.py --mode baselines --algo ppo --override total_timesteps=5650000   # match ReMDM online budget
+```
+
+### Architecture ablations
+
+```bash
+# 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
+
+25 registered ablations (same names as in the craftax repo). See `experiments/README.md`.
+
+```bash
+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
+```
+
+## Configuration
+
+One YAML config holds the experiment; the CLI holds the run.
+
+- **`configs/defaults.yaml`**: the **shared final paper recipe**, not a cheap baseline. Both clusters train exactly this; running with no `--config` trains it too.
+- **Config files** (`configs/*.yaml`): any file passed via `--config` is deep-merged onto `defaults.yaml`, so presets contain **only their deltas** — never re-state a default value. Presets are a single layer: they never inherit from one another.
+- **CLI flags**: per-invocation values — `--seed`, `--checkpoint`, `--data`, `--output`, `--episodes`, `--envs`, mode switches.
+- **`--override KEY=VALUE`** (repeatable): ad hoc config overrides. Keys are validated against `defaults.yaml` and values are cast to the key's type; a typo is an error, not a silent no-op.
+
+Precedence, lowest to highest: `configs/defaults.yaml` < `--config` file < `--override` and run flags.
+
+> **Hazard when writing a preset.** Four keys silently *override* an env-step-derived value when non-null, and `defaults.yaml` now 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 own `total_timesteps` to govern the offline budget must pin them back to **explicit `null`** — omitting them inherits the pins. Left unpinned, `smoke.yaml` would train 60,000 offline gradient steps instead of 19. `tests/test_config.py` enforces 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` | 24 GB GPU, larger model (384D, 6 heads) |
+| `configs/gpu_24gb_learning_behaviour.yaml` | 24 GB GPU 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](#key-hyperparameters) 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`. Checkpoints also upload as versioned W&B artifacts (type `model`) when `use_wandb` is on. All checkpoints store the W&B run ID, so passing them back via `--checkpoint` resumes the same W&B curve automatically.
+
+`checkpoints/` is gitignored.
+
+Released weights and the pre-computed ablation outputs are mirrored on the
+Hugging Face Hub at [`AnonMLuser/remdm-planner-minihack`](https://huggingface.co/AnonMLuser/remdm-planner-minihack),
+an anonymous account holding the layout below.
+
+```bash
+uv run hf download AnonMLuser/remdm-planner-minihack --include "checkpoints/**" --local-dir .
+```
+
+| 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 |
+
+Each published directory ships `<step>.pth` (full training state), `model.safetensors` (EMA weights only, no pickle), `config.yaml` (config snapshot) 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](#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"`. It was published by a version of `selection()` that read two config keys which had been renamed out of the config, so the values came back empty. 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, so nothing is lost. A publish from the 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 for any key it cannot read.
+
+### Publishing to the Hub
+
+`scripts/hf_upload.py` rediscovers and uploads three 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/`), and the eval JSONs in `results/inference/`. It drops W&B and hub config keys, shortens absolute paths and regenerates the model card.
+
+```bash
+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.
+
+**Checkpoint discovery expects the released layout**, `checkpoints/<role>/<name>/*.pth` — the layout the Hub repo mirrors. A training run writes to its own `checkpoints/dagger_<timestamp>/` directory, so copy the checkpoints 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 push already-published artefacts back up into a nested `checkpoints/hf/checkpoints/...` tree.
+
+## Results, citation, licence
+
+Results tables and the full method description are in the accompanying paper (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 probability `1 - alpha(t)`, `alpha(t)` linear or cosine. PAD (13) is never masked.
+- **Loss:** continuous-time MDLM NELBO: per sample `w(t) * sum_masked(CE) / L` with `w(t) = -alpha'(t) / (1 - alpha(t))` clipped to `[0, 1000]`; optional `label_smoothing`.
+- **Greedy sampling:** used for DAgger collection. Same MaskGIT loop, argmax decoding, no temperature/top-K/remasking, `diffusion_steps_collect` steps.
+
+**Reverse sampling (ReMDM Algorithm 1)**, over `K` steps (default 10):
+
+1. Predict logits; apply temperature and top-p (nucleus) filtering; sample predictions and record each committed token's decode probability `psi`.
+2. **Unmask:** each masked position commits independently with the posterior probability `(alpha_s - (1 - sigma) alpha_t) / (1 - alpha_t)`.
+3. **ReMDM remask:** each committed position re-masks with probability `sigma` from the configured Section-4.1 schedule.
+4. Final step: any remaining masked positions are committed by a greedy cleanup pass.
+
+| 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
+
+**Model**
+
+| Parameter | Default | Description |
+|---|---|---|
+| `n_embd` | 256 | Transformer hidden dimension |
+| `n_head` | 4 | Attention heads |
+| `n_layer` | 4 | Transformer blocks |
+| `n_global_tokens` | 8 | Global stream context tokens |
+| `seq_len` | 64 | Action plan length |
+| `dropout` | 0.0 | Forward masking already regularises |
+| `ema_decay` | 0.999 | EMA smoothing for inference weights |
+| `global_gate_init` | -3.0 | Initial global gate logit |
+| `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 |
+| `physics_aware_sampling` | false | Penalise hazardous actions at inference |
+
+**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`. The cosine LR `T_max` derives 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 at `total_timesteps`. `episodes_per_iteration` and `grad_steps_per_iteration` set 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). Shifting `total_timesteps` by more than ~2x changes the fraction of training the EMA window covers; set a matching decay manually for very short or long runs.
+
+**Training**
+
+| Parameter | Default | Description |
+|---|---|---|
+| `offline_lr` | 0.0003 | BC LR (cosine-decayed to 10%) |
+| `dagger_lr` | 0.00003 | DAgger LR (constant) |
+| `offline_batch_size` | 2048 | Offline BC batch size |
+| `dagger_batch_size` | 2048 | DAgger batch size |
+| `offline_grad_clip` | 1.0 | Gradient norm clip (offline) |
+| `dagger_grad_clip` | 1.0 | Gradient norm clip (DAgger) |
+| `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 |
+| `curriculum_preseed` | true | Pre-seed curriculum with 50/50 prior |
+| `curriculum_queue_size` | 100 | Curriculum window size per environment |
+
+**Collection, evaluation, performance, logging**
+
+| Parameter | Default | Description |
+|---|---|---|
+| `collect_episodes_per_env` | 5000 | Oracle episodes per ID environment |
+| `collect_num_workers` | 8 | Process workers for collection |
+| `collect_output` | `data/dataset.pt` | Collected dataset path (per-run: `--data`) |
+| `eval_episodes_per_env` | 50 | Episodes per env at eval (per-run: `--episodes`) |
+| `checkpoint_eval_episodes` | 50 | Episodes per env at checkpoint eval |
+| `use_amp` | true | Mixed precision via `torch.amp` |
+| `torch_compile` | true | `torch.compile` the model |
+| `num_collection_workers` | 8 | Workers for DAgger collection |
+| `use_wandb` | true | Enable W&B logging |
+| `wandb_project` | `remdm-planner-minihack` | W&B project |
+| `wandb_resume_id` | null | W&B run ID for resumption |
+| `offline_log_every` | 50 | Log frequency (offline steps) |
+| `seed` | null | RNG seed (null = random; per-run: `--seed`) |
+
+## DAgger training loop
+
+1. **Curriculum sampling:** pick an environment weighted by difficulty (low win-rate sampled more).
+2. **Model rollout:** EMA model, greedy sampling, replanning every 16 steps with the executed prefix locked, `episodes_per_iteration` episodes.
+3. **Oracle rollout:** BFS oracle on the **same seed**.
+4. **Efficiency filter:** add the oracle trajectory if the model failed or took >1.5x the oracle's steps.
+5. **Budget accounting:** `env_steps_total += model_steps + oracle_steps`; halt at `total_timesteps`.
+6. **Training:** sample the buffer, run `grad_steps_per_iteration` steps, 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
+
+```python
+# DAgger
+{
+    "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)
+{
+    "model_state_dict": ..., "ema_state_dict": ...,
+    "optimizer_state_dict": ..., "scheduler_state_dict": ...,
+    "step": int,
+    "env_steps": int,                  # step * offline_batch_size
+    "wandb_run_id": str | None,
+}
+
+# Offline BC, final (offline_final.pth)
+{"model_state_dict": ..., "ema_state_dict": ..., "wandb_run_id": str | None}
+```
+
+### Load programmatically
+
+```python
+# 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()
+```
+
+```python
+# From the full .pth, to resume or to pick training vs EMA weights
+import torch
+from src.config import load_config
+from src.models.denoiser import make_model, ModelEMA
+
+DIR = "checkpoints/online/Minihack-Online-Diffusion-DAgger-100M"
+cfg = load_config(f"{DIR}/config.yaml")
+ckpt = torch.load(f"{DIR}/iter563.pth", map_location="cpu", weights_only=False)
+
+model = make_model(cfg)
+model.load_state_dict(ckpt["model_state_dict"])
+
+ema = ModelEMA(model, decay=cfg.ema_decay)
+ema.load_state_dict(ckpt["ema_state_dict"])
+ema.apply_to(model)          # what evaluation uses by default
+model.eval()
+```
+
+### W&B artifacts and run resumption
+
+W&B model artifacts contain the `.pth` and a `config.yaml` snapshot. Reference format is `wandb:entity/project/artifact-name:version`, version being `latest`, `v0`, `v1`.
+
+All training loops store the W&B run ID in their checkpoints. Resuming extracts it and passes it to `wandb.init(resume="must")`, so curves continue with no gaps.
+
+```bash
+python main.py --mode online --checkpoint checkpoints/iter600.pth               # automatic
+# Manual override (checkpoint predates the feature)
+python main.py --mode online --checkpoint old.pth --override wandb_resume_id=abc123xyz
+```
+
+## W&B metric namespaces
+
+| Namespace | Contents |
+|---|---|
+| `diffusion/` | `loss`, `loss_diff`, `loss_aux` |
+| `train/` | `buffer_size`, `buffer_online_frac`, `model_won`, `added_to_buffer`, `episodes_collected`, `model_steps`, `oracle_steps`, `efficiency_ratio`, `lr`, `grad_norm`, `global_gate`, `env_steps`, `progress` |
+| `speed/` | `iter_time_sec`, `collect_time_sec`, `train_step_time_sec`, `samples_per_sec`, `env_steps_per_sec`, `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, avg steps, avg reward |
+| `eval_id/`, `eval_ood/` | `mean_win_rate` |
+| `curriculum/{env}/` | `win_rate` per training environment |
+| `ckpt_eval_id/`, `ckpt_eval_ood/` | Per-env metrics at checkpoint time |
+| `ckpt_eval/` | `id_winrate`, `ood_winrate` |
+| `offline/` | `final_loss`, `total_steps`, `total_timesteps` (summary only) |
+
+DAgger and offline BC both emit to `eval_id/` and `eval_ood/`, through the same `Evaluator` and EMA-weight path.
+
+## Performance tuning
+
+| Key | Default | Effect |
+|---|---|---|
+| `use_amp` | **true** | `torch.amp.autocast("cuda")` + `GradScaler` in both trainers. Roughly 2x on gradient steps, with loss and win rates statistically equivalent to FP32. No-op on CPU |
+| `torch_compile` | **true** | `torch.compile(model, mode="default")`. No measured gain beyond AMP |
+| `num_collection_workers` | 8 | Affects 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
+
+```bash
+uv run pytest            # 15 modules; `slow` deselected by default
+uv run pytest -m slow    # slow entry points only (BC + PPO baselines)
+```
+
+`conftest.py` forces CPU and disables W&B. `test_spec_*.py` and `test_method_spec*.py` pin each canonical statement of the method/training specification against the implementation; `test_config.py` and `test_recipe_values.py` guard the preset, delta-only and poolability rules and the shipped recipe values; `test_ablation_perf.py` and `test_gpu_step_perf.py` hold measured perf expectations. `test_smoke_src.py` and `test_smoke_experiments.py` cover both pipelines: modules import, the model builds from `configs/defaults.yaml`, a forward pass returns the expected shape and dtype with no NaNs, one training step gives a finite loss, save/reload reproduces identical output, each entry point runs, and all 25 registry ablations step. They assert things *run*, not that results are good. CPU-only, seeded, synthetic data; nothing written outside `tmp_path`. For a quality signal, use `--mode smoke`.
+
+## Implementation notes
 
-MIT, see `LICENSE`.
+- **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.
+- **Sampling paths:** evaluation uses stochastic ReMDM (temperature, top-p, remasking, `diffusion_steps_eval`); DAgger collection uses greedy argmax (`diffusion_steps_collect`).
+- **`remdm_sample`** guarantees a fully committed output via a final greedy cleanup of any remaining masked positions.
+- **EMA** updates after every gradient step, not per iteration. `DataCollector` syncs 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 `None` when 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).

ablation_assets/group_summary.csv

@@ -1,6 +1,6 @@
 Group,N,Mean,Best,Worst,StdDev
-Baseline,1,0.5833,0.5833,0.5833,0.0
-A,6,0.6208,0.6833,0.575,0.0372
-B,7,0.5429,0.6917,0.2583,0.1289
-C,7,0.6571,0.6667,0.6333,0.0104
-D,4,0.5917,0.625,0.55,0.0306
+Baseline,1,0.4375,0.4375,0.4375,0.0
+A,6,0.4278,0.4417,0.4125,0.0095
+B,7,0.3857,0.4875,0.1208,0.1164
+C,7,0.4399,0.4958,0.3917,0.0408
+D,4,0.3802,0.4417,0.3208,0.0428

ablation_assets/hypothesis_verdict.csv

@@ -1,26 +1,26 @@
 Method,Group,Score,Delta_Baseline,Verdict,Hypothesis
-action_diversity,D,0.6167,0.0333,NEUTRAL,If diversity filtering helps: degenerate plans corrupt training
-advantage_clip,B,0.5583,-0.025,NEUTRAL,If clipping helps: large advantage magnitudes destabilise training
-attention_only,C,0.6333,0.05,NEUTRAL,"If attention-only works: model needs routing updates, not feature updates"
-baseline_rl,Baseline,0.5833,0.0,NEUTRAL,Diagnoses whether the RL signal alone causes collapse
-bc_wins,B,0.4917,-0.0917,NEUTRAL,If BC on wins helps: the return weighting is the specific cause
-entropy_bonus,B,0.6167,0.0333,NEUTRAL,If entropy bonus helps: collapse is mode-collapse; not a gradient problem
-ewc,A,0.6167,0.0333,NEUTRAL,If EWC helps: forgetting pretrained representations is the proximate cause
-ffn_only,C,0.6583,0.075,IMPROVEMENT,If FFN-only works: stored knowledge (FFN as memory) needs updating; not attention
-frozen_backbone,C,0.6583,0.075,IMPROVEMENT,If frozen backbone helps: deep gradient flow into backbone causes collapse
-gradient_surgery,B,0.6917,0.1083,IMPROVEMENT,If PCGrad helps: gradients are conflicting and resolvable by projection
-head_only,C,0.6583,0.075,IMPROVEMENT,If head-only works: backbone representations are fine; only decision boundary needs updating
-kl_penalty,A,0.5833,0.0,NEUTRAL,If this helps: catastrophic forgetting is the primary cause; soft regularisation suffices
-layer_ablation_top1,C,0.6667,0.0833,IMPROVEMENT,Minimal unfrozen depth needed; collapse depth correlates with gradient flow depth
-layer_ablation_top2,C,0.6667,0.0833,IMPROVEMENT,Minimal unfrozen depth needed; collapse depth correlates with gradient flow depth
-layer_ablation_top3,C,0.6583,0.075,IMPROVEMENT,Minimal unfrozen depth needed; collapse depth correlates with gradient flow depth
-llrd,A,0.6167,0.0333,NEUTRAL,If LLRD helps: deep gradient flow into early layers corrupts representations
-lora,A,0.6833,0.1,IMPROVEMENT,If LoRA works: too many unconstrained degrees of freedom cause collapse
-low_t,B,0.6,0.0167,NEUTRAL,If low-t helps: high-t (coarse-structure) gradients are biased
-mixed_replay,A,0.575,-0.0083,NEUTRAL,If mixed replay helps: online data distribution alone is too corrupted
-normalized_adv,B,0.2583,-0.325,COLLAPSE,If std normalisation helps: simple mean normalisation is too loose
-reward_filtering,D,0.625,0.0417,NEUTRAL,If filtering helps: noisy/low-return data poisons gradients
-reward_model,D,0.55,-0.0333,NEUTRAL,If reward model helps: raw returns are too sparse; learned model smooths signal
-running_stats,D,0.575,-0.0083,NEUTRAL,If running stats help: batch normalisation is too noisy for small batches
-t_curriculum,B,0.5833,0.0,NEUTRAL,If curriculum helps: ordering of learning signals matters
-trust_region_kl,A,0.65,0.0667,IMPROVEMENT,If hard constraint helps: soft KL is insufficient -- a hard boundary is needed
+action_diversity,D,0.375,-0.0625,COLLAPSE,If diversity filtering helps: degenerate plans corrupt training
+advantage_clip,B,0.3375,-0.1,COLLAPSE,If clipping helps: large advantage magnitudes destabilise training
+attention_only,C,0.3958,-0.0417,NEUTRAL,"If attention-only works: model needs routing updates, not feature updates"
+baseline_rl,Baseline,0.4375,0.0,NEUTRAL,Diagnoses whether the RL signal alone causes collapse
+bc_wins,B,0.4583,0.0208,NEUTRAL,If BC on wins helps: the return weighting is the specific cause
+entropy_bonus,B,0.4458,0.0083,NEUTRAL,If entropy bonus helps: collapse is mode-collapse; not a gradient problem
+ewc,A,0.4333,-0.0042,NEUTRAL,If EWC helps: forgetting pretrained representations is the proximate cause
+ffn_only,C,0.4375,0.0,NEUTRAL,If FFN-only works: stored knowledge (FFN as memory) needs updating; not attention
+frozen_backbone,C,0.4667,0.0292,IMPROVEMENT,If frozen backbone helps: deep gradient flow into backbone causes collapse
+gradient_surgery,B,0.4875,0.05,IMPROVEMENT,If PCGrad helps: gradients are conflicting and resolvable by projection
+head_only,C,0.4958,0.0583,IMPROVEMENT,If head-only works: backbone representations are fine; only decision boundary needs updating
+kl_penalty,A,0.4417,0.0042,NEUTRAL,If this helps: catastrophic forgetting is the primary cause; soft regularisation suffices
+layer_ablation_top1,C,0.4875,0.05,IMPROVEMENT,Minimal unfrozen depth needed; collapse depth correlates with gradient flow depth
+layer_ablation_top2,C,0.4042,-0.0333,NEUTRAL,Minimal unfrozen depth needed; collapse depth correlates with gradient flow depth
+layer_ablation_top3,C,0.3917,-0.0458,NEUTRAL,Minimal unfrozen depth needed; collapse depth correlates with gradient flow depth
+llrd,A,0.4208,-0.0167,NEUTRAL,If LLRD helps: deep gradient flow into early layers corrupts representations
+lora,A,0.4333,-0.0042,NEUTRAL,If LoRA works: too many unconstrained degrees of freedom cause collapse
+low_t,B,0.4208,-0.0167,NEUTRAL,If low-t helps: high-t (coarse-structure) gradients are biased
+mixed_replay,A,0.4125,-0.025,NEUTRAL,If mixed replay helps: online data distribution alone is too corrupted
+normalized_adv,B,0.1208,-0.3167,COLLAPSE,If std normalisation helps: simple mean normalisation is too loose
+reward_filtering,D,0.4417,0.0042,NEUTRAL,If filtering helps: noisy/low-return data poisons gradients
+reward_model,D,0.3208,-0.1167,COLLAPSE,If reward model helps: raw returns are too sparse; learned model smooths signal
+running_stats,D,0.3833,-0.0542,COLLAPSE,If running stats help: batch normalisation is too noisy for small batches
+t_curriculum,B,0.4292,-0.0083,NEUTRAL,If curriculum helps: ordering of learning signals matters
+trust_region_kl,A,0.425,-0.0125,NEUTRAL,If hard constraint helps: soft KL is insufficient -- a hard boundary is needed

ablation_assets/main_results.csv

@@ -1,26 +1,26 @@
 Method,Group,Score,Score_Std,Delta_Pretrained,Delta_Baseline,Verdict
-gradient_surgery,B,0.6917,0.0118,-0.0083,0.1083,IMPROVEMENT
-lora,A,0.6833,0.0236,-0.0167,0.1,IMPROVEMENT
-layer_ablation_top1,C,0.6667,0.0118,-0.0333,0.0833,IMPROVEMENT
-layer_ablation_top2,C,0.6667,0.0118,-0.0333,0.0833,IMPROVEMENT
-frozen_backbone,C,0.6583,0.0236,-0.0417,0.075,IMPROVEMENT
-head_only,C,0.6583,0.0236,-0.0417,0.075,IMPROVEMENT
-ffn_only,C,0.6583,0.0118,-0.0417,0.075,IMPROVEMENT
-layer_ablation_top3,C,0.6583,0.0118,-0.0417,0.075,IMPROVEMENT
-trust_region_kl,A,0.65,0.0204,-0.05,0.0667,IMPROVEMENT
-attention_only,C,0.6333,0.0425,-0.0667,0.05,NEUTRAL
-reward_filtering,D,0.625,0.0204,-0.075,0.0417,NEUTRAL
-ewc,A,0.6167,0.0236,-0.0833,0.0333,NEUTRAL
-llrd,A,0.6167,0.0118,-0.0833,0.0333,NEUTRAL
-entropy_bonus,B,0.6167,0.0118,-0.0833,0.0333,NEUTRAL
-action_diversity,D,0.6167,0.0425,-0.0833,0.0333,NEUTRAL
-low_t,B,0.6,0.0354,-0.1,0.0167,NEUTRAL
-baseline_rl,Baseline,0.5833,0.0589,-0.1167,0.0,NEUTRAL
-kl_penalty,A,0.5833,0.0236,-0.1167,0.0,NEUTRAL
-t_curriculum,B,0.5833,0.0312,-0.1167,0.0,NEUTRAL
-mixed_replay,A,0.575,0.0204,-0.125,-0.0083,NEUTRAL
-running_stats,D,0.575,0.0,-0.125,-0.0083,NEUTRAL
-advantage_clip,B,0.5583,0.0118,-0.1417,-0.025,NEUTRAL
-reward_model,D,0.55,0.0408,-0.15,-0.0333,NEUTRAL
-bc_wins,B,0.4917,0.0825,-0.2083,-0.0917,NEUTRAL
-normalized_adv,B,0.2583,0.0656,-0.4417,-0.325,COLLAPSE
+head_only,C,0.4958,0.0386,0.0208,0.0583,IMPROVEMENT
+gradient_surgery,B,0.4875,0.0354,0.0125,0.05,IMPROVEMENT
+layer_ablation_top1,C,0.4875,0.0204,0.0125,0.05,IMPROVEMENT
+frozen_backbone,C,0.4667,0.0059,-0.0083,0.0292,IMPROVEMENT
+bc_wins,B,0.4583,0.0358,-0.0167,0.0208,NEUTRAL
+entropy_bonus,B,0.4458,0.0328,-0.0292,0.0083,NEUTRAL
+kl_penalty,A,0.4417,0.0212,-0.0333,0.0042,NEUTRAL
+reward_filtering,D,0.4417,0.0503,-0.0333,0.0042,NEUTRAL
+baseline_rl,Baseline,0.4375,0.0612,-0.0375,0.0,NEUTRAL
+ffn_only,C,0.4375,0.0102,-0.0375,0.0,NEUTRAL
+ewc,A,0.4333,0.0118,-0.0417,-0.0042,NEUTRAL
+lora,A,0.4333,0.0156,-0.0417,-0.0042,NEUTRAL
+t_curriculum,B,0.4292,0.0257,-0.0458,-0.0083,NEUTRAL
+trust_region_kl,A,0.425,0.0102,-0.05,-0.0125,NEUTRAL
+llrd,A,0.4208,0.0118,-0.0542,-0.0167,NEUTRAL
+low_t,B,0.4208,0.0524,-0.0542,-0.0167,NEUTRAL
+mixed_replay,A,0.4125,0.0468,-0.0625,-0.025,NEUTRAL
+layer_ablation_top2,C,0.4042,0.0156,-0.0708,-0.0333,NEUTRAL
+attention_only,C,0.3958,0.0236,-0.0792,-0.0417,NEUTRAL
+layer_ablation_top3,C,0.3917,0.0257,-0.0833,-0.0458,NEUTRAL
+running_stats,D,0.3833,0.0312,-0.0917,-0.0542,COLLAPSE
+action_diversity,D,0.375,0.0408,-0.1,-0.0625,COLLAPSE
+advantage_clip,B,0.3375,0.051,-0.1375,-0.1,COLLAPSE
+reward_model,D,0.3208,0.0118,-0.1542,-0.1167,COLLAPSE
+normalized_adv,B,0.1208,0.0412,-0.3542,-0.3167,COLLAPSE

ablation_assets/per_env.csv

@@ -1,26 +1,26 @@
 Method,MiniHack-Corridor-R2-v0,MiniHack-MazeWalk-9x9-v0,MiniHack-Room-Random-15x15-v0,MiniHack-Room-Random-5x5-v0
-action_diversity,0.5333,0.2333,0.6333,1.0
-advantage_clip,0.5333,0.2,0.4,1.0
-attention_only,0.6333,0.2,0.6667,1.0
-baseline_rl,0.5667,0.1333,0.5667,1.0
-bc_wins,0.5333,0.1,0.5667,1.0
-entropy_bonus,0.6,0.2,0.6333,1.0
-ewc,0.5333,0.2,0.7333,1.0
-ffn_only,0.7,0.2333,0.7,1.0
-frozen_backbone,0.5667,0.3,0.8,1.0
-gradient_surgery,0.6333,0.2667,0.8667,1.0
-head_only,0.5667,0.3,0.8,1.0
-kl_penalty,0.5667,0.2333,0.5667,1.0
-layer_ablation_top1,0.6667,0.2333,0.8,1.0
-layer_ablation_top2,0.6,0.2333,0.7333,1.0
-layer_ablation_top3,0.5667,0.2,0.8667,1.0
-llrd,0.5667,0.1667,0.6667,1.0
-lora,0.6,0.2,0.8,1.0
-low_t,0.6,0.1333,0.5333,0.9667
-mixed_replay,0.6,0.2,0.5667,1.0
-normalized_adv,0.1667,0.1667,0.3667,0.4333
-reward_filtering,0.5,0.2,0.7,1.0
-reward_model,0.5,0.1667,0.4333,0.9
-running_stats,0.6,0.1333,0.5333,1.0
-t_curriculum,0.6333,0.1667,0.5667,0.9667
-trust_region_kl,0.6333,0.2,0.7667,1.0
+action_diversity,0.4333,0.1,0.2667,0.7667
+advantage_clip,0.4333,0.15,0.1,0.6833
+attention_only,0.3667,0.1667,0.25,0.85
+baseline_rl,0.45,0.1333,0.3333,0.8333
+bc_wins,0.4333,0.1833,0.2833,0.8167
+entropy_bonus,0.45,0.1,0.3333,0.8333
+ewc,0.45,0.15,0.3333,0.8333
+ffn_only,0.4833,0.1667,0.4,0.7833
+frozen_backbone,0.4167,0.15,0.4833,0.9167
+gradient_surgery,0.45,0.15,0.4333,0.85
+head_only,0.45,0.15,0.4667,0.9
+kl_penalty,0.4167,0.15,0.25,0.8167
+layer_ablation_top1,0.4833,0.1833,0.45,0.8667
+layer_ablation_top2,0.4,0.15,0.3,0.7833
+layer_ablation_top3,0.4333,0.1667,0.25,0.7667
+llrd,0.4333,0.1333,0.2667,0.8667
+lora,0.4833,0.15,0.4167,0.85
+low_t,0.4667,0.1167,0.3333,0.85
+mixed_replay,0.4167,0.1167,0.3167,0.8
+normalized_adv,0.0833,0.0333,0.1,0.2833
+reward_filtering,0.4,0.1667,0.4333,0.8167
+reward_model,0.3667,0.1,0.1167,0.6833
+running_stats,0.4333,0.15,0.2833,0.7333
+t_curriculum,0.4167,0.1167,0.2667,0.7667
+trust_region_kl,0.4333,0.15,0.3167,0.9

checkpoints/offline/Minihack-Offline-Diffusion-BC-100M/config.yaml

@@ -19,9 +19,9 @@ baselines_eval_episodes_per_env: null
 baselines_eval_freq_env_steps: 10000
 baselines_n_envs_per_id: 2
 baselines_output_dir: outputs/baselines
-baselines_wandb_project: minihack-ReMDM-baselines
+baselines_wandb_project: remdm-planner-minihack-baselines
 buffer_capacity: 10000
-checkpoint_dir: offline_s0/offline_20260812_112835_ba36
+checkpoint_dir: checkpoints/offline_20260822_022901_0edc
 checkpoint_eval_episodes: 50
 checkpoint_every_timesteps: 940000
 collect_episodes_per_env: 5000
@@ -88,5 +88,4 @@ torch_compile: true
 total_timesteps: 5650000
 use_amp: true
 use_global_stream: true
-use_wandb: true
-weight_decay: 0.0001
+weight_decay: 0.0

checkpoints/offline/Minihack-Offline-Diffusion-BC-100M/model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:585f84cfb255804bf15eac2fa9e02c7cda626420c8af0b5ca38a5d93e1812e42
+oid sha256:a288ee5349613239757bfcc6e1847af89a9b5a28ff0a2843862476bdf3c9d3ed
 size 20974876

checkpoints/offline/Minihack-Offline-Diffusion-BC-100M/selection.json

@@ -3,7 +3,7 @@
   "selected": {
     "gradient_step": 50000
   },
-  "selection_metric": null,
+  "selection_metric": "mean ID win rate over 50 episodes/env, EMA weights",
   "candidates": {
     "unit": "gradient_steps",
     "every": 10000,

checkpoints/online/Minihack-Online-Diffusion-DAgger-100M/config.yaml

@@ -19,14 +19,14 @@ baselines_eval_episodes_per_env: null
 baselines_eval_freq_env_steps: 10000
 baselines_n_envs_per_id: 2
 baselines_output_dir: outputs/baselines
-baselines_wandb_project: minihack-ReMDM-baselines
+baselines_wandb_project: remdm-planner-minihack-baselines
 buffer_capacity: 10000
-checkpoint_dir: online_s0/dagger_20260811_210823_ee39
+checkpoint_dir: checkpoints/dagger_20260819_195240_ad84
 checkpoint_eval_episodes: 50
 checkpoint_every_timesteps: 940000
 collect_episodes_per_env: 5000
-collect_num_workers: 32
-collect_output: data/oracle_bc_gpu_h200.pt
+collect_num_workers: 8
+collect_output: data/oracle_bc_gpu_24gb.pt
 crop_size: 9
 curriculum_preseed: true
 curriculum_queue_size: 100
@@ -60,7 +60,7 @@ n_global_tokens: 8
 n_head: 4
 n_layer: 4
 noise_schedule: linear
-num_collection_workers: 32
+num_collection_workers: 8
 num_diffusion_steps: 100
 offline_batch_size: 2048
 offline_buffer_capacity: 1500000
@@ -88,5 +88,4 @@ torch_compile: true
 total_timesteps: 5650000
 use_amp: true
 use_global_stream: true
-use_wandb: true
 weight_decay: 0.0001

checkpoints/online/Minihack-Online-Diffusion-DAgger-100M/model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1a25616b8ebd7673982fea78ba33030a552cebd0356cbe9a6ff55c42380a3709
+oid sha256:045148d40b8bbccf0cbad981209adbae1da1caf3e11f6b57ae00490501fe7b4d
 size 20974876

checkpoints/online/Minihack-Online-Diffusion-DAgger-100M/selection.json

@@ -3,11 +3,11 @@
   "selected": {
     "dagger_iteration": 563
   },
-  "selection_metric": null,
+  "selection_metric": "mean ID win rate over 50 episodes/env, EMA weights",
   "candidates": {
-    "unit": "dagger_iterations",
-    "every": null,
-    "configured_max": null
+    "unit": "env_steps",
+    "every": 940000,
+    "configured_max": 5650000
   },
   "eval_protocol": {
     "episodes_per_env": 50,

demo_minihack.ipynb

@@ -0,0 +1,989 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# ReMDM MiniHack — Demo Notebook\n",
+    "\n",
+    "\n",
+    "\n",
+    "This notebook is a fully self-contained demonstration of **Remasking Discrete Diffusion Models (ReMDM) for action-sequence planning in MiniHack navigation environments**. It accompanies the paper *Return-Weighted ELBO Fine-Tuning Degrades Masked Diffusion Planners*. We study what the standard tractable objective for aligning a masked discrete diffusion planner with a reward actually does: across 25 ablation conditions in MiniHack and Craftax Classic, three seeds each, it **makes the planner worse** — no condition on Craftax Classic recovers the checkpoint it started from — and an exact decomposition of its gradient shows that the return weighting is **not** what does the damage.\n",
+    "\n",
+    "Running this notebook top-to-bottom in a fresh Google Colab runtime downloads the source code, the pre-trained DAgger checkpoint, and the pre-computed RL fine-tuning ablation assets from a public HuggingFace repository, then evaluates the model live on procedurally generated MiniHack layouts and visualises both the agent's behaviour and the underlying diffusion denoising process.\n",
+    "\n",
+    "**How to use this notebook**\n",
+    "\n",
+    "1. Open this notebook in Colab (`File > Upload notebook` > select `demo_minihack.ipynb`).\n",
+    "2. *(Optional)* Tweak the configuration constants in **Cell 1** below — change `ID_ENVS` / `OOD_ENVS` to test other registered environments, change `SEED` to test fresh procedurally-generated layouts, raise `EPISODES_PER_ENV` for tighter error bars, or set `CUSTOM_DES_FILE` to a hand-authored MiniHack `.des` level uploaded to Colab.\n",
+    "3. *(Optional)* `Runtime > Change runtime type > T4 GPU` for the fastest run.\n",
+    "4. `Runtime > Run all`.\n",
+    "\n",
+    "**Runtime budget** (default `EPISODES_PER_ENV=20`): ≈10 min on a Colab T4 GPU, ≈25 min on CPU. The first cell installs NLE which compiles C code — this takes ~3 min by itself.\n",
+    "\n",
+    "**Which evaluation path this notebook uses.** Cell 5 calls `Evaluator._run_episodes_batched` from `src/planners/inference.py` — the *same* evaluator the ablation harness calls (`experiments/rl_finetuning/ablations/training.py` imports `Evaluator` and calls `evaluator.evaluate(...)` at every eval point). The notebook and the paper's MiniHack numbers therefore come from one code path, with one sampler and one replanning rule. The only difference is episode count and seed offset.\n",
+    "\n",
+    "**Reproducibility statement.** Training is done **offline** — this notebook only runs inference on the supplied checkpoint, so that the notebook stays runnable on a free Colab runtime and can be tested on unseen inputs. Because MiniHack environments are procedurally generated, every run with a different `SEED` exposes the model to layouts it has never seen, including in the four in-distribution maps it was trained on."
+   ],
+   "id": "d866ea1602067016"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 1 — Configuration\n",
+    "\n",
+    "Everything you may want to change lives in this single cell."
+   ],
+   "id": "bed5b72429321b9d"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ============================================================================\n",
+    "# Configuration — change these to test on unseen inputs\n",
+    "# ============================================================================\n",
+    "\n",
+    "# Public HuggingFace repo containing source code, the stripped pre-trained\n",
+    "# checkpoint, and the pre-computed ablation assets. No auth required.\n",
+    "# This is an anonymous mirror published for double-blind review, so it carries\n",
+    "# no account name or URL that would identify the authors. Point it at your own\n",
+    "# Hub repo holding the same layout, or train from source (see README.md).\n",
+    "HF_REPO_ID: str = \"AnonMLuser/remdm-planner-minihack\"\n",
+    "\n",
+    "# Reproducibility seed. The live evaluator generates env seeds as\n",
+    "# (SEED + episode_index) per environment, so changing this number tests the\n",
+    "# model on a fresh batch of procedurally generated layouts.\n",
+    "SEED: int = 42\n",
+    "\n",
+    "# Episodes per environment for the live evaluation pass. Higher = tighter\n",
+    "# error bars but slower. The reported numbers in the paper used 50.\n",
+    "EPISODES_PER_ENV: int = 20\n",
+    "\n",
+    "# In-distribution training environments (4 maps).\n",
+    "ID_ENVS: list[str] = [\n",
+    "    \"MiniHack-Room-Random-5x5-v0\",\n",
+    "    \"MiniHack-Room-Random-15x15-v0\",\n",
+    "    \"MiniHack-Corridor-R2-v0\",\n",
+    "    \"MiniHack-MazeWalk-9x9-v0\",\n",
+    "]\n",
+    "\n",
+    "# Out-of-distribution zero-shot evaluation environments (3 maps).\n",
+    "OOD_ENVS: list[str] = [\n",
+    "    \"MiniHack-Room-Dark-15x15-v0\",\n",
+    "    \"MiniHack-Corridor-R5-v0\",\n",
+    "    \"MiniHack-MazeWalk-45x19-v0\",\n",
+    "]\n",
+    "\n",
+    "# Optional path to a custom .des MiniHack scenario file uploaded to Colab.\n",
+    "# Leave as None to skip; set to e.g. \"/content/my_level.des\" to evaluate the\n",
+    "# model on a hand-authored MiniHack level alongside the registry envs.\n",
+    "CUSTOM_DES_FILE: str | None = None\n",
+    "\n",
+    "# Inference device. None = auto-detect (CUDA if available, else CPU).\n",
+    "INFERENCE_DEVICE: str | None = None\n",
+    "\n",
+    "# Local directory the HF snapshot is downloaded into.\n",
+    "SNAPSHOT_DIR: str = \"remdm-planner-minihack\""
+   ],
+   "id": "3b72b898aab60cdf"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 2 — Setup & installation\n",
+    "\n",
+    "The next three cells install NetHack Learning Environment (NLE), MiniHack, PyTorch, and the supporting libraries; verify that NLE compiled correctly; and download the project source + checkpoint from HuggingFace. **NLE is the highest-risk failure point on Colab** because it has to compile C code under the hood. If the verification cell fails with an import error, restart the Colab runtime (`Runtime > Restart runtime`) and re-run from the install cell."
+   ],
+   "id": "34dd96e5ef1cb0ef"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ── 1. System dependencies + Python packages ─────────────────────────────\n",
+    "import os\n",
+    "import subprocess\n",
+    "import sys\n",
+    "\n",
+    "print(\"[1/3] Installing system dependencies for NLE (NetHack)...\")\n",
+    "try:\n",
+    "    _apt = subprocess.run(\n",
+    "        [\n",
+    "            \"apt-get\", \"install\", \"-y\", \"-q\",\n",
+    "            \"cmake\", \"build-essential\", \"bison\", \"flex\", \"libbz2-dev\",\n",
+    "        ],\n",
+    "        check=False, capture_output=True, text=True,\n",
+    "    )\n",
+    "    if _apt.returncode != 0:\n",
+    "        # Non-Colab environments (e.g. local Jupyter) will fail apt-get; that is\n",
+    "        # fine if the deps are already installed. Print and continue.\n",
+    "        print(\"  apt-get returned non-zero (likely already installed or non-Debian env):\")\n",
+    "        print(\"  \" + (_apt.stderr or \"\").strip()[-400:])\n",
+    "except FileNotFoundError:\n",
+    "    # No apt-get at all (macOS, non-Debian Linux). Fine if the build deps are\n",
+    "    # already present; NLE's compilation is verified in the next cell either way.\n",
+    "    print(\"  apt-get not found (non-Debian environment) - skipping system deps.\")\n",
+    "\n",
+    "print(\"[2/3] Installing NLE (compiles NetHack from source — slow first run)...\")\n",
+    "subprocess.run(\n",
+    "    [sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"nle>=1.2.0\"],\n",
+    "    check=True,\n",
+    ")\n",
+    "\n",
+    "print(\"[3/3] Installing MiniHack + supporting libraries...\")\n",
+    "subprocess.run(\n",
+    "    [\n",
+    "        sys.executable, \"-m\", \"pip\", \"install\", \"-q\",\n",
+    "        \"minihack>=1.0.2\",\n",
+    "        \"torch>=2.4\",\n",
+    "        \"huggingface_hub>=0.25\",\n",
+    "        \"polars>=1.0\",\n",
+    "        \"matplotlib>=3.8\",\n",
+    "        \"pyyaml>=6.0\",\n",
+    "        \"gymnasium>=0.29\",\n",
+    "        \"numpy>=1.26\",\n",
+    "    ],\n",
+    "    check=True,\n",
+    ")\n",
+    "print(\"Install step complete.\")"
+   ],
+   "id": "264dea581a6422cd"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ── 2. Verify NLE/MiniHack actually loaded ───────────────────────────────\n",
+    "# Fail loudly with a clear message if the C compilation step did not work.\n",
+    "try:\n",
+    "    import nle  # noqa: F401\n",
+    "    import minihack  # noqa: F401\n",
+    "    import gymnasium as gym\n",
+    "    _env = gym.make(\n",
+    "        \"MiniHack-Room-Random-5x5-v0\",\n",
+    "        observation_keys=(\"glyphs\", \"chars\"),\n",
+    "    )\n",
+    "    _obs, _info = _env.reset(seed=0)\n",
+    "    _env.close()\n",
+    "    print(f\"NLE OK. Glyphs shape: {_obs['glyphs'].shape}\")\n",
+    "except Exception as exc:\n",
+    "    raise RuntimeError(\n",
+    "        \"NLE/MiniHack failed to import or instantiate. On Colab this almost \"\n",
+    "        \"always means the C build of NetHack did not finish. Please restart \"\n",
+    "        \"the runtime (Runtime > Restart runtime) and re-run the install \"\n",
+    "        \"cell from a clean state.\"\n",
+    "    ) from exc"
+   ],
+   "id": "2a3dde194076e9be"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ── 3. Download source + checkpoint + ablation assets from HuggingFace ──\n",
+    "from huggingface_hub import snapshot_download\n",
+    "\n",
+    "if HF_REPO_ID == \"UNSET_HF_REPO_ID\":\n",
+    "    raise RuntimeError(\n",
+    "        \"HF_REPO_ID is unset: checkpoints are withheld for anonymous \"\n",
+    "        \"review. Edit the constant in Cell 1 to point at a \"\n",
+    "        \"public HuggingFace repo containing the source + checkpoint + assets.\"\n",
+    "    )\n",
+    "\n",
+    "snapshot_path = snapshot_download(repo_id=HF_REPO_ID, local_dir=SNAPSHOT_DIR)\n",
+    "print(f\"Snapshot at: {snapshot_path}\")\n",
+    "\n",
+    "# Make `import src.*` resolve against the downloaded snapshot.\n",
+    "if snapshot_path not in sys.path:\n",
+    "    sys.path.insert(0, snapshot_path)\n",
+    "\n",
+    "# Smoke imports — fail here if the snapshot is incomplete.\n",
+    "from src.config import load_config\n",
+    "from src.models.denoiser import make_model, ModelEMA\n",
+    "from src.diffusion.sampling import remdm_sample\n",
+    "from src.envs.minihack_env import make_env\n",
+    "from src.planners.inference import Evaluator, format_eval_results\n",
+    "print(\"Imports OK.\")"
+   ],
+   "id": "d90516ceca2cc1e3"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 3 — Project overview\n",
+    "\n",
+    "**Problem.** MiniHack navigation environments are sparse-reward gridworlds with procedurally generated dungeon layouts, locked doors, mazes, and partial observability. Off-the-shelf model-free reinforcement learning baselines (PPO, A2C, DQN, recurrent PPO) do not learn these layouts from scratch within our compute budget, reaching between 0.3% and 4.7% in-distribution after 5M environment steps at default hyperparameters — which is why we train by imitation. They are reported to establish that fact, not as a matched comparison.\n",
+    "\n",
+    "**Approach.** We treat planning as **discrete masked diffusion over action sequences**. A dual-stream transformer denoiser, conditioned on (i) a 9×9 local glyph crop centred on the agent and (ii) the full 21×79 global dungeon map, generates 64-step action plans by iteratively denoising a sequence of `[MASK]` tokens. At inference time we use **ReMDM** (Remasking Discrete Diffusion Models — Wang et al. 2025): MaskGIT-style progressive unmasking interleaved with stochastic confidence-weighted **remasking**, which lets the model revise low-confidence commitments mid-trajectory rather than baking in early mistakes.\n",
+    "\n",
+    "**Architecture (`LocalDiffusionPlannerWithGlobal`, ≈5.2M parameters, PyTorch).**\n",
+    "\n",
+    "```\n",
+    "Local stream:    9×9 glyphs   → Embed(6000,64) → CNN(64→32→64) → Linear → 1 token  (256-D)\n",
+    "Global stream:   21×79 glyphs → Embed(6000,32) → CNN(32→32→64) → AdaptivePool(2,4)\n",
+    "                              → Linear(64,256) → 8 spatial tokens (256-D)\n",
+    "                              + auxiliary goal head: mean(global) → MLP → [B,2]  (normalised staircase coords)\n",
+    "                              × sigmoid(learnable scalar gate, init logit = −3.0)  ← keeps global stream nearly closed early in training\n",
+    "Action stream:   action_emb(14,256) + timestep_emb(100,256) + position_emb(64,256)\n",
+    "Transformer:     concat [local(1) + global(8) + actions(64) = 73 tokens]\n",
+    "                 → 4-layer bidirectional TransformerEncoder (256-D, 4 heads, GELU, pre-norm) → last 64 tokens\n",
+    "Action head:     Linear(256, 12) → action logits  (12 movement / interact actions; MASK=12 and PAD=13 are inputs only)\n",
+    "```\n",
+    "\n",
+    "The auxiliary staircase-coordinate head **has no counterpart in the offline baselines**, and the paper lists it as one of four confounds in the imitation comparison (Cell 8).\n",
+    "\n",
+    "**Diffusion process.** Forward: each token is independently replaced with `[MASK]` with probability σ(t) under a linear schedule. Reverse (inference): K=10 denoising steps with confidence-based progressive unmasking interleaved with stochastic remasking at probability p_remask = η · (1 − r), where η = 0.15 and r is the progress through the chain. DAgger collection uses greedy argmax sampling at K=5 steps. The planner runs in **receding horizon** mode: it commits a 64-action plan, executes the first 16 actions, then replans on the next observation.\n",
+    "\n",
+    "**Training pipeline.** DAgger online training with **BFS oracle** supervision (`main.py --mode online`). Iteration 0 seeds the replay buffer with three oracle trajectories per training environment as an implicit BC warm-start; subsequent iterations alternate model rollouts → BFS oracle relabelling on the same seed → efficiency-filtered buffer insertion → AdamW gradient steps. The exponentially decaying β starts at 1.0; total budget ≈5.65M env steps with 30 episodes/iteration and 100 grad steps/iteration. Inference uses EMA-shadowed weights. The checkpoint loaded by this notebook is the iter-600 DAgger snapshot from a GPU-H200 run, with all optimiser/scheduler state stripped (the inference path only needs `ema_state_dict`).\n",
+    "\n",
+    "### The objective under study\n",
+    "\n",
+    "Fine-tuning runs on the planner's **own rollouts**. Each iteration collects\n",
+    "trajectories under the current policy and cuts them into `H`-step windows. Window\n",
+    "`i` is assigned the return `R_i` — the reward summed over exactly the `H` steps that\n",
+    "window trains on, *not* the episode total broadcast to every window — and receives a\n",
+    "scalar weight\n",
+    "\n",
+    "```\n",
+    "A_i = clip( max(R_i, 0) / (µ_batch + ε),  c_min, c_max ),   [c_min, c_max] = [0.1, 5.0]\n",
+    "```\n",
+    "\n",
+    "The training loss is the ELBO of the denoising objective with each window's term\n",
+    "scaled by `A_i`, treated as a constant under a stop-gradient. This is a **clipped\n",
+    "return ratio**, in the lineage of reward-weighted regression — it is *not* the\n",
+    "advantage-weighted regression weight, which subtracts a baseline and applies an\n",
+    "exponential temperature `exp(A/β)`. The paper's conclusions characterise this weight\n",
+    "function and do not carry over automatically to the exponential form, which was not\n",
+    "tested.\n",
+    "\n",
+    "### The exact decomposition\n",
+    "\n",
+    "Write `Ā` for the mean weight and `δ_i = A_i/Ā − 1`, so that `Σ_i δ_i = 0`. The\n",
+    "gradient then factors **exactly**, with no approximation:\n",
+    "\n",
+    "```\n",
+    "∇L_RW  =  Ā · [  ∇L_BC   +   (1/B) Σ_i δ_i ∇ℓ_i  ]\n",
+    "                 ‾‾‾‾‾‾        ‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾\n",
+    "                imitation      g_δ — the return's entire contribution\n",
+    "```\n",
+    "\n",
+    "`Ā` is a scalar that rescales the step size and leaves the direction unchanged, so\n",
+    "everything the return does enters through `g_δ`, whose RMS scale is the weights'\n",
+    "coefficient of variation `CV_A = sqrt(B/ESS − 1)`. Both terms are measurable at a\n",
+    "single parameter point, and both diagnostics cost one batch and no accelerator.\n",
+    "\n",
+    "### Three contributions (paper headline)\n",
+    "\n",
+    "1. **Fine-tuning degrades the planner.** On Craftax Classic, across 25 conditions\n",
+    "   with three seeds each, **no condition recovers the checkpoint it started from**\n",
+    "   and the plain objective gives up **3.59 of 11.81** points. Collected return\n",
+    "   falls alongside held-out score (5.21 → 3.73 while eval goes 12.06 → 8.49), so\n",
+    "   this is not reward hacking under an evaluation mismatch. The conditions that\n",
+    "   lose least are the ones that barely train: LoRA (11.63), a hard KL trust region\n",
+    "   (11.53) and head-only updates (11.30).\n",
+    "2. **The return term is large and points away from imitation.** `CV_A` is **1.16**\n",
+    "   on Craftax Classic (effective sample size 437 of a 1024-window batch) and\n",
+    "   **2.68** on MiniHack (591 of 4608), so the weights do rank windows and the\n",
+    "   familiar explanation — that sparse returns fail to discriminate — is\n",
+    "   unavailable. Measured directly, `‖g_δ‖/‖∇L_BC‖ = 0.49 ± 0.01` at cosine\n",
+    "   **0.02 ± 0.05**, against **0.893 ± 0.010** between two independent noise draws\n",
+    "   of `∇L_BC` itself. In `D = 9.33 × 10⁶` dimensions a random direction gives\n",
+    "   `3.3 × 10⁻⁴`, so `g_δ` is structure, not noise.\n",
+    "3. **But shrinking it makes things worse.** Advantage clipping confines the weights\n",
+    "   to [0.8, 1.2] and cuts the return term fivefold, to **0.097 ± 0.003**. It scores\n",
+    "   **5.06** on Craftax Classic — 3.16 below the unclipped baseline and second worst\n",
+    "   in the suite. A condition with almost no return term degrades *further* than one\n",
+    "   with a large one, so the return weighting is not what does the damage. What is\n",
+    "   left is the data: advantage clipping is also the nearest thing in the suite to\n",
+    "   unweighted training on the model's own rollouts, and the suite is ordered by how\n",
+    "   much plasticity each condition allows rather than by any property of the reward.\n",
+    "\n",
+    "*(Live evaluation in Cell 5; imitation context in Cell 8; ablation evidence in Cells 9–10.)*"
+   ],
+   "id": "14890835349143"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 4 — Load the pre-trained model\n",
+    "\n",
+    "We load the architecture from `src/models/denoiser.py` using the project's own `make_model(cfg)` factory, then pour the EMA shadow weights from the stripped checkpoint into it. This is the same EMA-evaluation path that `main.py --mode inference` uses internally."
+   ],
+   "id": "2516d108471c4cc6"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "\n",
+    "# Load the project's authoritative config (the same defaults.yaml the\n",
+    "# training run used). All sampling / arch hyperparameters live there.\n",
+    "_overrides = {\"device\": INFERENCE_DEVICE} if INFERENCE_DEVICE else {}\n",
+    "cfg = load_config(cli_overrides=_overrides)\n",
+    "device = torch.device(cfg.device)\n",
+    "print(f\"Inference device: {device}\")\n",
+    "print(\n",
+    "    f\"Model arch: n_embd={cfg.n_embd}, n_head={cfg.n_head}, \"\n",
+    "    f\"n_layer={cfg.n_layer}, n_global_tokens={cfg.n_global_tokens}, \"\n",
+    "    f\"seq_len={cfg.seq_len}\"\n",
+    ")\n",
+    "print(\n",
+    "    f\"Diffusion: K_eval={cfg.diffusion_steps_eval}, schedule={cfg.noise_schedule}, \"\n",
+    "    f\"remask={cfg.remask_strategy}, eta={cfg.eta}, T={cfg.temperature}, top_p={cfg.top_p}\"\n",
+    ")\n",
+    "\n",
+    "# Build the architecture and load the stripped EMA checkpoint. The HF repo\n",
+    "# ships an inference-only checkpoint that contains only the EMA shadow weights\n",
+    "# (≈21 MB) — no optimiser, scheduler, RNG, or curriculum state.\n",
+    "ckpt_path = os.path.join(SNAPSHOT_DIR, \"checkpoint_inference.pth\")\n",
+    "if not os.path.exists(ckpt_path):\n",
+    "    raise FileNotFoundError(\n",
+    "        f\"Expected stripped checkpoint at {ckpt_path}. The HF repo must \"\n",
+    "        f\"contain checkpoint_inference.pth at its root.\"\n",
+    "    )\n",
+    "\n",
+    "model = make_model(cfg).to(device)\n",
+    "ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)\n",
+    "ema_sd = ckpt[\"ema_state_dict\"] if isinstance(ckpt, dict) and \"ema_state_dict\" in ckpt else ckpt\n",
+    "ema = ModelEMA(model, decay=cfg.ema_decay)\n",
+    "ema.load_state_dict(ema_sd)\n",
+    "ema.apply_to(model)\n",
+    "model.eval()\n",
+    "\n",
+    "n_params = sum(p.numel() for p in model.parameters())\n",
+    "print(\n",
+    "    f\"Loaded EMA weights into {type(model).__name__} \"\n",
+    "    f\"({n_params:,} parameters ≈ {n_params / 1e6:.2f} M).\"\n",
+    ")"
+   ],
+   "id": "71470180a5f12e46"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 5 — Live inference on procedurally generated layouts ⭐\n",
+    "\n",
+    "We now run the model live on `EPISODES_PER_ENV` episodes per environment, on the 4 in-distribution maps **and** the 3 out-of-distribution maps **and** the optional custom `.des` file. Each episode uses a per-`SEED` RNG so changing `SEED` in Cell 1 produces a fresh batch of procedurally generated layouts that the model has never trained or evaluated on.\n",
+    "\n",
+    "We call the project's batched evaluator directly (`Evaluator._run_episodes_batched`), which runs all `EPISODES_PER_ENV` rollouts of a given environment in lockstep and batches every replanning forward pass through the GPU. Episodes that fail to construct (e.g. an invalid `.des` file) count as losses individually rather than crashing the pass."
+   ],
+   "id": "717e4b710eba9ccc"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import time\n",
+    "from pathlib import Path\n",
+    "\n",
+    "import numpy as np\n",
+    "import polars as pl\n",
+    "\n",
+    "evaluator = Evaluator()\n",
+    "\n",
+    "\n",
+    "def run_live_eval(\n",
+    "    env_ids: list[str],\n",
+    "    n_episodes: int,\n",
+    "    seed: int,\n",
+    "    des_files: list[str] | None = None,\n",
+    ") -> dict[str, dict]:\n",
+    "    \"\"\"Evaluate the model on each env id with seeds derived from *seed*.\n",
+    "\n",
+    "    Mirrors ``Evaluator.evaluate`` but uses ``seed + ep`` so that changing\n",
+    "    the notebook ``SEED`` constant actually rotates the underlying\n",
+    "    procedural layouts.\n",
+    "    \"\"\"\n",
+    "    targets: list[tuple[str, str | None]] = [(eid, None) for eid in env_ids]\n",
+    "    if des_files:\n",
+    "        for p in des_files:\n",
+    "            with open(p) as fh:\n",
+    "                targets.append((Path(p).stem, fh.read()))\n",
+    "\n",
+    "    out: dict[str, dict] = {}\n",
+    "    for env_id, des_content in targets:\n",
+    "        seeds = [seed + ep for ep in range(n_episodes)]\n",
+    "        eps = evaluator._run_episodes_batched(\n",
+    "            model, env_id, n_episodes, cfg, device,\n",
+    "            seeds=seeds,\n",
+    "            des_content=des_content,\n",
+    "            blind_global=False,\n",
+    "        )\n",
+    "        wins = sum(1 for r in eps if r[\"won\"])\n",
+    "        n = max(len(eps), 1)\n",
+    "        out[env_id] = {\n",
+    "            \"win_rate\": wins / n,\n",
+    "            \"wins\": wins,\n",
+    "            \"avg_reward\": sum(r[\"total_reward\"] for r in eps) / n,\n",
+    "            \"avg_steps\": sum(r[\"steps\"] for r in eps) / n,\n",
+    "            \"n_episodes\": len(eps),\n",
+    "        }\n",
+    "    return out\n",
+    "\n",
+    "\n",
+    "all_envs = list(ID_ENVS) + list(OOD_ENVS)\n",
+    "des_files = [CUSTOM_DES_FILE] if CUSTOM_DES_FILE else None\n",
+    "n_total = len(all_envs) * EPISODES_PER_ENV\n",
+    "if des_files:\n",
+    "    n_total += EPISODES_PER_ENV\n",
+    "\n",
+    "print(\n",
+    "    f\"Evaluating on {len(all_envs)} registry envs\"\n",
+    "    + (f\" + {len(des_files)} custom .des\" if des_files else \"\")\n",
+    "    + f\" × {EPISODES_PER_ENV} episodes (= {n_total} rollouts)...\"\n",
+    ")\n",
+    "t0 = time.time()\n",
+    "live_results = run_live_eval(\n",
+    "    all_envs, EPISODES_PER_ENV, SEED, des_files=des_files,\n",
+    ")\n",
+    "elapsed = time.time() - t0\n",
+    "print(f\"Done in {elapsed:.1f}s ({elapsed / max(n_total, 1):.2f}s per rollout).\\n\")\n",
+    "\n",
+    "print(format_eval_results(live_results, label=\"Live ReMDM Inference\"))"
+   ],
+   "id": "812ffba082dc4a3a"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Tabulate and split-aware aggregate.\n",
+    "rows = []\n",
+    "for env_id, stats in live_results.items():\n",
+    "    if env_id in ID_ENVS:\n",
+    "        split = \"ID\"\n",
+    "    elif env_id in OOD_ENVS:\n",
+    "        split = \"OOD\"\n",
+    "    else:\n",
+    "        split = \"custom\"\n",
+    "    rows.append(\n",
+    "        {\n",
+    "            \"split\": split,\n",
+    "            \"env\": env_id,\n",
+    "            \"win_rate\": round(stats[\"win_rate\"], 3),\n",
+    "            \"avg_steps\": round(stats[\"avg_steps\"], 1),\n",
+    "            \"avg_reward\": round(stats[\"avg_reward\"], 2),\n",
+    "            \"n_episodes\": stats[\"n_episodes\"],\n",
+    "        }\n",
+    "    )\n",
+    "df_live = pl.DataFrame(rows).sort([\"split\", \"env\"])\n",
+    "print(df_live)\n",
+    "\n",
+    "id_wr = float(np.mean([live_results[e][\"win_rate\"] for e in ID_ENVS if e in live_results]))\n",
+    "ood_wr = float(np.mean([live_results[e][\"win_rate\"] for e in OOD_ENVS if e in live_results]))\n",
+    "print(\n",
+    "    f\"\\nMean ID  win rate (this run): {id_wr:.2%}\"\n",
+    "    f\"   <- reference (paper Table 6, ReMDM DAgger, 50 episodes/env): 48.5%\"\n",
+    ")\n",
+    "print(\n",
+    "    f\"Mean OOD win rate (this run): {ood_wr:.2%}\"\n",
+    "    f\"   <- reference (paper Table 6, ReMDM DAgger): 4.7% mean OOD\"\n",
+    "    f\" (Room-Dark-15x15 12%, Corridor-R5 2%, MazeWalk-45x19 0%)\"\n",
+    ")\n",
+    "print(\n",
+    "    \"\\nNote: small deviations from the reported numbers are expected -- this run \"\n",
+    "    f\"uses {EPISODES_PER_ENV} episodes per env (vs 50) and a different seed offset, \"\n",
+    "    \"so it samples a different (smaller) batch of procedurally generated layouts.\"\n",
+    ")"
+   ],
+   "id": "631ddcc627be0ac5"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 6 — Visualise the agent's behaviour ⭐\n",
+    "\n",
+    "We roll out a single episode on a representative ID environment and capture the dual-stream observations the model actually sees: the **9×9 local glyph crop** (centred on the agent `@`) and the **full 21×79 global dungeon map**. We display the start, mid, and end of the episode side-by-side."
+   ],
+   "id": "57f89da930062058"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import random\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "\n",
+    "random.seed(SEED)\n",
+    "np.random.seed(SEED)\n",
+    "torch.manual_seed(SEED)\n",
+    "\n",
+    "VIZ_ENV = ID_ENVS[2]  # MiniHack-Corridor-R2-v0 — visually richer than a single room\n",
+    "\n",
+    "env = make_env(VIZ_ENV, None, cfg)\n",
+    "(local_obs, global_obs), _ = env.reset(seed=SEED)\n",
+    "\n",
+    "trajectory_local: list[np.ndarray] = [local_obs.copy()]\n",
+    "trajectory_global: list[np.ndarray] = [global_obs.copy()]\n",
+    "actions_taken: list[int] = []\n",
+    "plan = None\n",
+    "step_in_plan = 0\n",
+    "total_reward = 0.0\n",
+    "won = False\n",
+    "\n",
+    "for step in range(200):\n",
+    "    if step_in_plan == 0 or step_in_plan >= cfg.replan_every:\n",
+    "        # Replan: full ReMDM denoising over the current local + global obs.\n",
+    "        local_t = torch.from_numpy(local_obs).long().unsqueeze(0).to(device)  # [1, 9, 9]\n",
+    "        glb_t = torch.from_numpy(global_obs).long().unsqueeze(0).to(device)  # [1, 21, 79]\n",
+    "        plan = remdm_sample(\n",
+    "            model, local_t, glb_t, cfg, device,\n",
+    "            physics_aware=False, blind_global=False,\n",
+    "        )[0].cpu().numpy()  # [seq_len]\n",
+    "        step_in_plan = 0\n",
+    "\n",
+    "    action = int(plan[step_in_plan])\n",
+    "    step_in_plan += 1\n",
+    "    actions_taken.append(action)\n",
+    "    (local_obs, global_obs), reward, term, trunc, info = env.step(action)\n",
+    "    total_reward += reward\n",
+    "    trajectory_local.append(local_obs.copy())\n",
+    "    trajectory_global.append(global_obs.copy())\n",
+    "    if info.get(\"won\"):\n",
+    "        won = True\n",
+    "    if term or trunc:\n",
+    "        break\n",
+    "\n",
+    "env.close()\n",
+    "n_steps = len(actions_taken)\n",
+    "print(\n",
+    "    f\"{VIZ_ENV} (seed={SEED}): {n_steps} steps, won={won}, total_reward={total_reward:.2f}\"\n",
+    ")"
+   ],
+   "id": "ef5071a3e18492aa"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "snapshots = [\n",
+    "    (\"Start\", 0),\n",
+    "    (\"Mid\", n_steps // 2),\n",
+    "    (\"End\", n_steps),\n",
+    "]\n",
+    "\n",
+    "fig, axes = plt.subplots(2, 3, figsize=(13, 7))\n",
+    "for col, (label, idx) in enumerate(snapshots):\n",
+    "    axes[0, col].imshow(trajectory_local[idx], cmap=\"viridis\")\n",
+    "    axes[0, col].set_title(f\"{label} — local 9×9 (step {idx})\")\n",
+    "    axes[0, col].set_xticks([])\n",
+    "    axes[0, col].set_yticks([])\n",
+    "\n",
+    "    axes[1, col].imshow(trajectory_global[idx], cmap=\"viridis\")\n",
+    "    axes[1, col].set_title(f\"{label} — global 21×79 (step {idx})\")\n",
+    "    axes[1, col].set_xticks([])\n",
+    "    axes[1, col].set_yticks([])\n",
+    "\n",
+    "fig.suptitle(\n",
+    "    f\"Dual-stream observations during a live rollout on {VIZ_ENV}\",\n",
+    "    fontsize=12,\n",
+    ")\n",
+    "plt.tight_layout()\n",
+    "plt.show()\n",
+    "\n",
+    "print(\n",
+    "    \"Each colour above is a NetHack glyph ID (walls, floor, agent @, staircase >, \"\n",
+    "    \"doors +, etc.). The local crop is the 9×9 window the local CNN stream sees; \"\n",
+    "    \"the global stream sees the full 21×79 grid and contributes the auxiliary \"\n",
+    "    \"staircase-coordinate prediction through the goal head.\"\n",
+    ")"
+   ],
+   "id": "14d3f4ba57c10610"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 7 — Visualise the ReMDM denoising process ⭐\n",
+    "\n",
+    "This is the heart of the method. Given a single observation, the planner starts from a sequence of 64 `[MASK]` tokens and over `K = diffusion_steps_eval = 10` reverse-diffusion steps it (a) predicts an action distribution at every position, (b) commits the highest-confidence positions via MaskGIT-style progressive unmasking, and (c) **stochastically remasks** previously committed positions whose confidence is low so they can be re-decoded later. We surface the per-step state by calling `remdm_sample(..., return_analytics=True)`, which returns the full denoising trajectory and per-step diagnostics."
+   ],
+   "id": "936ac1409c9efa8a"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Use a fresh observation from the same env so the visualisation stays focused\n",
+    "# on the token-level denoising process, not on action quality.\n",
+    "env = make_env(ID_ENVS[2], None, cfg)\n",
+    "(local_obs, global_obs), _ = env.reset(seed=SEED)\n",
+    "env.close()\n",
+    "\n",
+    "local_t = torch.from_numpy(local_obs).long().unsqueeze(0).to(device)\n",
+    "glb_t = torch.from_numpy(global_obs).long().unsqueeze(0).to(device)\n",
+    "\n",
+    "seq, path_per_step, conf_track, masked_track = remdm_sample(\n",
+    "    model, local_t, glb_t, cfg, device,\n",
+    "    physics_aware=False, blind_global=False,\n",
+    "    return_analytics=True,\n",
+    ")\n",
+    "K = len(path_per_step)\n",
+    "seq_len = cfg.seq_len\n",
+    "mask_token = cfg.mask_token\n",
+    "\n",
+    "# Build a [K+1, seq_len] visualisation: row 0 = the all-MASK initial state,\n",
+    "# row k = the sequence after denoising step k.\n",
+    "vis = np.full((K + 1, seq_len), mask_token, dtype=np.int64)\n",
+    "for k, state in enumerate(path_per_step):\n",
+    "    vis[k + 1] = state\n",
+    "# Encode masked positions as -1 so they get a distinct colour band.\n",
+    "display_grid = np.where(vis == mask_token, -1, vis).astype(float)\n",
+    "\n",
+    "print(\n",
+    "    f\"Denoising trajectory: K={K} steps, seq_len={seq_len}, \"\n",
+    "    f\"final masked-token count={int((seq[0] == mask_token).sum().item())} (must be 0)\"\n",
+    ")"
+   ],
+   "id": "59079c3462bbfaf6"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, (ax1, ax2) = plt.subplots(\n",
+    "    2, 1, figsize=(13, 6),\n",
+    "    gridspec_kw={\"height_ratios\": [3, 1]},\n",
+    ")\n",
+    "\n",
+    "im = ax1.imshow(\n",
+    "    display_grid, aspect=\"auto\", cmap=\"viridis\",\n",
+    "    vmin=-1, vmax=cfg.action_dim - 1,\n",
+    "    interpolation=\"nearest\",\n",
+    ")\n",
+    "ax1.set_yticks(range(K + 1))\n",
+    "ax1.set_yticklabels([\"init\"] + [f\"k={k+1}\" for k in range(K)])\n",
+    "ax1.set_xlabel(\"token position (0..63)\")\n",
+    "ax1.set_ylabel(\"denoising step\")\n",
+    "ax1.set_title(\n",
+    "    f\"ReMDM denoising trajectory ({K} steps) — \"\n",
+    "    \"dark band (-1) = masked, colours = committed action ids 0..11\"\n",
+    ")\n",
+    "cbar = plt.colorbar(im, ax=ax1, ticks=[-1] + list(range(cfg.action_dim)))\n",
+    "cbar.set_label(\"token id (-1 = MASK)\")\n",
+    "\n",
+    "ax2.plot(range(1, K + 1), conf_track, marker=\"o\", label=\"avg confidence (committed tokens)\")\n",
+    "ax2.plot(\n",
+    "    range(1, K + 1),\n",
+    "    [m / seq_len for m in masked_track],\n",
+    "    marker=\"s\",\n",
+    "    label=\"fraction still masked\",\n",
+    ")\n",
+    "ax2.set_xlabel(\"denoising step k\")\n",
+    "ax2.set_ylabel(\"value\")\n",
+    "ax2.set_xticks(range(1, K + 1))\n",
+    "ax2.legend(loc=\"upper right\")\n",
+    "ax2.grid(alpha=0.3)\n",
+    "\n",
+    "plt.tight_layout()\n",
+    "plt.show()\n",
+    "\n",
+    "print(f\"Final committed plan (first 16 actions): {seq[0, :16].tolist()}\")\n",
+    "print(\n",
+    "    \"Read the heatmap top-to-bottom: row 0 is fully masked, the next rows commit \"\n",
+    "    \"high-confidence positions while occasionally re-masking low-confidence ones \"\n",
+    "    \"(the ReMDM 'conf' strategy with eta=0.15), and the last row is the fully \"\n",
+    "    \"committed plan that gets executed for the next 16 env steps before replanning.\"\n",
+    ")"
+   ],
+   "id": "ee28eb905ce0956"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 8 — Imitation learning context (baselines)\n",
+    "\n",
+    "The table below reports this project's MiniHack baseline runs, under the **imitation protocol**: 50 episodes per environment at inference settings. Baseline rows are means ± std over 3 training seeds; the two ReMDM rows are a single evaluation seed and are read directly from the shipped `results/inference/` artefacts.\n",
+    "\n",
+    "**Provenance note.** The current paper carries the baseline *configurations* (Table 11, \"MiniHack baseline hyperparameters\") but no baseline *results* table — that comparison was cut when the paper narrowed to the fine-tuning result. The model-free and offline-BC rows below are therefore project context rather than a paper number. The two ReMDM rows *are* paper numbers: the DAgger row is paper Table 6.\n",
+    "\n",
+    "The four model-free rows are included for **one purpose**: to establish that these layouts are not solvable from scratch within our compute budget, which is why we train by imitation. They use default hyperparameters and are **not a matched comparison**.\n",
+    "\n",
+    "**This is not an architecture result.** Any gap to the offline CNN+MLP and Decision Transformer confounds four design differences: an auxiliary staircase-prediction head the baselines lack, ten denoising passes per decision against one, chunked replanning against greedy per-step selection, and untuned baselines.\n",
+    "\n",
+    "**Two evaluation protocols appear in this work and their numbers are not interchangeable.** The direct evaluation of the DAgger checkpoint (this cell, 50 episodes per environment) gives **48.5%** in-distribution — this is paper Table 6. The ablation-harness protocol (Cells 9–10) runs 20 episodes per environment per seed inside the training loop at its own sampling settings, and gives the same checkpoint **47.5%**. The two agree closely here, but no claim crosses between them."
+   ],
+   "id": "e0023c1ce6dbbaee"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import polars as pl\n",
+    "\n",
+    "# Baseline runs under the imitation protocol: 50 episodes per environment at\n",
+    "# inference settings. Baseline rows are means +/- std over 3 training seeds.\n",
+    "#\n",
+    "# PROVENANCE: the current paper keeps the baseline *configurations* (Table 11)\n",
+    "# but no baseline *results* table, so the first six rows are project context,\n",
+    "# not paper numbers. The two ReMDM rows are read from the shipped\n",
+    "# results/inference/*.json artefacts; the DAgger row is paper Table 6.\n",
+    "baselines = [\n",
+    "    {\"method\": \"PPO\",                            \"params\": \"~0.4 M\", \"id_win_pct\": \"0.3 +/- 0.5\",  \"id_steps\":  \"396.8 +/- 10.1\", \"ood_win_pct\": \"0.0 +/- 0.0\", \"ood_steps\": \"767.6 +/- 6.0\"},\n",
+    "    {\"method\": \"A2C\",                            \"params\": \"~0.4 M\", \"id_win_pct\": \"2.8 +/- 2.3\",  \"id_steps\":  \"392.9 +/- 8.1\",  \"ood_win_pct\": \"0.4 +/- 0.8\", \"ood_steps\": \"763.8 +/- 5.0\"},\n",
+    "    {\"method\": \"DQN\",                            \"params\": \"~0.4 M\", \"id_win_pct\": \"3.8 +/- 1.2\",  \"id_steps\":  \"388.2 +/- 5.2\",  \"ood_win_pct\": \"1.3 +/- 0.2\", \"ood_steps\": \"762.7 +/- 1.3\"},\n",
+    "    {\"method\": \"PPO-RNN\",                        \"params\": \"~0.5 M\", \"id_win_pct\": \"4.7 +/- 2.9\",  \"id_steps\":  \"385.4 +/- 13.2\", \"ood_win_pct\": \"0.0 +/- 0.0\", \"ood_steps\": \"766.7 +/- 3.1\"},\n",
+    "    {\"method\": \"CNN+MLP (Offline BC)\",           \"params\": \"~0.6 M\", \"id_win_pct\": \"6.3 +/- 2.5\",  \"id_steps\":  \"378.8 +/- 6.4\",  \"ood_win_pct\": \"0.7 +/- 0.7\", \"ood_steps\": \"764.7 +/- 2.0\"},\n",
+    "    {\"method\": \"Decision Transformer (Off. BC)\", \"params\": \"~1.0 M\", \"id_win_pct\": \"9.3 +/- 2.1\",  \"id_steps\":  \"160.7 +/- 4.2\",  \"ood_win_pct\": \"1.8 +/- 0.4\", \"ood_steps\": \"196.8 +/- 0.6\"},\n",
+    "    {\"method\": \"ReMDM (Offline BC)\",             \"params\": \"~5.2 M\", \"id_win_pct\": \"73.0\",         \"id_steps\":  \"107.1\",          \"ood_win_pct\": \"6.0\",         \"ood_steps\": \"418.5\"},\n",
+    "    {\"method\": \"ReMDM (DAgger)\",                 \"params\": \"~5.2 M\", \"id_win_pct\": \"48.5\",         \"id_steps\":  \"173.9\",          \"ood_win_pct\": \"4.7\",         \"ood_steps\": \"426.2\"},\n",
+    "]\n",
+    "df_baselines = pl.DataFrame(baselines)\n",
+    "with pl.Config(tbl_rows=10, fmt_str_lengths=40):\n",
+    "    print(df_baselines)\n",
+    "print(\n",
+    "    \"\\nThe four model-free rows are NOT a matched comparison -- they use default \"\n",
+    "    \"hyperparameters and are reported only to establish that these layouts are not \"\n",
+    "    \"solvable from scratch within our budget (0.3-4.7% ID after 5M env steps), \"\n",
+    "    \"which is why we train by imitation.\\n\"\n",
+    "    \"\\nThe DAgger row (48.5% ID / 4.7% OOD) is paper Table 6, and it is the \"\n",
+    "    \"checkpoint every ablation condition in Cells 9-10 fine-tunes. The planner \"\n",
+    "    \"retains little zero-shot transfer to held-out layouts, so the paper makes no \"\n",
+    "    \"structural generalisation claim.\"\n",
+    ")"
+   ],
+   "id": "cbd0984cc1f7787"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 9 — RL fine-tuning ablation findings (pre-computed)\n",
+    "\n",
+    "On top of the DAgger checkpoint, we ran a **25-condition suite** (`experiments/rl_finetuning/run_ablations.py --all`) that asks what return-weighted ELBO fine-tuning of the diffusion planner actually does. The 25 conditions are organised into four mechanism groups (A: regularisation, B: training signal, C: architectural freezing, D: data quality / weight computation) plus a `baseline_rl` reference, all initialised from the same checkpoint and run for 500 iterations with 3 seeds. All numbers here use the **ablation-harness protocol**, under which the pretrained checkpoint scores **47.5%** — against 48.5% for the direct evaluation in Cell 8.\n",
+    "\n",
+    "**In MiniHack the objective takes the checkpoint from 47.5% to 43.8% (±6.1).** Three conditions finish nominally above the checkpoint — `head_only` 49.6%, `gradient_surgery` 48.8%, `layer_ablation_top1` 48.8% — 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.\n",
+    "\n",
+    "**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 (11.81 → 8.22, no condition above the checkpoint). The claim rests on Craftax Classic.\n",
+    "\n",
+    "**The reward pathway is not the explanation.** `CV_A` averages **2.68** here (effective sample size 591 of a 4,608-window batch) against 1.16 on Craftax Classic, so in both environments the weights rank windows within a batch and the familiar explanation that sparse returns fail to discriminate is unavailable. The direct measurement of the return term was only possible on Craftax Classic — MiniHack needs the NetHack Learning Environment, which would not build on the machine used for that measurement — and there it gives `‖g_δ‖/‖∇L_BC‖ = 0.49 ± 0.01` at cosine `0.02 ± 0.05`. The control that settles it is also Craftax-side: advantage clipping cuts the return term fivefold and scores **3.16 points worse**, so the return weighting is not what does the damage.\n",
+    "\n",
+    "**What orders the outcome is permitted displacement.** Final KL divergence from the checkpoint correlates with final score at Spearman **−0.56** here and **−0.79** on Craftax Classic (−0.56 and −0.75 with LoRA excluded, which matters because the drift probe reads LoRA's frozen base weights and records near-zero). The sign is the same in both environments now: more displacement, lower score. Group C (architectural freezing) has the highest group mean here (**44.0%**) and Group A the tightest spread (std 0.95 points) — freezing neither destabilises the model nor helps it past the checkpoint.\n",
+    "\n",
+    "**Group-level summary** (mean ID win rate over each group, Table 7 from the paper):\n",
+    "\n",
+    "| Group | N | Mean | Δ from Baseline RL | Best | Worst | Std |\n",
+    "|---|---|---|---|---|---|---|\n",
+    "| Pretrained checkpoint | – | **47.50** | – | – | – | – |\n",
+    "| Baseline RL | 1 | 43.75 | – | – | – | 0.00 |\n",
+    "| A (Regularisation) | 6 | 42.78 | −0.97 | 44.17 | 41.25 | 0.95 |\n",
+    "| B (Training signal) | 7 | 38.57 | −5.18 | 48.75 | 12.08 | 11.64 |\n",
+    "| C (Architectural freezing) | 7 | 43.99 | +0.24 | 49.58 | 39.17 | 4.08 |\n",
+    "| D (Data quality / weights) | 4 | 38.02 | −5.73 | 44.17 | 32.08 | 4.28 |\n",
+    "\n",
+    "Every group mean is below the pretrained checkpoint, in both environments.\n",
+    "\n",
+    "Group B degrades with high variance (std 11.64), driven by the worst condition in the suite: **normalised advantages** at **12.1%** (3.73 on Craftax Classic). It alone mean-centres the weights, making roughly half of them **negative** and violating the non-negativity assumption (A1) the decomposition needs. A negative weight on a cross-entropy term is gradient *ascent* on that sequence's likelihood — unbounded and without a trust region — and its drift lands two to four orders of magnitude beyond any other condition. Supplying the repulsive direction a policy gradient would have does not rescue the objective; it destroys it.\n",
+    "\n",
+    "**On its own terms.** Craftax Classic collected return falls (5.21 → 3.73) alongside eval score, so the objective is not quietly succeeding at what it optimises. On MiniHack collected return is flat (−0.53 → −0.34, no trend across quartiles) while the training-window win rate rises from 6.3% to 13.5% and eval is flat. Neither environment looks like reward hacking under an evaluation mismatch.\n",
+    "\n",
+    "**Statistics.** With 3 seeds per arm, an exact two-sided permutation test enumerates C(6,3) = 20 relabellings, so the smallest attainable p-value is 2/20 = **0.10** — a reported p = 0.10 is the most extreme outcome the design can produce, not a null result. The shipped `significance_test.txt` compares the condition furthest from baseline (`normalized_adv`) and reports p at that floor with a bootstrap CI of **[−39.6, −23.8] points**; the paper's **[−2.1, +13.7]** interval is the different, best-versus-baseline comparison. Cell 10 prints the shipped file verbatim so the two are not confused.\n",
+    "\n",
+    "Each condition also collects gradient-alignment, representation-drift, CKA-similarity, and per-t-bin gradient-norm diagnostics. All training was done **offline** on the GPU-H200 machine — this notebook only displays the saved figures (downloaded with the HF snapshot)."
+   ],
+   "id": "bf4148dfab969b4a"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from IPython.display import Image, display\n",
+    "\n",
+    "assets_dir = os.path.join(SNAPSHOT_DIR, \"ablation_assets\")\n",
+    "if not os.path.isdir(assets_dir):\n",
+    "    raise FileNotFoundError(\n",
+    "        f\"Expected ablation assets at {assets_dir}. The HF repo must contain \"\n",
+    "        \"an ablation_assets/ directory with the pre-computed PNGs and CSVs.\"\n",
+    "    )\n",
+    "\n",
+    "figures = [\n",
+    "    (\n",
+    "        \"final_score_comparison.png\",\n",
+    "        \"Final ID win rate per condition, sorted, against the pretrained \"\n",
+    "        \"checkpoint at 47.5% under the harness protocol. Baseline RL falls to \"\n",
+    "        \"43.8% (+/-6.1). Three conditions -- head_only 49.6, gradient_surgery \"\n",
+    "        \"48.8, layer_ablation_top1 48.8 -- finish nominally above the \"\n",
+    "        \"checkpoint, but the best-versus-baseline bootstrap interval is \"\n",
+    "        \"[-2.1, +13.7] points at p = 0.40, so they are not read as \"\n",
+    "        \"improvements. The worst condition is normalized_adv at 12.1%, and it \"\n",
+    "        \"collapses because it mean-centres the weights and makes half of them \"\n",
+    "        \"negative -- gradient ascent on below-average rollouts -- not because \"\n",
+    "        \"the reward signal is uninformative. CV_A here is 2.68: the weights do \"\n",
+    "        \"rank windows.\",\n",
+    "    ),\n",
+    "    (\n",
+    "        \"group_comparison.png\",\n",
+    "        \"Score distribution by intervention group, all below the pretrained \"\n",
+    "        \"47.5%. Group C (architectural freezing) has the highest mean at 44.0, \"\n",
+    "        \"Group A (regularisation) the tightest spread at 42.8 +/- 0.95, and \"\n",
+    "        \"Groups B (38.6) and D (38.0) sit lowest. Group B's spread (std 11.6) \"\n",
+    "        \"is carried by the normalized_adv collapse. Freezing parameter subsets \"\n",
+    "        \"neither destabilises the model nor carries it past the checkpoint.\",\n",
+    "    ),\n",
+    "    (\n",
+    "        \"score_delta_over_baseline_rl.png\",\n",
+    "        \"Sorted improvement vs the baseline_rl condition. A handful of \"\n",
+    "        \"conditions beat baseline return-weighted ELBO fine-tuning; on Craftax \"\n",
+    "        \"Classic, where the effect is unambiguous, none of the 25 reaches the \"\n",
+    "        \"checkpoint it started from. What orders this ranking is how far each \"\n",
+    "        \"condition is permitted to displace the parameters (Spearman -0.56 \"\n",
+    "        \"between final KL and final score here, -0.79 on Craftax Classic), not \"\n",
+    "        \"how the return is turned into a weight.\",\n",
+    "    ),\n",
+    "    (\n",
+    "        \"per_env_delta.png\",\n",
+    "        \"Per-environment change from the pretrained checkpoint. The easy rooms \"\n",
+    "        \"sit near ceiling and move little; the losses concentrate in the \"\n",
+    "        \"layouts that need longer routing. This mirrors Craftax Classic, where \"\n",
+    "        \"deep tech-tree achievements are the first to go (tier 3 falls 23% to \"\n",
+    "        \"10%) while shallow behaviours hold.\",\n",
+    "    ),\n",
+    "    (\n",
+    "        \"gradient_alignment.png\",\n",
+    "        \"Cosine similarity between the return-weighted RL loss and the BC loss \"\n",
+    "        \"gradients over fine-tuning. NOTE: the two gradients are taken at \"\n",
+    "        \"*different parameter points*, so this is a retention diagnostic only \"\n",
+    "        \"-- it confounds weighting with drift and is unrelated to the \"\n",
+    "        \"single-parameter-point measurement of ||g_delta||/||grad L_BC|| = \"\n",
+    "        \"0.49 +/- 0.01 at cosine 0.02 +/- 0.05 quoted above.\",\n",
+    "    ),\n",
+    "]\n",
+    "\n",
+    "for fname, caption in figures:\n",
+    "    path = os.path.join(assets_dir, fname)\n",
+    "    if os.path.exists(path):\n",
+    "        print(f\"\\n=== {fname} ===\\n{caption}\\n\")\n",
+    "        display(Image(filename=path))\n",
+    "    else:\n",
+    "        print(f\"\\n[missing figure: {fname}]\")"
+   ],
+   "id": "8d35602d5ddf2c60"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 10 — Ablation results tables\n",
+    "\n",
+    "The tables below are loaded directly from the saved CSV outputs of the ablation pipeline. `main_results.csv` is the canonical sortable score table; `hypothesis_verdict.csv` attaches each condition to the mechanism hypothesis it tests, so the verdict column reads as a direct answer to *what does return-weighted ELBO fine-tuning of the diffusion planner actually do?*\n",
+    "\n",
+    "**Read the `Verdict` column carefully.** It is scored against **baseline RL**, not against the pretrained checkpoint. A row marked `IMPROVEMENT` beat the plain return-weighted objective; with three exceptions here (and none at all on Craftax Classic) it did not beat the checkpoint it started from. `Delta_Pretrained` is the column that answers that, and it is negative for 22 of the 25 conditions in MiniHack and for all 25 on Craftax Classic.\n",
+    "\n",
+    "The collective answer across all 25 conditions is that the outcome does not track how the return is turned into a weight — the Craftax-side control cuts the return term fivefold and scores 3.16 *worse* — but does track how far each condition is permitted to displace the parameters. What that leaves pointing at is the data: fine-tuning on the model's own rollouts, with the return weighting a large but incidental passenger."
+   ],
+   "id": "9cdbeab0b5332cf"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import polars as pl\n",
+    "\n",
+    "main_csv = os.path.join(assets_dir, \"main_results.csv\")\n",
+    "verdicts_csv = os.path.join(assets_dir, \"hypothesis_verdict.csv\")\n",
+    "\n",
+    "df_main = pl.read_csv(main_csv).sort(\"Score\", descending=True)\n",
+    "print(\"=== main_results.csv (sorted by Score, descending) ===\")\n",
+    "with pl.Config(tbl_rows=30):\n",
+    "    print(df_main)"
+   ],
+   "id": "75782007e322494c"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df_verdicts = pl.read_csv(verdicts_csv).sort(\"Delta_Baseline\", descending=True)\n",
+    "print(\"=== hypothesis_verdict.csv (sorted by improvement over baseline_rl) ===\")\n",
+    "with pl.Config(tbl_rows=30, fmt_str_lengths=80):\n",
+    "    print(df_verdicts)"
+   ],
+   "id": "c166b4e53f680bf2"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 11 — Conclusions\n",
+    "\n",
+    "**Empirical findings.**\n",
+    "\n",
+    "1. **The planner learns these layouts; the model-free baselines do not.** The DAgger-trained ReMDM planner reaches **48.5% in-distribution** over 50 episodes per environment (paper Table 6), and the offline-BC variant **73.0%**, against 0.3–4.7% for PPO, A2C, DQN and PPO-RNN trained from scratch. Those four are **not a matched comparison** — they use default hyperparameters and are reported only to establish that the layouts are not solvable from scratch within our budget. Any gap to the offline CNN+MLP and Decision Transformer confounds an auxiliary staircase-prediction head the baselines lack, ten denoising passes per decision against one, chunked replanning against greedy per-step selection, and untuned baselines. **This is not an architecture result**, and the current paper carries these baselines only as configurations (Table 11), not as a results table.\n",
+    "\n",
+    "2. **Zero-shot OOD transfer is weak, and we make no structural transfer claim.** The DAgger planner reaches **4.7%** mean win rate on the 3 held-out OOD maps against 48.5% in distribution (Room-Dark-15x15 12%, Corridor-R5 2%, MazeWalk-45x19 0%). The offline-BC variant reaches 6.0%. What transfer exists is to a partially observed room, not to longer routing or larger mazes.\n",
+    "\n",
+    "3. **Return-weighted ELBO fine-tuning degrades the planner — and the reward signal is not why.** In MiniHack the objective takes the checkpoint from **47.5% to 43.8% (±6.1)**; three conditions finish nominally above it, but the best-versus-baseline interval is [−2.1, +13.7] at p = 0.40, so they are not read as improvements. On Craftax Classic, where the effect is large and unambiguous, **no condition of the 25 recovers the checkpoint** and the plain objective gives up **3.59 of 11.81** points. The explanation is *not* that sparse returns fail to discriminate: `CV_A` is **2.68** here and **1.16** on Craftax Classic, and the return term itself measures `‖g_δ‖/‖∇L_BC‖ = 0.49 ± 0.01` at cosine **0.02 ± 0.05** — large and pointing away from imitation, against 0.893 ± 0.010 for `∇L_BC` against itself. **But shrinking it makes things worse.** Advantage clipping cuts the return term fivefold, to 0.097 ± 0.003, and scores **3.16 points below** the unclipped baseline. A condition with almost no return term degrades further than one with a large one.\n",
+    "\n",
+    "4. **What is left is the data.** The suite is ordered by plasticity — Spearman **−0.56** here and **−0.79** on Craftax Classic between final KL and final score — and advantage clipping, the nearest thing in the suite to unweighted training on the model's own rollouts, sits near the bottom. Together these point at fine-tuning on self-generated rollouts as the damaging ingredient, with the return weighting neither causing the harm nor preventing it. The paper is explicit that this is an inference from a near-substitute: the unweighted-ELBO-on-all-rollouts arm was not run.\n",
+    "\n",
+    "5. **It is not reward hacking under an evaluation mismatch.** On Craftax Classic collected return falls from **5.21 to 3.73** over the same 500 iterations in which eval score falls from 12.06 to 8.49. On MiniHack collected return is flat (−0.53 → −0.34) while the training-window win rate rises from 6.3% to 13.5% and eval is flat. The objective is not quietly succeeding at what it optimises.\n",
+    "\n",
+    "**Why the sign constraint is not the whole story.** The objective is a bound only for `A_i ≥ 0`, so it can re-rank sampled behaviour but never push mass away from it, which caps how much *improvement* re-ranking can deliver. That does not explain *degradation*, and it cannot be the whole story, because the same non-negativity holds of reward-weighted and advantage-weighted regression, which work elsewhere. The difference we can point to is the data: those methods are usually applied to a fixed dataset rather than to rollouts fed back through a denoising objective at every step.\n",
+    "\n",
+    "**A second, separate cost of the surrogate.** Not offered as evidence for the above: the ELBO surrogate is badly conditioned across diffusion time, and markedly worse on the harder environment. On Craftax Classic the gradient norm in the top third of the `t` range is **2.5×** the bottom third with low-`t`/high-`t` cosine similarity **0.06**; on MiniHack the ratio is **1.14** and the cosine **0.52**. Repairing the conditioning does not recover the checkpoint — on Craftax Classic `low_t` (8.12) and `t_curriculum` (8.41) sit at the baseline of 8.22 rather than above it.\n",
+    "\n",
+    "**What these results do not establish.** MiniHack is **not a second confirmation**: the effect here is small against a 6.1-point seed standard deviation, and the claim rests on Craftax Classic. There is no unweighted-ELBO-on-all-rollouts arm and no continued-DAgger arm. The gradient measurement is at one parameter point, on Craftax Classic only — MiniHack would need the NetHack Learning Environment, which would not build on the machine used for that measurement. The null characterises the clipped return ratio, not the exponential advantage weight `exp(A/β)`. And the gradient-alignment cosine logged by the training harness takes its two gradients at *different* parameter points, so it is a **retention diagnostic only**.\n",
+    "\n",
+    "**Two cheap diagnostics, worth taking first.** For anyone reaching for this objective: `CV_A` comes free from an effective-sample-size counter and says whether the returns rank anything at all. The ratio `‖g_δ‖/‖∇L_BC‖` costs two backward passes and says how much of the update the reward is responsible for. Ours said 0.49, which looked like a mechanism until the clipping control said otherwise.\n",
+    "\n",
+    "**Open problem.** Per-step formulations over the denoising chain (d1, DiffPO) admit **signed** advantages, and so change the object being optimised rather than the weighting inside a regression. The obstacle for planners like ours is that those estimators assume monotone unmasking, which ReMDM's inference-time remasking violates. What these results argue for is an estimator that tolerates remasking *and* admits a repulsive direction; **categorical flow matching** (Campbell et al., 2024) is one route to the first half."
+   ],
+   "id": "9bb18fc96b4b459a"
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

experiments/rl_finetuning/outputs/minihack_ablations/diagnosis.md

@@ -2,21 +2,21 @@
 
 ## Hypothesis Ranking
 
-### 1. Catastrophic Forgetting [*] (0%)
+### 1. Gradient Conflict [*] (33%)
 
-**Description:** Pretrained representations are corrupted by RL gradients.
+**Description:** RL and BC gradients point in conflicting directions.
 
-**Evidence:** 0/5 supporting ablations improved over baseline.
+**Evidence:** 1/3 supporting ablations improved over baseline.
 
-**Recommendation:** Implement strong parameter regularisation (EWC + LLRD) or use LoRA to restrict update space.
+**Recommendation:** Apply PCGrad and investigate t-distribution bias.
 
-### 2. Gradient Conflict [*] (0%)
+### 2. Catastrophic Forgetting [*] (17%)
 
-**Description:** RL and BC gradients point in conflicting directions.
+**Description:** Pretrained representations are corrupted by RL gradients.
 
-**Evidence:** 0/3 supporting ablations improved over baseline.
+**Evidence:** 1/6 supporting ablations improved over baseline.
 
-**Recommendation:** Apply PCGrad and investigate t-distribution bias.
+**Recommendation:** Implement strong parameter regularisation (EWC + LLRD) or use LoRA to restrict update space.
 
 ### 3. Signal Sparsity [*] (0%)
 
@@ -54,65 +54,65 @@
 
 | Ablation | Group | Score | Delta vs Pretrained |
 |---|---|---|---|
-| action_diversity | D | 0.6167 | -0.0833 |
-| advantage_clip | B | 0.5583 | -0.1417 |
-| attention_only | C | 0.6333 | -0.0667 |
-| baseline_rl | Baseline | 0.5833 | -0.1167 |
-| bc_wins | B | 0.4917 | -0.2083 |
-| entropy_bonus | B | 0.6167 | -0.0833 |
-| ewc | A | 0.6167 | -0.0833 |
-| ffn_only | C | 0.6583 | -0.0417 |
-| frozen_backbone | C | 0.6583 | -0.0417 |
-| gradient_surgery | B | 0.6917 | -0.0083 |
-| head_only | C | 0.6583 | -0.0417 |
-| kl_penalty | A | 0.5833 | -0.1167 |
-| layer_ablation_top1 | C | 0.6667 | -0.0333 |
-| layer_ablation_top2 | C | 0.6667 | -0.0333 |
-| layer_ablation_top3 | C | 0.6583 | -0.0417 |
-| llrd | A | 0.6167 | -0.0833 |
-| lora | A | 0.6833 | -0.0167 |
-| low_t | B | 0.6000 | -0.1000 |
-| mixed_replay | A | 0.5750 | -0.1250 |
-| normalized_adv | B | 0.2583 | -0.4417 |
-| reward_filtering | D | 0.6250 | -0.0750 |
-| reward_model | D | 0.5500 | -0.1500 |
-| running_stats | D | 0.5750 | -0.1250 |
-| t_curriculum | B | 0.5833 | -0.1167 |
-| trust_region_kl | A | 0.6500 | -0.0500 |
+| action_diversity | D | 0.3750 | -0.1000 |
+| advantage_clip | B | 0.3375 | -0.1375 |
+| attention_only | C | 0.3958 | -0.0792 |
+| baseline_rl | Baseline | 0.4375 | -0.0375 |
+| bc_wins | B | 0.4583 | -0.0167 |
+| entropy_bonus | B | 0.4458 | -0.0292 |
+| ewc | A | 0.4333 | -0.0417 |
+| ffn_only | C | 0.4375 | -0.0375 |
+| frozen_backbone | C | 0.4667 | -0.0083 |
+| gradient_surgery | B | 0.4875 | +0.0125 |
+| head_only | C | 0.4958 | +0.0208 |
+| kl_penalty | A | 0.4417 | -0.0333 |
+| layer_ablation_top1 | C | 0.4875 | +0.0125 |
+| layer_ablation_top2 | C | 0.4042 | -0.0708 |
+| layer_ablation_top3 | C | 0.3917 | -0.0833 |
+| llrd | A | 0.4208 | -0.0542 |
+| lora | A | 0.4333 | -0.0417 |
+| low_t | B | 0.4208 | -0.0542 |
+| mixed_replay | A | 0.4125 | -0.0625 |
+| normalized_adv | B | 0.1208 | -0.3542 |
+| reward_filtering | D | 0.4417 | -0.0333 |
+| reward_model | D | 0.3208 | -0.1542 |
+| running_stats | D | 0.3833 | -0.0917 |
+| t_curriculum | B | 0.4292 | -0.0458 |
+| trust_region_kl | A | 0.4250 | -0.0500 |
 
 | Ablation | Score | Delta vs Baseline | Verdict |
 |---|---|---|---|
-| action_diversity | 0.6167 | +0.0333 | NEUTRAL |
-| advantage_clip | 0.5583 | -0.0250 | NEUTRAL |
-| attention_only | 0.6333 | +0.0500 | NEUTRAL |
-| baseline_rl | 0.5833 | +0.0000 | NEUTRAL |
-| bc_wins | 0.4917 | -0.0917 | NEUTRAL |
-| entropy_bonus | 0.6167 | +0.0333 | NEUTRAL |
-| ewc | 0.6167 | +0.0333 | NEUTRAL |
-| ffn_only | 0.6583 | +0.0750 | IMPROVEMENT |
-| frozen_backbone | 0.6583 | +0.0750 | IMPROVEMENT |
-| gradient_surgery | 0.6917 | +0.1083 | IMPROVEMENT |
-| head_only | 0.6583 | +0.0750 | IMPROVEMENT |
-| kl_penalty | 0.5833 | +0.0000 | NEUTRAL |
-| layer_ablation_top1 | 0.6667 | +0.0833 | IMPROVEMENT |
-| layer_ablation_top2 | 0.6667 | +0.0833 | IMPROVEMENT |
-| layer_ablation_top3 | 0.6583 | +0.0750 | IMPROVEMENT |
-| llrd | 0.6167 | +0.0333 | NEUTRAL |
-| lora | 0.6833 | +0.1000 | IMPROVEMENT |
-| low_t | 0.6000 | +0.0167 | NEUTRAL |
-| mixed_replay | 0.5750 | -0.0083 | NEUTRAL |
-| normalized_adv | 0.2583 | -0.3250 | COLLAPSE |
-| reward_filtering | 0.6250 | +0.0417 | NEUTRAL |
-| reward_model | 0.5500 | -0.0333 | NEUTRAL |
-| running_stats | 0.5750 | -0.0083 | NEUTRAL |
-| t_curriculum | 0.5833 | +0.0000 | NEUTRAL |
-| trust_region_kl | 0.6500 | +0.0667 | IMPROVEMENT |
-
-*Pretrained score: 0.7000*
-*Baseline RL score: 0.5833*
+| action_diversity | 0.3750 | -0.0625 | COLLAPSE |
+| advantage_clip | 0.3375 | -0.1000 | COLLAPSE |
+| attention_only | 0.3958 | -0.0417 | NEUTRAL |
+| baseline_rl | 0.4375 | +0.0000 | NEUTRAL |
+| bc_wins | 0.4583 | +0.0208 | NEUTRAL |
+| entropy_bonus | 0.4458 | +0.0083 | NEUTRAL |
+| ewc | 0.4333 | -0.0042 | NEUTRAL |
+| ffn_only | 0.4375 | +0.0000 | NEUTRAL |
+| frozen_backbone | 0.4667 | +0.0292 | IMPROVEMENT |
+| gradient_surgery | 0.4875 | +0.0500 | IMPROVEMENT |
+| head_only | 0.4958 | +0.0583 | IMPROVEMENT |
+| kl_penalty | 0.4417 | +0.0042 | NEUTRAL |
+| layer_ablation_top1 | 0.4875 | +0.0500 | IMPROVEMENT |
+| layer_ablation_top2 | 0.4042 | -0.0333 | NEUTRAL |
+| layer_ablation_top3 | 0.3917 | -0.0458 | NEUTRAL |
+| llrd | 0.4208 | -0.0167 | NEUTRAL |
+| lora | 0.4333 | -0.0042 | NEUTRAL |
+| low_t | 0.4208 | -0.0167 | NEUTRAL |
+| mixed_replay | 0.4125 | -0.0250 | NEUTRAL |
+| normalized_adv | 0.1208 | -0.3167 | COLLAPSE |
+| reward_filtering | 0.4417 | +0.0042 | NEUTRAL |
+| reward_model | 0.3208 | -0.1167 | COLLAPSE |
+| running_stats | 0.3833 | -0.0542 | COLLAPSE |
+| t_curriculum | 0.4292 | -0.0083 | NEUTRAL |
+| trust_region_kl | 0.4250 | -0.0125 | NEUTRAL |
+
+*Pretrained score: 0.4750*
+*Baseline RL score: 0.4375*
 
 ## Aggregate Verdict
 
-Mixed results: 8/24 ablations collapsed. Check individual verdicts above.
+Mixed results: 5/24 ablations collapsed. Check individual verdicts above.
 
-**Gradient alignment** = +0.5747: RL gradient has useful signal.
+**Gradient alignment** = +0.1107: RL gradient has useful signal.