ReMDM Planner — Discrete Diffusion Planning on Craftax
A JAX implementation of ReMDM (Remasking Discrete Diffusion Model) for action-sequence planning in the Craftax environment. A bidirectional transformer learns to generate action plans by iteratively denoising masked token sequences, conditioned on the current environment observation.
Description
The planner starts from a fully-masked action sequence and iteratively unmasks tokens over T denoising steps, producing a plan_horizon-length plan. The ReMDM framework extends standard Masked Discrete Language Modelling (MDLM) with remasking strategies that allow committed tokens to be re-predicted, improving plan coherence.
Training follows a four-stage pipeline:
[Stage 1] Train PPO agent Craftax_Baselines/ppo_rnn.py | ppo_rnd.py
|
v checkpoint
[Stage 2a] Collect trajectories main.py --mode collect (optional)
|
v .npz file
[Stage 2b] Train offline main.py --mode offline
| (from .npz or live PPO rollouts)
v diffusion checkpoint
[Stage 3] Online fine-tuning main.py --mode online
|
v fine-tuned checkpoint
[Stage 4] Evaluate main.py --mode inference
Installation
Prerequisites (system-level)
uv manages Python packages only. The following must be installed at the OS level before
running on a GPU node — they are not in pyproject.toml:
CUDA 13 driver and toolkit (libcuda.so, libcudnn)
On HPC clusters these are typically loaded via module load cuda/13.x.
1. Create the virtual environment
bash
1# CPU-only (local development / macOS)2uv sync34# NVIDIA CUDA 13 (GPU node — Linux only)5uv sync --extra cuda
67# Activate8source .venv/bin/activate
uv sync reads pyproject.toml, resolves a fully-reproducible lockfile (uv.lock),
and installs into .venv/. Commit uv.lock to pin the exact dependency graph.
2. Initialise the submodule
git submodule update --init --recursive
Dependencies
Package
Version
Role
jax
>=0.9.2
JIT compilation and functional arrays
flax
>=0.12.6
Neural network definitions
optax
>=0.2.8
Adam optimiser and gradient clipping
craftax
>=1.5.0
Procedurally-generated Minecraft-like environment
chex
>=0.1.91
JAX testing and assertion utilities
distrax
>=0.1.7
Probability distributions
orbax
>=0.1.9
Model checkpointing
wandb
>=0.25.1
Experiment logging
numpy
>=2.4.4
Array operations
matplotlib
>=3.10.8
Plotting
polars
>=1.39.3
DataFrame analysis
orjson
>=3.11.8
Fast JSON serialisation
pyyaml
>=6.0.3
Config file parsing
Full specification in pyproject.toml. Exact transitive pins are in uv.lock.
Usage
All modes share the same entry point. Defaults are loaded from configs/defaults.yaml; any value can be overridden on the command line.
The file stores arrays shaped [num_envs, num_iters, ...], preserving per-environment contiguity so episode boundaries are respected during window sampling.
Stage 2b — Train offline from live PPO rollouts
Roll out the PPO agent live at each update step and train the diffusion model on the collected windows. Windows that cross episode boundaries are masked out; windows with higher cumulative reward receive proportionally larger gradient contributions (clipped to [0.1, return_weight_cap]).
The diffusion model (learner) is fine-tuned via DAgger (Dataset Aggregation). At each iteration a mixed policy blends the PPO expert and the diffusion learner (controlled by an exponentially decaying beta). The mixed policy rolls out trajectories; the expert labels every visited state with the action it would take. These (obs, expert_plan) pairs are appended to a growing circular replay buffer, and the diffusion model is retrained on the full buffer with the standard MDLM ELBO loss (pure behavioural cloning — no reward weighting).
Prints mean episode return, completed episodes, steps per second, and per-achievement unlock counts. Uses historical inpainting: the first hist_len plan positions are locked to observed history.
Loading checkpoints from W&B artifacts
Any checkpoint path argument (--checkpoint_path, --offline_checkpoint_path, --ppo_checkpoint_path) accepts a W&B artifact reference prefixed with wandb:. The artifact is downloaded automatically before training or evaluation begins.
Control the download location with --wandb_download_dir (defaults to ./artifacts/).
Resuming a Training Run
A completed training checkpoint can be used as the starting point for a new run that continues where the previous one left off. This is useful when extending the training budget or when a preempted job needs to be restarted.
Offline resume:
bash
1# Auto-detect step and wandb run ID from checkpoint metadata2python main.py --mode offline \3 --ppo_checkpoint_path /path/to/ppo_checkpoint \4 --resume_checkpoint_path /path/to/completed_offline_checkpoint \5 --total_timesteps 200000000\6 --save_policy
78# Explicit step and wandb run ID override9python main.py --mode offline \10 --ppo_checkpoint_path /path/to/ppo_checkpoint \11 --resume_checkpoint_path /path/to/completed_offline_checkpoint \12 --resume_step 1525\13 --resume_wandb_run_id abc123xyz \14 --total_timesteps 200000000\15 --save_policy
1617# Resume from a W&B artifact18python main.py --mode offline \19 --ppo_checkpoint_path /path/to/ppo_checkpoint \20 --resume_checkpoint_path wandb:my-team/remdm-craftax/policy:latest \21 --total_timesteps 200000000\22 --save_policy
The DAgger replay buffer is not persisted across resumes. It starts empty and refills within the first few iterations.
JIT compilation is fully preserved. Resume only affects initialisation outside jax.jit (loading checkpoint, setting the optimizer step counter, adjusting scan length).
The cosine LR schedule is constructed for the full num_updates range. The optimizer step counter is set to the resume offset so the learning rate picks up exactly where the previous run stopped.
When resume_checkpoint_path points to a checkpoint with a metadata sidecar, resume_step and resume_wandb_run_id are auto-detected. Explicit CLI flags override the metadata values.
Checkpoints without a metadata sidecar (created before this feature) still load; provide --resume_step explicitly.
Configuration
All hyperparameters are in configs/defaults.yaml. Override any value on the command line:
JAX functional purity: training closures (make_train, make_train_dagger) are fully JIT-compatible. Environment construction and checkpoint I/O happen outside jax.jit.
Offline training: --mode offline rolls out the PPO agent live at each update step via make_train. Use --mode collect to save a trajectory .npz for inspection or analysis; re-feeding it to --mode offline is not supported — pass --ppo_checkpoint_path instead.
Episode-boundary masking: the offline sampler pre-computes a validity mask over all (env, time) positions. A window at (e, t) is valid only if dones[e, t+1:t+H-1] are all False.
Return weighting: valid windows are weighted by their cumulative reward, normalised by the batch mean and clipped to [0.1, RETURN_WEIGHT_CAP]. Weights are passed as per-sample multipliers into the MDLM loss before reduction, so they correctly scale each sample's gradient contribution.
LR schedule: cosine decay from lr to lr * 0.1 over all gradient steps. Set lr_warmup_steps > 0 to prepend a linear warm-up phase.
Loss weight clipping: the MDLM SUBS weight -alpha'(t) / (1 - alpha_t) is clipped to 1000 to prevent numerical instability when alpha_t ≈ 1.
Validation rollouts: during offline training, a held-out rollout runs every val_interval steps. It uses the same sampling parameters as inference (remask_strategy, eta, use_loop, t_on, t_off, temperature, top_p) with val_diffusion_steps denoising steps and val_replan_every env steps per plan, for a total of val_steps environment steps.
W&B logging: all metric aggregation is centralised in src/planners/logging.py. Metric namespaces: diffusion/ (loss, accuracy), train/ (data quality, throughput), env/ (episode returns, achievements), val/ (validation rollouts, emitted every val_interval steps), dagger/ (online DAgger training: beta, buffer fill, reward mean, valid fraction). train/sps (environment frames/sec) is only logged in modes that perform live environment interaction.
DAgger dataset aggregation: online training (--mode online) implements DAgger (Ross et al., 2011). A circular replay buffer accumulates (obs, expert_plan) pairs across all iterations. Each update samples uniformly from the full buffer, not just the latest batch. Training samples that cross episode boundaries (any done within the plan-horizon window) are marked invalid. The expert (PPO agent) receives correct done flags so its RNN hidden state resets on episode boundaries.
Denoising step indexing: the reverse scan runs from step_idx = 0 to T-1, mapping to diffusion time t = (T - step_idx) / T (high noise to low noise).
Submodule PPO agents: PPO training lives entirely in Craftax_Baselines/. Planner scripts only consume pre-trained checkpoints via --ppo_checkpoint_path.
