experiments/README.md

ReMDM Experiments

Research and diagnostic scripts for investigating RL fine-tuning of the ReMDM diffusion planner. These scripts are standalone research code — they import from src/ but do not modify it.

---

rl_finetuning/ — RL Fine-Tuning Ablation Suite

Diagnoses why RL fine-tuning of the diffusion model collapses and which interventions fix it. Implements 26 ablations: a baseline plus four groups (A: Regularisation, B: Training Signal, C: Architecture, D: Data Quality), with a comprehensive diagnostic and analysis pipeline.

Training data is on-policy. Each iteration rolls the current model out under its EMA weights (diffusion_steps_collect denoising steps per plan, num_steps // plan_horizon plan cycles) and trains on those windows, weighted by each window's own H-step reward sum. The suite therefore needs no expert: --ppo-checkpoint is gone, and only the pretrained diffusion --checkpoint is required.

Directory structure

rl_finetuning/
├── run_ablations.py          # CLI entry point
├── ablations/
│   ├── losses.py             # All loss/objective variants as factory functions
│   ├── optimizers.py         # LLRD, LoRA, gradient surgery, param masking
│   ├── registry.py           # AblationSpec dataclass + REGISTRY (26 ablations)
│   └── training.py           # make_run_ablation() factory + AblationHistory dataclass
├── diagnostics/
│   ├── gradient.py           # Grad alignment, per-layer norms, surgery metrics
│   ├── representation.py     # KL drift, CKA similarity, activation norms
│   └── timestep.py           # t-bin gradient norms, per-t loss decomposition
├── analysis/
│   ├── gdelta.py             # Return term g_delta of the decomposition (no training)
│   ├── plots.py              # 16 matplotlib figure generators
│   ├── tables.py             # Summary tables as polars DataFrames + LaTeX export
│   ├── report.py             # diagnosis.md + decision tree figure
│   └── action_distribution.py  # Pre/post-RL action distribution analysis
└── configs/
    ├── ablations_default.yaml   # Base: all ablation hyperparameters
    ├── ablations_fast.yaml      # Smoke-test overlay (50 iterations, 16 envs)
    ├── ablations_final_craftax_classic_gpu_24gb.yaml   # RTX 3090 Ti, seed 42 (reference)
    ├── ablations_final_craftax_classic_gpu_h200.yaml   # H200, seed 43
    ├── ablations_final_craftax_gpu_24gb.yaml           # Full Craftax, 24GB, seed 42
    └── ablations_final_craftax_gpu_h200.yaml           # Full Craftax, H200, seed 43

Each ablations_final_* matches the configs/final_* of the same name.

ablations_default.yaml carries the released checkpoints' architecture — 384-dim, 8 heads, 6 layers, d_ff 768, plan_horizon 32, the same for Classic and Full — so the ablations_final_* presets need not restate it. A run against a differently-shaped checkpoint must override those keys or the model build fails on a shape mismatch.

Config layering

Lowest to highest:

configs/defaults.yaml -> ablations_default.yaml -> machine config -> ablations_fast.yaml (--fast only) -> CLI flags

Any file given to --ablations-config layers on top of ablations_default.yaml automatically, so the ablations_final_* presets carry only their own deltas. An ablations config never inherits from another ablations config.

ablations_fast.yaml is deliberately not layered that way: --fast reads it raw and overlays it last.

Presets hold only deltas, never restate a default — tests/test_config.py enforces it.

The two craftax presets each restate the three keys where Full Craftax departs from the Classic base (env_name, val_diffusion_steps, temperature); with no inheritance between configs there is nowhere shared to put them, so a change to those must be made in both files.

Compilation cache

The graph is identical across the seeds of one ablation — only the PRNG key differs, and that is a runtime argument — so at num_seeds: 3 two runs in three are a cache hit, as are reruns and the per-GPU processes of --merge. Off unless jax_compilation_cache_dir is set, and run_ablations.py has no --override, so it must be set in configs/defaults.yaml. Point it at local disk, not an NFS home:

# configs/defaults.yaml
jax_compilation_cache_dir: /var/tmp/your-user/jax-cache

Usage

--checkpoint takes the pretrained diffusion checkpoint from either --mode offline or --mode online; for DAgger, the final ({env}-policy) or best-validation ({env}-policy-best) artifact is consumed directly. It also accepts wandb:team/project/artifact:latest references, downloaded automatically before training begins.

List all ablations:

python experiments/rl_finetuning/run_ablations.py --list

Smoke test (2 ablations, fast config):

python experiments/rl_finetuning/run_ablations.py \
    --ablations baseline_rl kl_penalty \
    --fast \
    --checkpoint $PRETRAINED_CKPT

Full suite (all 26 ablations):

python experiments/rl_finetuning/run_ablations.py \
    --config configs/defaults.yaml \
    --ablations-config experiments/rl_finetuning/configs/ablations_default.yaml \
    --all \
    --num-seeds 3 \
    --checkpoint $PRETRAINED_CKPT \
    --use-wandb

Full suite against a pinned final_* checkpoint:

# Craftax Classic, GPU-24GB hardware (seed 42 checkpoint)
python experiments/rl_finetuning/run_ablations.py \
    --ablations-config experiments/rl_finetuning/configs/ablations_final_craftax_classic_gpu_24gb.yaml \
    --all --num-seeds 3 \
    --checkpoint wandb:my-team/remdm-planner-craftax/Craftax-Classic-Symbolic-v1-policy-best:latest \
    --use-wandb

Specific ablations:

python experiments/rl_finetuning/run_ablations.py \
    --ablations ewc lora gradient_surgery trust_region_kl \
    --checkpoint $PRETRAINED_CKPT

Re-plot from saved results (no training):

python experiments/rl_finetuning/run_ablations.py \
    --analyze-only \
    --results-path experiments/rl_finetuning/outputs/run_20250101_120000/results.json

Merge multi-GPU results:

python experiments/rl_finetuning/run_ablations.py \
    --merge outputs/gpu0/results.json outputs/gpu1/results.json \
    --output-dir experiments/rl_finetuning/outputs/merged/

--merge only pools runs from configs that agree on result-affecting keys. It compares the configs the results files recorded and refuses, naming every diverging key with both values; a file that records no config is refused too. Each family's GPU-24GB config is its reference, and the GPU-H200 sibling is not poolable with it:

Key	Classic GPU-24GB	Classic GPU-H200	Craftax GPU-24GB	Craftax GPU-H200	Effect
num_envs	192	64	128	64	rollout diversity per iteration
batch_size	1024	256	1024	512	per-update SNR
eval_steps	1024	512	1024	512	noisier score
mixed_replay_buffer_size	20000	10000	10000	10000	replay horizon

Values above are post-layering: a config's own value, or what it inherits from ablations_default.yaml. Differences in diagnostic cadence (eval_every, cka_every, cka_batch_size, per_layer_every, repr_drift_every, grad_align_every, t_analysis_every) are wall-clock only and do not affect poolability.

tests/test_config.py enforces this: every ablations_final_*.yaml must be declared poolable or not, configs declared poolable must match their family reference on the result-affecting keys, and the recorded GPU-H200 divergences must stay accurate. Aligning a GPU-H200 config later fails the test until it is moved to the poolable set. The key set itself is declared once, in run_ablations._RESULT_AFFECTING, so the classification the tests check and the refusal --merge performs are the same policy.

Ablations

Group	Name	Tests
Baseline	baseline_rl	Standard return-weighted ELBO
A: Regularisation	kl_penalty	Soft KL constraint vs. pretrained
	ewc	Elastic Weight Consolidation (Fisher diagonal)
	llrd	Layer-wise Learning Rate Decay
	lora	Low-Rank Adaptation of attention projections
	mixed_replay	Self-replay: the run's own past online windows resampled into each batch
	trust_region_kl	Hard KL trust region via quadratic barrier
B: Training Signal	t_curriculum	Anneal t range high→low over training
	entropy_bonus	Entropy regularisation for action diversity
	gradient_surgery	PCGrad: project conflicting RL/BC gradients
	advantage_clip	PPO-style advantage clipping [1-ε, 1+ε]
	normalized_adv	Std-normalised advantages
	bc_wins	Uniform ELBO on win windows (no advantage weighting)
	bc_all	Uniform ELBO on all rollout windows (no advantage weighting)
	low_t	ELBO restricted to low-t (fine-detail) regime
C: Architecture	frozen_backbone	Train the action head + token embeddings (backbone frozen)
	head_only	Train only the final action projection
	attention_only	Train only the attention projections (Q/K/V/O)
	ffn_only	Train only the per-block FFN layers
	layer_ablation_top1	Train only the top-1 transformer block + head
	layer_ablation_top2	Train only the top-2 transformer blocks + head
	layer_ablation_top3	Train only the top-3 transformer blocks + head
D: Data Quality	reward_filtering	Top-75th-percentile return windows only
	running_stats	EMA running mean/std for advantage normalisation
	action_diversity	Discard degenerate (all-same-action) plans
	reward_model	MLP reward model soft-weighting of advantages

Output structure

experiments/rl_finetuning/outputs/{run_id}/
├── results.json               # All histories + final scores (machine-readable; see schema below)
├── diagnosis.md               # Human-readable verdict + evidence + recommendations
├── checkpoint_{name}/         # Per-ablation fine-tuned params, last seed (Orbax)
├── figures/
│   ├── curves_{name}.png                  # Per-ablation training curves (2×3 grid)
│   ├── final_score_comparison.png
│   ├── eval_scores_over_training.png
│   ├── score_delta_over_baseline_rl.png
│   ├── gradient_alignment.png
│   ├── gradient_conflict_map.png
│   ├── per_layer_grad_heatmap_{name}.png
│   ├── representation_drift.png
│   ├── cka_similarity.png
│   ├── t_distribution_analysis.png
│   ├── t_bin_grad_norms_{name}.png
│   ├── t_bin_norms_heatmap.png            # Per-t-bin gradient norms, final iteration
│   ├── group_comparison.png               # Boxplot of scores by ablation group
│   ├── win_rate_and_effective_batch_size.png
│   ├── achievement_breakdown.png          # Start vs end achievement rates (stacked bars)
│   ├── achievement_collapse_{name}.png    # Per-ablation achievement heatmap over time
│   ├── diagnosis_decision_tree.png
│   └── action_dist/                        # On by default; --no-action-dist skips it
│       ├── action_freq_{name}.png         # Side-by-side pre/post action frequency bars
│       ├── transition_matrix_{name}.png   # 3-panel heatmap (pre, post, difference)
│       ├── action_metrics_{name}.png      # 2x2 dashboard (entropy in nats, effective, Gini, divergences)
│       └── js_divergence_comparison.png   # Cross-ablation JS divergence bar chart
├── gdelta/                    # --measure-gdelta only
│   ├── gdelta_seed{n}.json    # Per rollout seed; +/- within is across that seed's draws
│   └── gdelta_aggregate.json  # Across seeds; the dispersion the paper's table prints
└── tables/
    ├── main_results.{csv,tex}
    ├── significance_test.txt              # Max-statistic permutation test + p floor + bootstrap CI
    ├── group_summary.{csv,tex}            # Group-level summary table
    ├── gradient_analysis.{csv,tex}
    ├── t_distribution.{csv,tex}
    ├── repr_drift.{csv,tex}               # KL drift values at the final iteration
    ├── per_env.{csv,tex}                  # Per-achievement rates; needs pretrained_ach_rates
    ├── forgetting_analysis.{csv,tex}
    ├── hypothesis_verdict.{csv,tex}
    ├── achievement_summary.{csv,tex}      # Per-achievement final unlock rates
    ├── gdelta.{csv,tex}                   # --measure-gdelta only: the decomposition per weight transform
    └── results.tex                        # --emit-tex-macros only: \newcommand per headline number

Action distribution analysis runs by default and is disabled with --no-action-dist. The rollout is one vectorised scan over num_envs, sized from the config rather than by an episode count, so it costs a fraction of a training run — which is why the default differs from the minihack twin, where the same flag defaults off because MiniHack rollouts are not vectorised. It reads each ablation's final_params, so it only covers ablations that completed.

results.json schema:

{
  "pretrained_score": 0.1234,
  "pretrained_ach_rates": {"achievement_collect_wood": 0.42, ...},
  "config": {"MAX_ITER": 1000, ...},   // the merged config, keys uppercase
  "merge_provenance": { ... },         // --merge only: inputs + which supplied config
  "ablations": {
    "kl_penalty": {
      "score": 0.1456,        // mean across seeds
      "score_std": 0.008,     // std across seeds (0.0 if num_seeds=1)
      "all_scores": [0.1456], // per-seed scores
      "base_seed": 42, "seeds": [42], // seeding actually used
      "wall_clock_s": 812.4,
      "per_seed_finals": [{...}],     // per-seed end-of-run metrics
      "final_ach_rates": {"achievement_collect_wood": 0.40, ...},
                              // achievement detail of the same post-loop
                              // evaluations that produced all_scores, seed-averaged
      "all_final_ach_rates": [{...}],   // per-seed, before averaging
      "history": { ... }      // AblationHistory serialised
    }
  }
}

results.json is written incrementally after each ablation completes — a partial file with N of 26 ablations is fully valid and loadable by --analyze-only --results-path.

Measuring the return term (--measure-gdelta)

Splits the return-weighted ELBO gradient into an imitation term and a return term at a single parameter point:

grad L_RW  =  Abar * ( grad L_BC + g_delta ),
g_delta    =  (1/B) sum_i delta_i grad l_i,    delta_i = A_i/Abar - 1.

Loads the pretrained checkpoint, collects one on-policy batch from it, and evaluates grad L_BC, g_delta and grad L_RW on that batch at those parameters under a shared (z_t, t) draw, so the only difference between the three is the weight vector. It repeats for the four weighting ablations (baseline_rl, advantage_clip, normalized_adv, bc_wins) and reports CV_A, Abar, ESS/B, the norm ratio and the cosine, plus a shuffled-delta null that keeps the weight multiset and destroys its association with each window's return. No training and no optimiser step occur; it runs on a laptop CPU.

Results land in gdelta/ under the run's own output directory, beside results.json, and the aggregate additionally produces tables/gdelta.{csv,tex}. With --emit-tex-macros, the analysis pass picks the aggregate up and emits the measured quantities as \rwGdelta* macros. Those are kept separate from the \rwCvA* macros, which recover CV_A from the ESS logged during training: the two are measured on different batches and do not agree.

Config comes from --results-path, so the weight transforms measured are the ones that run trained under; without it the standard layering applies.

Reproduction (three rollout seeds, aggregated in one pass):

python experiments/rl_finetuning/run_ablations.py --measure-gdelta --gdelta-seeds 0 1 2 \
    --checkpoint checkpoints/online/Craftax-Classic-Symbolic-v1-Online-Diffusion-DAgger-100M \
    --results-path experiments/rl_finetuning/outputs/craftax_classic_ablations/results.json \
    --output-dir experiments/rl_finetuning/outputs/craftax_classic_ablations

Seeds run on separate machines are aggregated afterwards with --gdelta-inputs, the counterpart to --merge:

python experiments/rl_finetuning/run_ablations.py --run-id gdelta \
    --gdelta-inputs experiments/rl_finetuning/outputs/gdelta/gdelta_seed{0,1,2}.json

A single seed's ratio_std_draws / cos_std_draws are dispersions over that seed's eight (z_t, t) draws. The aggregate averages the per-seed means and reports the standard deviation across seeds, which is what the paper's table prints.

CLI reference

Flag	Description
--checkpoint PATH	Pretrained diffusion checkpoint, offline or DAgger (path or wandb: artifact)
--config PATH	Main pipeline config (default: configs/defaults.yaml)
--ablations-config PATH	Ablations config, layered on ablations_default.yaml (default: ablations_default.yaml)
--all	Run all 26 ablations
--ablations NAME [NAME ...]	Run specific ablations by name
--list	Print registered ablations and exit
--fast	Smoke-test overlay: ablations_fast.yaml applied last (max_iter 50, num_envs 16)
--num-seeds N	Seeds per ablation (overrides num_seeds, default 3)
--seed N	Base random seed
--output-dir DIR	Root output directory (default: outputs/{run_id}/)
--run-id ID	Run identifier (default: run_{timestamp})
--analyze-only	Skip training, regenerate analysis from existing results
--results-path PATH	Explicit path to results.json (with --analyze-only or --measure-gdelta)
--merge PATH [PATH ...]	Merge multiple results.json files and regenerate analysis
--measure-gdelta	Measure the return term at the pretrained checkpoint; no training
--gdelta-seeds N [N ...]	Rollout seeds to measure (default 0); the reported +/- is across these
--gdelta-draws N	Independent (z_t, t) draws per seed (default 8)
--gdelta-inputs PATH [PATH ...]	Aggregate per-seed gdelta JSONs from separate machines
--emit-tex-macros	Also write tables/results.tex, one \newcommand per headline number
--action-dist / --no-action-dist	Pre/post action-distribution analysis (default on here; off in the minihack twin)
--use-wandb / --no-use-wandb	Enable/disable W&B logging (overrides use_wandb, default false)
--wandb-project NAME	W&B project (default remdm-planner-craftax-ablations)
--wandb-entity NAME	W&B entity
--max-iter N	Override max training iterations
--num-envs N	Override rollout environments per iteration
--batch-size N	Override batch size
--eval-every N	Override evaluation frequency
--lr FLOAT	Override learning rate

There is no --override: keys that are not flags are set in the config files.

W&B logging

Three metrics per ablation, logged under ablations/{name}/ against iteration: train_loss, env_score (both every logged iteration) and eval_score (every eval_every). They are written after the arm finishes, not during it, and there is no wandb.summary write.

Every other quantity in the table below — gradient alignment, per-layer norms, KL drift, CKA, the t-bin norms — is collected into AblationHistory and reaches results.json only. Read those from the run directory, not from W&B.

Diagnostic metrics collected

Metric	Frequency	What it answers
Eval score	every eval_every iters	Primary performance
Training loss	every 10 iters	Optimisation signal
Env score	every 10 iters	Online rollout quality
Gradient alignment (cos sim)	every grad_align_every	Is the RL gradient useful?
Per-layer gradient norms	every per_layer_every	Which layers collapse?
KL drift from pretrained	every repr_drift_every	How much has the model changed?
CKA similarity	every cka_every	Representational drift (activation level)
t-bin gradient norms	every t_analysis_every	Is high-t gradient biased?
Win rate	every 10 iters	Signal sparsity
Effective batch size	every 10 iters	Gradient concentration
Gradient surgery fraction	every grad_align_every	PCGrad projected mass
Action dist JS divergence	post-training	Mode collapse vs drift?
Action dist KL / TV	post-training	Magnitude of behavioural shift
Action transition matrix	post-training	Bigram structure change