remdm-planner-craftax

ReMDM Planner — Craftax Classic artefacts

Anonymous release accompanying Return-Weighted ELBO Fine-Tuning Degrades Masked Diffusion Planners (under double-blind review). It contains the code, the trained checkpoints and the pre-computed results behind every Craftax Classic number in the paper.

Anonymity notice. This repository is an anonymised mirror prepared for double-blind review. Author names, institutional affiliations, experiment-tracking identifiers and absolute filesystem paths have been removed or replaced with placeholders. Cluster hardware is referred to only as GPU-24GB and GPU-H200.

What is here

src/  Craftax_Baselines/  configs/  experiments/  scripts/  tests/
main.py  pyproject.toml  uv.lock  Dockerfile  LICENSE  demo_craftax.ipynb

checkpoints/
  offline/Craftax-Classic-Symbolic-v1-Offline-Diffusion-BC-100M/   Offline BC planner
  online/Craftax-Classic-Symbolic-v1-Online-Diffusion-DAgger-100M/ DAgger planner (the
                                                                   checkpoint every
                                                                   ablation fine-tunes)
  ppo_agents/                                                      PPO-RNN experts

experiments/rl_finetuning/outputs/craftax_classic_ablations/   The published 25-condition suite
  figures/   113 pre-computed PNGs
  tables/    main_results.csv, group_summary.csv, hypothesis_verdict.csv,
             achievement_summary.csv, gradient_analysis.csv, significance_test.txt,
             gdelta.csv, results.tex, ...
  gdelta/    gdelta_{seed0,seed1,seed2,aggregate}.json   <- the return-term measurement
  results.json, diagnosis.md

experiments/rl_finetuning/outputs/review_*/    Five review arms, outside the 25-condition suite
  review_anchor_baseline_rl/          same-stack anchor for the four below
  review_run1_bc_all/                 unweighted control (every weight fixed at one)
  review_run2_advclip_lr_matched/     advantage_clip at a matched effective step
  review_run4_baseline_lr1e-4/        learning-rate sweep
  review_run4_baseline_lr1e-5/

results/inference/     Single-seed evaluations through the inpainting sampler (see below),
                       plus PPO-RNN expert evaluations
results/paper_figures/ The 11 manuscript figures as vector PDF

Headline numbers (Craftax Classic)

No condition of the 25 recovers the checkpoint it started from. Group means are 9.74 (A), 7.07 (B), 10.53 (C) and 7.97 (D).

The review arms

These are separate runs on a different host and are quoted against each other, not against the table above: the suite does not reproduce across hosts, and the same checkpoint and configuration score the pretrained planner at 11.9754 here against 11.8084 above, a consistent offset of 1.41%. baseline_rl was therefore re-run as a same-stack anchor.

Removing the return weighting while keeping the rollouts costs more than the weighting itself (−3.60 against the anchor). Lowering the learning rate by a factor of thirty recovers 68% of the degradation without removing it — 10.8104 still sits 1.1650 below the same-stack checkpoint, more than ten times the seed sd at that rate. Matching the effective step does not close the advantage-clipping gap.

The MiniHack counterparts of all three arms were run and are reported in the paper, but their run directories are not part of either release — only the Craftax Classic ones are published here. The MiniHack numbers are quoted in that repo's notebook.

The return term

experiments/rl_finetuning/outputs/craftax_classic_ablations/gdelta/ holds the direct measurement of the return term g_δ at the pretrained checkpoint, over 3 rollout seeds × 8 noise draws on 1,024-window batches. Reproduce it on a CPU with run_ablations.py --measure-gdelta. The Eq.-4 correctness residual is 4.8 × 10⁻⁷.

Two evaluation paths — read this before comparing numbers

Every number in the paper comes from the ablation harness (experiments/rl_finetuning/ablations/training.py::build_eval_fn): sample_plan with no locked prefix, 50 denoising steps, 8 actions executed per plan, scored as returned_episode_returns.

main.py --mode inference can also take a different path, sample_plan_inpainting, which freezes every executed action as an inpainting prefix and replans every step. It is a different planner at evaluation time and scores far lower on the same weights (3.26 against 11.81 for the DAgger checkpoint). The single-seed JSONs under results/inference/ were produced by that path and are kept as an ablation on the planning-as-inpainting design choice — they are not paper numbers. configs/defaults.yaml now sets inference_sampler: sample_plan, so the two paths agree by default.

demo_craftax.ipynb evaluates through the harness path (Cell 5) and runs the inpainting path only as a clearly labelled ablation (Cell 5b).

Quick start

from huggingface_hub import snapshot_download
path = snapshot_download(repo_id="AnonMLuser/remdm-planner-craftax", local_dir="remdm-planner-craftax")

Then open demo_craftax.ipynb, or see the project documentation below.

Citation

@inproceedings{anonymous2026returnweighted,
  title     = {Return-Weighted ELBO Fine-Tuning Degrades Masked Diffusion Planners},
  author    = {Anonymous},
  booktitle = {Submitted to NeurIPS 2026},
  year      = {2026},
  note      = {Under double-blind review}
}

---

ReMDM Planner for Craftax

JAX implementation of ReMDM (Remasking Discrete Diffusion Model) for action-sequence planning in Craftax, a JAX-accelerated, procedurally generated open-world survival game. A bidirectional transformer generates plan_horizon-length action plans by iteratively denoising masked token sequences, conditioned on the current symbolic observation. Trained under a pre-trained PPO expert, either offline (behavioural cloning on live rollouts) or online (DAgger).

The sibling repository remdm-planner-minihack implements the same method in PyTorch on MiniHack. 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 T denoising steps; ReMDM extends MDLM with remasking strategies that let committed tokens be re-predicted, improving plan coherence.

Two independent training pipelines are compared head-to-head, both fed by one pre-trained PPO expert checkpoint: --mode offline behaviour-clones from live expert rollouts, --mode online runs DAgger from scratch against expert labels. Either output is scored with --mode inference.

Setup

Prerequisites: Python 3.12+, uv. Linux GPU use needs NVIDIA driver >= 580 for CUDA 13, or >= 525 with --extra cuda12. CUDA and cuDNN come from the pip wheels, so no OS-level toolkit is required; if module load cuda/13.x is in your shell profile, unset LD_LIBRARY_PATH, which otherwise shadows the wheel libraries.

git clone https://github.com/ANONYMOUS/remdm-planners.git
cd remdm-planner-craftax
# Or, if already cloned without submodules:
git submodule update --init --recursive

# Default: CPU-only JAX (macOS, or Linux without a GPU).
# Installs the dev group (pytest) too.
uv sync

# Linux GPU, CUDA 13 (driver >= 580)
uv sync --extra cuda13

# Linux GPU, CUDA 12 fallback (driver >= 525, or Maxwell/Pascal cards)
uv sync --extra cuda12

Extras: cuda13 and cuda12 are mutually exclusive and Linux-only. JAX ships GPU support only through these extras, so a GPU node needs one explicitly.

Repo layout

remdm-planner-craftax/
├── Craftax_Baselines/       Git submodule — PPO expert training and env wrappers
├── configs/                 Experiment configs (defaults.yaml + presets, see Configuration)
├── src/                     Model, diffusion, planner pipelines
├── experiments/
│   └── rl_finetuning/       RL fine-tuning ablation suite (run_ablations.py)
├── scripts/                 Param counter, PPO evaluator, paper figures, HF upload, provisioning
├── tests/                   Smoke suite — uv run pytest
├── checkpoints/             Gitignored — offline/, online/, ppo_agents/ (see Checkpoints)
├── results/                 Gitignored, created on demand — inference/ eval JSONs and
│                            paper_figures/ manuscript PDFs, both published (see Checkpoints)
├── demo_craftax.ipynb       Demo notebook
├── main.py                  CLI entry point
└── pyproject.toml           uv project — deps, cuda12/cuda13 extras, dev group

Quickstart

Full DAgger pipeline (rollout, expert labelling, gradient updates, validation) under configs/smoke.yaml, ~25 s on CPU. The expert is randomly initialised unless --ppo-checkpoint is given, so this runs on a clean clone with no downloads. Watch mean step reward, loss and all metrics finite; returns stay at 0.000, since no episode terminates in so short a run.

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. All training modes need a PPO expert checkpoint.

Stage 1 — Train the PPO expert (submodule)

cd Craftax_Baselines
python ppo_rnn.py --env_name Craftax-Classic-Symbolic-v1 \
    --total_timesteps 1000000000 --save_policy --use_wandb
cd ..

(ppo_rnd.py for Random Network Distillation.) Released experts are on the HF Hub, see Checkpoints.

Offline BC

Rolls out the PPO agent live at each update.

python main.py --mode offline --ppo-checkpoint /path/to/ppo_checkpoint

Online DAgger

Trained from scratch. Per iteration a mixed expert/learner policy rolls out, the expert labels every visited state, and the model trains on the aggregated buffer.

python main.py --mode online --ppo-checkpoint /path/to/ppo_checkpoint

# Optional: warm-start from a pre-trained offline checkpoint
python main.py --mode online --ppo-checkpoint /path/to/ppo_checkpoint \
    --checkpoint /path/to/offline_checkpoint

With save_policy: true (default) and W&B on, training uploads two artifacts, either consumable via --checkpoint wandb:…: {env_name}-policy (final) and {env_name}-policy-best (highest validation return).

Collect trajectories to disk

Rolls out the PPO checkpoint and saves (obs, actions, rewards, dones) as .npz, for inspection; --mode offline does not consume it (it rolls out live).

python main.py --mode collect --ppo-checkpoint /path/to/ppo_checkpoint \
    --data data/trajectories.npz \
    --override collect_num_steps=1000000 --override collect_num_envs=128

Resuming a training run

# Same shape for --mode online. --resume also accepts a wandb: artifact reference.
python main.py --mode offline --ppo-checkpoint /path/to/ppo_checkpoint \
    --resume /path/to/completed_offline_checkpoint \
    --override offline_total_timesteps=200000000

The DAgger replay buffer is not persisted; it refills within a few iterations. The cosine LR schedule spans the full num_updates, offset so the LR resumes where it stopped. resume_step and resume_wandb_run_id come from the metadata sidecar; without one, pass --resume-step (or --resume-wandb-run-id) explicitly.

--resume restores the optimiser state, so it needs a checkpoint written by the current AdamW chain; an older one fails loudly and there is no compatibility path. Use --checkpoint instead — parameters only, warm-starting a fresh run.

Evaluation from a checkpoint

python main.py --mode inference --checkpoint /path/to/checkpoint --output results/inference/eval.json

Prints steps per second, per-achievement unlock counts, and two returns that must not be quoted against one another; --output also writes both as JSON:

By default this replans from scratch every eval_replan (8) steps, conditioned only on the current observation — the same sampler and cadence as build_eval_fn in the ablation harness, so it is the protocol behind the published numbers. Length and width come from eval_steps / eval_num_envs.

--override inference_sampler=inpainting switches to the historical-inpainting sampler, which replans every step with the executed actions locked as a fixed prefix, leaving fewer free positions the further into a window it gets. It is kept as an ablation on the planning-as-inpainting design choice and scores far lower on the same weights; no published number comes from it.

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 restore. Every released diffusion checkpoint carries the defaults.yaml architecture, so evaluate with the matching final_* config, which also sets the right env_name and recipe values:

python main.py --mode inference \
    --config configs/final_craftax_classic_gpu_24gb.yaml \
    --checkpoint checkpoints/online/Craftax-Classic-Symbolic-v1-Online-Diffusion-DAgger-100M

Any checkpoint flag (--checkpoint, --ppo-checkpoint, --resume) accepts a W&B artifact reference prefixed wandb:; the artifact downloads automatically (location: wandb_download_dir, default ./artifacts/).

python main.py --mode inference \
    --checkpoint wandb:my-team/remdm-planner-craftax/Craftax-Classic-Symbolic-v1-policy:latest

Baselines and ablations

RL baselines

PPO baselines (the expert family: ppo, ppo_rnn, ppo_rnd) train in the Craftax_Baselines submodule, see Training. Evaluate an expert with scripts/eval_ppo_expert.py:

uv run python scripts/eval_ppo_expert.py \
    --path checkpoints/ppo_agents/Craftax-Classic-Symbolic-v1-PPO_RNN-1000M \
    --env-name Craftax-Classic-Symbolic-v1

Method ablations (named configs)

Each paper experiment is a named config; pass it via --config:

python main.py --mode online --ppo-checkpoint <ppo> --config configs/classic_exp_a_beta_fix.yaml
python main.py --mode online --ppo-checkpoint <ppo> --config configs/classic_exp_b_beta_big_model.yaml
python main.py --mode online --ppo-checkpoint <ppo> --config configs/classic_exp_c_full_recipe.yaml
python main.py --mode online --ppo-checkpoint <ppo> --config configs/classic_exp_d_850K_model.yaml

RL fine-tuning ablation suite

26 registered ablations (same names as in the minihack repo). See experiments/README.md.

python experiments/rl_finetuning/run_ablations.py --list
python experiments/rl_finetuning/run_ablations.py \
    --checkpoint $PRETRAINED_CKPT --all
python experiments/rl_finetuning/run_ablations.py \
    --checkpoint wandb:my-team/remdm-planner-craftax/Craftax-Classic-Symbolic-v1-policy-best:latest \
    --ablations baseline_rl kl_penalty --fast

The same entry point measures the return term of the gradient decomposition at the pretrained checkpoint, with no training and no accelerator:

python experiments/rl_finetuning/run_ablations.py --measure-gdelta --gdelta-seeds 0 1 2 \
    --checkpoint $PRETRAINED_CKPT --results-path $RUN/results.json --output-dir $RUN

Configuration

One YAML config holds the experiment; the CLI holds the run.

Precedence, lowest to highest: configs/defaults.yaml < --config preset < --override and run flags. Exactly two config layers — a preset never inherits from another preset.

• Config files (configs/*.yaml): hyperparameters, model and method settings, ablation definitions.
• Run flags: --seed, --checkpoint, --ppo-checkpoint, --data, --output, --resume*, --jit/--no-jit (disable JIT for debugging).
• --override KEY=VALUE (repeatable): keys are validated against defaults.yaml and cast to the key's type, so a typo is an error, not a silent no-op.

defaults.yaml is the final Craftax Classic recipe, not a neutral baseline. Run main.py with no --config and you get the paper's Classic DAgger run: a 384-dim, 6-layer model over 100M env frames.

Presets hold only deltas. A key belongs in a preset only if its value differs from defaults.yaml; restating one silently pins the preset when the recipe later moves. tests/test_config.py enforces this.

Schedule keys are denominated in env frames, not update steps. Six settings — lr_warmup_frames, offline_total_timesteps, online_total_timesteps, dagger_beta_final, dagger_buffer_cycles, val_interval_frames — declare the hardware-invariant quantity; resolve_num_updates() and resolve_scaled_hyperparams() derive the update-step forms the runners consume (num_updates, LR_WARMUP_STEPS, DAGGER_BETA_DECAY, DAGGER_BUFFER_MAX, VAL_INTERVAL) from them at load. Set the frame-denominated key; the derived ones are outputs, not inputs.

python main.py --mode offline --ppo-checkpoint <ppo> --no-jit \
    --override lr=1e-4 --override plan_horizon=64 --override num_envs=4

Within each family the two machine configs differ only in num_envs and seed, guarded by test_cluster_siblings_differ_only_in_num_envs_and_seed rather than by the loader. A Full Craftax hyperparameter change must be made in both final_craftax_* files — with no inheritance those 8 keys are duplicated verbatim in each; a Classic one belongs in defaults.yaml.

Key hyperparameters are documented inline in configs/defaults.yaml; the appendix tabulates the load-bearing ones. Ablation-suite hyperparameters live in experiments/rl_finetuning/configs/, loaded by run_ablations.py, not main.py.

Checkpoints

With save_policy: true (the default), training saves Orbax checkpoints to policies (final) and policies_best (highest validation return) — under wandb.run.dir with W&B on, uploaded as {env_name}-policy and {env_name}-policy-best, and under {checkpoint_dir}/{mode}/{run_name}/ with W&B off, so a run never discards its weights. Diffusion checkpoints carry a resume_metadata.json sidecar recording the producing run's config, which is what --resume reads; PPO checkpoints carry config.yaml and wandb-summary.json.

Pass the checkpoint directory, not the step subdirectory — CheckpointManager resolves the latest step itself. Offline checkpoints save at the resolved env-frame budget: 99,942,400 for the Classic recipe at 512 envs (1525 updates × 512 × 128).

checkpoints/ is gitignored; released weights live on the Hub at AnonMLuser/remdm-planner-craftax, mirroring the layout below.

Full-Craftax diffusion planner checkpoints are not released: no full-Craftax training run has completed; the released Full Craftax expert can still supervise a new run:

python main.py --mode online \
    --config configs/final_craftax_gpu_24gb.yaml \
    --ppo-checkpoint checkpoints/ppo_agents/Craftax-Symbolic-v1-PPO_RNN-1000M

# All four (~470 MB); narrow the --include glob for a single checkpoint.
uv run hf download AnonMLuser/remdm-planner-craftax --include "checkpoints/**" --local-dir .

Keep the --include. The Hub repo carries its own README.md (the generated model card), LICENSE and .gitattributes; dropping the glob and pulling into --local-dir . overwrites this repository's copies of all three. To fetch everything, add --exclude "README.md" "LICENSE" ".gitattributes", or use a separate --local-dir. Publishing is safe either way — hf_upload.py stages LICENSE and the demo README.md from git, not the working tree.

Experiment outputs

Ablation figures, tables and results.json are regenerated output, so experiments/rl_finetuning/outputs/ and results/inference/ are gitignored. Obtain them either way:

# Fetch the published run (figures, tables, results.json, diagnosis.md)
uv run hf download AnonMLuser/remdm-planner-craftax \
    --include "experiments/rl_finetuning/outputs/**" --local-dir .

# Or regenerate from a checkpoint; writes to outputs/{run_id}/
python experiments/rl_finetuning/run_ablations.py --checkpoint $PRETRAINED_CKPT --all

scripts/hf_upload_demo.py reads outputs/craftax_classic_ablations/{figures,tables} from the working copy, so fetch or regenerate first. demo_craftax.ipynb needs no local copy — it reads them through its own snapshot_download.

Paper figures

Each manuscript figure puts Craftax Classic and MiniHack side by side, so they are built by scripts/paper_figures.py rather than by the single-environment experiments/rl_finetuning/analysis/plots.py. It reads both repositories' results.json and emits vector PDF at NeurIPS column width:

uv run python scripts/paper_figures.py \
    --minihack-results ../remdm-planner-minihack/experiments/rl_finetuning/outputs/minihack_ablations/results.json \
    --outdir results/paper_figures

The MiniHack path defaults to that sibling checkout. Pass --emit-tex-macros to run_ablations.py to also write tables/results.tex, one \newcommand per headline quantity, so the manuscript cites generated numbers instead of retyping them. Macros from this repository are prefixed rw and the sibling suite's mh, so both files can be \input together.

Publishing to the Hub

scripts/hf_upload.py rediscovers and uploads four things, each keeping its repo-relative path: checkpoints/, every experiments/rl_finetuning/outputs/<run>/ holding a results.json (with diagnosis.md, tables/, figures/, gdelta/), the eval JSONs in results/inference/, and the manuscript figure PDFs in results/paper_figures/. It drops W&B and hub config keys, shortens absolute paths and regenerates the model card.

HF_TOKEN=hf_xxx uv run python scripts/hf_upload.py --repo-id <ANON_HF_REPO_ID> --dry-run

--dry-run prints the staged tree and card without uploading; drop it to upload. Also --inference-results <FILE|DIR> ... (eval JSONs kept elsewhere), --private, --yes.

Checkpoint discovery expects the released layout, checkpoints/<role>/<name>/<step>/. A training run writes elsewhere, so copy its wandb.run.dir/policies directory to checkpoints/{offline,online}/<name> first, or nothing is staged. checkpoints/hf/ is skipped — that is where a Hub download lands, and publishing from it would nest already-published artefacts under checkpoints/hf/checkpoints/....

Results, citation, licence

Results tables and the full method description are in Return-Weighted ELBO Fine-Tuning Degrades Masked Diffusion Planners (under submission); demo_craftax.ipynb reproduces the headline evaluation. Citation to be added on publication. Licence: MIT, see LICENSE.

---

Appendix: benchmark-specific detail

Environments

Set via the env_name config key.

Remasking strategies

Selected by remask_strategy, on top of the three-phase loop controlled by use_loop, t_on and t_off.

Key hyperparameters

configs/defaults.yaml is authoritative and commented inline. Tabulated here are the keys that change a result, carry a hazard, or are named elsewhere in this README.

Environment. env_name selects the benchmark: Craftax-Classic-Symbolic-v1 (default) or Craftax-Symbolic-v1 for Full Craftax.

Diffusion model

Transformer architecture. d_model 384, n_heads 8, n_layers 6, d_ff 768, obs_encoder_layers 2, obs_encoder_width 768, dropout_rate 0.1 — the shape every released checkpoint carries. A checkpoint restores only against a matching config.

Offline training

Online DAgger training

Data collection / inference

Checkpointing / logging

The resume_*, use_wandb and wandb_* keys mirror the run flags documented under Configuration.

Persistent compilation cache

The whole training run is one jax.jit, so every process pays one large compilation before any work happens, and multi-seed runs, resumed runs and the ablation suite each repeat it. jax_compilation_cache_dir makes the second and later runs of the same graph skip it. The cache is keyed on the lowered HLO, so a hit is bit-identical to a miss. Point it at local disk, not an NFS home:

python main.py --mode online --ppo-checkpoint <ppo> \
    --config configs/final_craftax_classic_gpu_24gb.yaml \
    --override jax_compilation_cache_dir=/var/tmp/$USER/jax-cache

Environment wrappers

From Craftax_Baselines/wrappers.py (submodule):

Stack (identical for training and inference): env -> LogWrapper -> AutoResetEnvWrapper -> BatchEnvWrapper.

Testing

uv run pytest

A CPU-only suite, 14 modules. Tiny synthetic data and a shrunken model throughout — no real checkpoints, datasets or network calls, and nothing written outside tmp_path. conftest.py forces JAX_PLATFORMS=cpu and disables W&B; there are no custom markers.

Implementation notes