PyTorch implementation of ReMDM (Remasking Discrete Diffusion Model) for action-sequence planning in MiniHack navigation environments. A dual-stream transformer generates 64-step action plans by iteratively denoising masked token sequences, conditioned on a 9x9 local crop and the full 21x79 dungeon map.
The primary training method is DAgger with BFS oracle supervision: the buffer is seeded with pure expert trajectories on the first iteration, providing an implicit behavioural cloning warm-start. An optional standalone offline BC mode is available for pre-training on collected datasets. Generalises zero-shot from 4 in-distribution environments to 3 out-of-distribution environments.
Pipeline
[Primary] DAgger online training main.py --mode dagger
| (seed buffer with oracle demos on iter 0,
| collect with model, label with oracle,
| efficiency filter, curriculum sampling)
v checkpoint
[Evaluate] ID + OOD evaluation main.py --mode inference --checkpoint iter8000.pth
DAgger with implicit warm-start is the recommended pipeline. The --mode collect + --mode offline path is available for explicit two-stage pre-training on oracle demonstrations before DAgger.
Environments
In-distribution (training):
Environment
Description
MiniHack-Room-Random-5x5-v0
Small random room
MiniHack-Room-Random-15x15-v0
Large random room
MiniHack-Corridor-R2-v0
Two-room corridor
MiniHack-MazeWalk-9x9-v0
Small maze
Out-of-distribution (zero-shot evaluation):
Environment
Description
MiniHack-Room-Dark-15x15-v0
Dark room (limited visibility)
MiniHack-Corridor-R5-v0
Five-room corridor
MiniHack-MazeWalk-45x19-v0
Large maze
Installation
Prerequisites
Python 3.12+ is required.
macOS (arm64): Install cmake via Homebrew (needed to compile nle from source):
brew install cmake
Linux (x86_64): Pre-built wheels are available, but if building from source:
This installs all dependencies from the lockfile, including nle>=1.2.0 (from the maintained NetHack-LE fork), minihack, torch>=2.11.0, wandb, polars, orjson, and scipy.
GPU support (optional)
By default PyTorch runs on CPU. For NVIDIA CUDA 12:
Collects a few oracle trajectories, trains for 30 iterations, and prints ID evaluation results.
python main.py --mode smoke
Collect oracle demonstrations
Run the BFS oracle across all 4 ID environments and save the trajectories as a .pt dataset for offline BC training. Uses multiprocessing for parallelism.
bash
1# Default: 5000 episodes per env, output to data/dataset.pt2python main.py --mode collect
34# Custom episode count and output5python main.py --mode collect collect_episodes_per_env=2000\6collect_output=data/small_dataset.pt
78# Fewer workers (default: 8)9python main.py --mode collect collect_num_workers=41011# Reproducible with fixed seed12python main.py --mode collect seed=42
The output .pt file is directly consumable by --mode offline:
The model takes (local_obs, global_obs, noisy_action_seq, t_discrete) and returns {"actions": [B,64,12], "goal_pred": [B,2]}.
A LocalDiffusionPlanner variant (no global stream, no goal head) is also available for ablation studies.
Diffusion
Forward process (MDLM): Each action token is independently replaced with MASK (token 12) with probability 1 - alpha(t), where alpha(t) follows a linear or cosine schedule. PAD tokens (13) are never masked.
Loss: Cross-entropy on masked positions only, averaged globally across the batch. By default uses a flat average (matching the reference implementation). Optional SUBS importance weighting w(t) = -alpha'(t) / (1 - alpha(t)), clipped to [0, 1000], can be enabled via use_importance_weighting: true. Optional label smoothing via label_smoothing (default 0.0).
Reverse sampling (ReMDM): Over K denoising steps (default 10):
Model predicts logits; apply temperature scaling and top-K filtering.
Sample predictions; compute per-token confidence.
MaskGIT unmask: commit the n_unmask highest-confidence masked positions.
ReMDM remask: stochastically re-mask committed positions to allow refinement.
Final step: commit all remaining positions.
Greedy sampling: Used during DAgger data collection for deterministic rollouts. Same MaskGIT progressive unmasking loop but with argmax decoding (no temperature, no top-K, no remasking). Uses fewer denoising steps (diffusion_steps_collect: 5) for faster collection.
UCL GPU learning behaviour study (eta=0.18, B=6144)
configs/ucl_gpu_no_amp.yaml
UCL GPU without AMP (B=3584, 32 workers)
DAgger Training Loop
Each DAgger iteration:
Curriculum sampling: Select an environment weighted by difficulty (low win-rate environments sampled more).
Model rollout: Generate plans with the EMA model using greedy sampling; execute with replanning every 16 steps. Collects episodes_per_iteration (default 30) episodes per iteration.
Oracle rollout: Run the BFS oracle on the same seed for comparison.
Efficiency filter: Add the oracle trajectory to the buffer if the model failed or took >1.5x the oracle's steps.
Training: Sample from the replay buffer; run grad_steps_per_iteration gradient steps, updating EMA weights after each gradient step.
Collection uses GPU-batched rollouts when on CUDA with episodes_per_iteration > 1, falling back to threaded CPU collection or sequential collection as appropriate.
The BFS oracle uses a 5-tier priority: (1) kick adjacent doors, (2) BFS to staircase, (3) BFS to frontier, (4) BFS to farthest tile, (5) random cardinal.
Reward Shaping
The environment wrapper applies shaped rewards to guide learning:
Inference uses EMA weights by default. Pass --no-ema to use training weights.
W&B Artifacts
Checkpoints are automatically uploaded as versioned W&B artifacts (type "model") at each checkpoint save. Each artifact contains the .pth weights and a config.yaml snapshot of all hyperparameters used.
The artifact reference format is entity/project/artifact-name:version where version is latest, v0, v1, etc.
W&B Run Resumption
All training loops save the W&B run ID in their checkpoints. When resuming from a checkpoint, the run ID is automatically extracted and passed to wandb.init(resume="must"), so metrics continue on the same W&B curves with no gaps.
bash
1# DAgger: automatic -- run ID is read from the checkpoint2python main.py --mode dagger --checkpoint checkpoints/iter2000.pth
34# Offline BC: automatic5python main.py --mode offline --data dataset.pt \6 --checkpoint checkpoints/offline_epoch10.pth
78# Manual override (e.g. checkpoint saved before this feature was added):9python main.py --mode dagger --checkpoint old_checkpoint.pth \10wandb_resume_id=abc123xyz
1112# Ablation suite:13python experiments/rl_finetuning/run_ablations.py \14 --checkpoint path/to/ckpt.pth --all --use_wandb \15 --wandb_resume_id abc123xyz
The run ID is visible in the W&B dashboard URL: wandb.ai/.../runs/<run-id>.
Performance Tuning
Three config keys control performance optimisations. Defaults are set for GPU training; override for CPU or different hardware.
Mixed precision (use_amp: true)
Wraps training forward/backward in torch.amp.autocast("cuda") with GradScaler. Active in both offline BC and DAgger training.
Measured speedup: 2.2x on gradient steps, 1.7x on full smoke test wall-clock
Memory: peak GPU stays ~16 GB at B=3584 (same as FP32 due to embedding-heavy model)
Correctness: loss trajectory and win rates statistically equivalent to FP32
When to use: always on GPU. No effect on CPU (autocast is a no-op)
Default:false in defaults.yaml; enabled in GPU-specific configs
torch.compile (torch_compile: true)
Applies torch.compile(model, mode="default") before training. Falls back gracefully if no C compiler is found (common on managed GPU nodes).
Measured speedup: none beyond AMP alone. Not recommended for primary training.
Default:true in defaults.yaml
When to use: experimental only. May help on future PyTorch versions with better dynamic shape support.
Parallel collection (num_collection_workers: N)
DAgger episode collection supports three strategies (auto-selected):
GPU-batched (default on CUDA with episodes_per_iteration > 1): all envs in lockstep
Threaded CPU (fallback when num_collection_workers > 0): ThreadPoolExecutor with CPU model copies
Sequential (reference behaviour): one episode at a time
Default:8 workers in defaults.yaml
When to use: GPU-batched is preferred; workers primarily affect the CPU fallback path
Profiling
Run python scripts/profile_dagger.py [key=value ...] to profile DAgger iteration components. Supports all config overrides (e.g., use_amp=true).
Implementation Notes
MDLM loss returns 0.0 (not NaN) when no masked positions exist in the batch. Uses global averaging by default; SUBS importance weighting is opt-in via use_importance_weighting: true.
PAD tokens are never masked during the forward process and are excluded from the loss.
Sampling paths: Evaluation uses stochastic ReMDM sampling (temperature, top-K, remasking) with diffusion_steps_eval (default 10) steps. DAgger collection uses greedy argmax sampling (deterministic, no remasking) with diffusion_steps_collect (default 5) steps for faster rollouts.
remdm_sample guarantees a fully committed output (no MASK tokens) via a final-step commit and an assertion check. A min-keep 10% safety net prevents degenerate all-masked states.
EMA shadow weights are updated after every gradient step (not per iteration). The DataCollector syncs the latest EMA weights before each rollout.
Curriculum initialises with a 50/50 prior per environment (configurable via curriculum_preseed) and uses bucket-based weights: low win-rate (0.2), medium (1.0), high (0.1).
Replay buffer pins offline data at the front; only online samples are FIFO-evicted. Returns None on empty buffer (callers handle gracefully).
Global gate initialises at sigmoid(-3.0) ~ 0.047, starting nearly closed to prevent the global stream from destabilising early training.
Dropout is set to 0.0 by default. The discrete diffusion forward masking already regularises; dropout on top is redundant.
DAgger warm-start: On iteration 0, the buffer is seeded with 3 oracle trajectories per ID environment (12 total), giving the curriculum and training loop data to work with immediately.
