Refresh artefacts and code for the second review release

.gitattributes

@@ -33,237 +33,3 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
-checkpoints/offline/Craftax-Classic-Symbolic-v1-Offline-Diffusion-BC-100M/99942400/default/ocdbt.process_0/d/503349be2d11cb1bf04a17e7b3e2d1e0 filter=lfs diff=lfs merge=lfs -text
-checkpoints/offline/Craftax-Classic-Symbolic-v1-Offline-Diffusion-BC-100M/99942400/default/ocdbt.process_0/d/781b85e3ec447139f7d23160643a9c21 filter=lfs diff=lfs merge=lfs -text
-checkpoints/offline/Craftax-Classic-Symbolic-v1-Offline-Diffusion-BC-100M/99942400/default/ocdbt.process_0/d/92adc334c11103c7741baed0d71e66cb filter=lfs diff=lfs merge=lfs -text
-checkpoints/online/Craftax-Classic-Symbolic-v1-Online-Diffusion-DAgger-100M/100000000/default/ocdbt.process_0/d/35a15297533c9b116fa22b2bb1975990 filter=lfs diff=lfs merge=lfs -text
-checkpoints/online/Craftax-Classic-Symbolic-v1-Online-Diffusion-DAgger-100M/100000000/default/ocdbt.process_0/d/9778194f691de46821092616cc3a1ff2 filter=lfs diff=lfs merge=lfs -text
-checkpoints/ppo_agents/Craftax-Classic-Symbolic-v1-PPO_RNN-1000M/1000000000/default/ocdbt.process_0/d/369457b7c6608f1adf28eb88024d6b91 filter=lfs diff=lfs merge=lfs -text
-checkpoints/ppo_agents/Craftax-Classic-Symbolic-v1-PPO_RNN-1000M/1000000000/default/ocdbt.process_0/d/6bb110a840637eede93e25f5356236f9 filter=lfs diff=lfs merge=lfs -text
-checkpoints/ppo_agents/Craftax-Symbolic-v1-PPO_RNN-1000M/1000000000/default/ocdbt.process_0/d/e5a4020f50167115120fe5dac41c20fb filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_breakdown.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_action_diversity.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_advantage_clip.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_attention_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_baseline_rl.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_bc_wins.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_entropy_bonus.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_ewc.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_ffn_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_frozen_backbone.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_gradient_surgery.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_head_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_kl_penalty.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_layer_ablation_top1.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_layer_ablation_top2.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_layer_ablation_top3.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_llrd.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_lora.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_low_t.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_mixed_replay.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_normalized_adv.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_reward_filtering.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_reward_model.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_running_stats.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_t_curriculum.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/achievement_collapse_trust_region_kl.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/action_dist/js_divergence_comparison.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/cka_similarity.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_action_diversity.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_advantage_clip.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_attention_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_baseline_rl.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_bc_wins.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_entropy_bonus.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_ewc.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_ffn_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_frozen_backbone.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_gradient_surgery.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_head_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_kl_penalty.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_layer_ablation_top1.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_layer_ablation_top2.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_layer_ablation_top3.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_llrd.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_lora.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_low_t.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_mixed_replay.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_normalized_adv.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_reward_filtering.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_reward_model.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_running_stats.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_t_curriculum.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/curves_trust_region_kl.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/eval_scores_over_training.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/final_score_comparison.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/gradient_alignment.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/gradient_conflict_map.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_action_diversity.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_advantage_clip.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_attention_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_baseline_rl.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_bc_wins.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_entropy_bonus.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_ewc.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_ffn_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_frozen_backbone.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_gradient_surgery.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_head_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_kl_penalty.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_layer_ablation_top1.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_layer_ablation_top2.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_layer_ablation_top3.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_llrd.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_lora.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_low_t.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_mixed_replay.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_normalized_adv.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_reward_filtering.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_reward_model.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_running_stats.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_t_curriculum.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/per_layer_grad_heatmap_trust_region_kl.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/representation_drift.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/score_delta_over_baseline_rl.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_action_diversity.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_advantage_clip.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_attention_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_baseline_rl.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_bc_wins.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_entropy_bonus.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_ewc.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_ffn_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_frozen_backbone.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_gradient_surgery.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_head_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_kl_penalty.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_layer_ablation_top1.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_layer_ablation_top2.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_layer_ablation_top3.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_llrd.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_lora.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_low_t.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_mixed_replay.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_normalized_adv.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_reward_filtering.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_reward_model.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_running_stats.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_t_curriculum.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_grad_norms_trust_region_kl.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_bin_norms_heatmap.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/t_distribution_analysis.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_final_results/analysis/figures/win_rate_and_effective_batch_size.png filter=lfs diff=lfs merge=lfs -text
-checkpoints/offline/Craftax-Classic-Symbolic-v1-Offline-Diffusion-BC-100M/99942400/default/ocdbt.process_0/d/1b7ef41abe2bf7d133b20680543c1738 filter=lfs diff=lfs merge=lfs -text
-checkpoints/offline/Craftax-Classic-Symbolic-v1-Offline-Diffusion-BC-100M/99942400/default/ocdbt.process_0/d/8f77d89a5c0896490dcc8e0a7a3e001f filter=lfs diff=lfs merge=lfs -text
-checkpoints/online/Craftax-Classic-Symbolic-v1-Online-Diffusion-DAgger-100M/40370176/default/ocdbt.process_0/d/8bf3ebb1b364b9527375993e4a94621d filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_breakdown.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_action_diversity.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_advantage_clip.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_attention_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_baseline_rl.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_bc_wins.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_entropy_bonus.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_ewc.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_ffn_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_frozen_backbone.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_gradient_surgery.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_head_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_kl_penalty.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_layer_ablation_top1.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_layer_ablation_top2.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_layer_ablation_top3.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_llrd.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_lora.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_low_t.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_mixed_replay.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_normalized_adv.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_reward_filtering.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_reward_model.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_running_stats.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_t_curriculum.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/achievement_collapse_trust_region_kl.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/cka_similarity.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_action_diversity.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_advantage_clip.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_attention_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_baseline_rl.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_bc_wins.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_entropy_bonus.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_ewc.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_ffn_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_frozen_backbone.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_gradient_surgery.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_head_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_kl_penalty.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_layer_ablation_top1.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_layer_ablation_top2.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_layer_ablation_top3.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_llrd.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_lora.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_low_t.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_mixed_replay.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_normalized_adv.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_reward_filtering.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_reward_model.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_running_stats.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_t_curriculum.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/curves_trust_region_kl.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/eval_scores_over_training.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/final_score_comparison.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/gradient_alignment.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/gradient_conflict_map.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_action_diversity.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_advantage_clip.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_attention_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_baseline_rl.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_bc_wins.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_entropy_bonus.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_ewc.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_ffn_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_frozen_backbone.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_gradient_surgery.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_head_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_kl_penalty.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_layer_ablation_top1.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_layer_ablation_top2.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_layer_ablation_top3.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_llrd.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_lora.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_low_t.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_mixed_replay.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_normalized_adv.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_reward_filtering.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_reward_model.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_running_stats.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_t_curriculum.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/per_layer_grad_heatmap_trust_region_kl.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/representation_drift.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/score_delta_over_baseline_rl.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_action_diversity.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_advantage_clip.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_attention_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_baseline_rl.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_bc_wins.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_entropy_bonus.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_ewc.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_ffn_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_frozen_backbone.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_gradient_surgery.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_head_only.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_kl_penalty.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_layer_ablation_top1.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_layer_ablation_top2.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_layer_ablation_top3.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_llrd.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_lora.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_low_t.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_mixed_replay.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_normalized_adv.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_reward_filtering.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_reward_model.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_running_stats.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_t_curriculum.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_grad_norms_trust_region_kl.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_bin_norms_heatmap.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/t_distribution_analysis.png filter=lfs diff=lfs merge=lfs -text
-experiments/rl_finetuning/outputs/craftax_classic_ablations/figures/win_rate_and_effective_batch_size.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -246,3 +246,4 @@ CLI_CHANGES.md
 
 # Claude Code: per-machine settings stay local; the rest of .claude/ is shared
 .claude/settings.local.json
+results/

Craftax_Baselines/Dockerfile

@@ -33,9 +33,9 @@ USER duser
 WORKDIR /home/duser
 
 # Install Python packages
-ENV PATH="/home/duser/.local/bin:$PATH"
+ENV PATH="<path omitted for anonymous review>"
 RUN python3 -m pip install --upgrade pip
 ARG REQS
 RUN pip install $REQS -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
-WORKDIR /home/duser/Craftax
+WORKDIR <path omitted for anonymous review>

Craftax_Baselines/run_docker.sh

@@ -14,7 +14,7 @@ for gpu in $gpus; do
     docker run \
         --gpus device=$gpu \
         -e WANDB_API_KEY=$WANDB_API_KEY \
-        -v $(pwd):/home/duser/Craftax \
+        -v $(pwd):<path omitted for anonymous review> \
         --name craftax_$gpu \
         --user $(id -u) \
         --rm \

Dockerfile

@@ -1,4 +1,4 @@
-# craftax: vast.ai image.
+# remdm-planner-craftax: vast.ai image.
 #
 # Build and push (amd64 is mandatory: vast GPU hosts are x86, an arm64 image
 # from an Apple Silicon Mac will not run):
@@ -59,7 +59,12 @@ RUN cp uv.lock /opt/venv/baked-uv.lock
 WORKDIR /workspace
 
 # --- optional: bake the released checkpoints ------------------------------
-# Left out on purpose. 470 MB unpacks into a large small-file layer, which is
-# the slowest kind to pull and a likely cause of stalled image pulls. Mount or
-# copy a local checkpoints/ directory into /opt/checkpoints instead. Released
-# weights are withheld for anonymous review; see README.md.
+# Left out on purpose. 470 MB on the Hub unpacks into a large small-file layer,
+# which is the slowest kind to pull and a likely cause of stalled image pulls.
+# Fetching them at boot takes under a minute. If you do want them baked, use an
+# isolated env so the pinned huggingface_hub in /opt/venv is not disturbed:
+#
+# RUN HF_HOME=/tmp/hf uvx --from "huggingface_hub[cli]" \
+#         hf download AnonMLuser/remdm-planner-craftax \
+#         --include "checkpoints/**" --local-dir /opt/checkpoints \
+#     && rm -rf /tmp/hf

LICENSE

@@ -1,6 +1,7 @@
 MIT License
 
-Copyright (c) 2026 The authors of "Return-Weighted ELBO Fine-Tuning of Discrete Diffusion Planners Reduces to Self-Imitation"
+Copyright (c) 2026 The authors of "Return-Weighted ELBO Fine-Tuning Degrades
+Masked Diffusion Planners"
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

README.md

@@ -38,13 +38,24 @@ checkpoints/
                                                                    ablation fine-tunes)
   ppo_agents/                                                      PPO-RNN experts
 
-experiments/rl_finetuning/outputs/craftax_classic_ablations/
+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, ...
-  results.json
-
-results/inference/   Single-seed evaluations through the inpainting sampler (see below)
+             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)
@@ -59,6 +70,39 @@ results/inference/   Single-seed evaluations through the inpainting sampler (see
 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.
+
+| Arm | lr | Score |
+|---|---|---|
+| `baseline_rl` — same-stack anchor | 3e-4 | 8.3395 ± 0.1468 |
+| `bc_all` — uniform weights, all rollout windows | 3e-4 | **4.7348 ± 0.1869** |
+| `advantage_clip` — matched effective step | 3.223e-4 | **4.9246 ± 0.2573** |
+| `baseline_rl` | 1e-4 | **10.1486 ± 0.1811** |
+| `baseline_rl` | 1e-5 | **10.8104 ± 0.1024** |
+
+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**
@@ -66,12 +110,14 @@ Every number in the paper comes from the **ablation harness**
 `sample_plan` with no locked prefix, 50 denoising steps, 8 actions executed per
 plan, scored as `returned_episode_returns`.
 
-`main.py --mode inference` takes 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.
+`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).
@@ -103,13 +149,13 @@ Then open `demo_craftax.ipynb`, or see the project documentation below.
 
 JAX implementation of **ReMDM** (Remasking Discrete Diffusion Model) for action-sequence planning in [Craftax](https://github.com/MichaelTMatthews/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 [`minihack/`](../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.
+The sibling repository [`remdm-planner-minihack`](../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 in the accompanying paper (under submission; citation to follow), both supervised by a pre-trained PPO expert:
+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.
 
-One PPO expert checkpoint feeds both pipelines: `--mode offline` behaviour-clones from live expert rollouts, `--mode online` runs DAgger from scratch against expert labels. Either output is scored with `--mode inference`.
+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
 
@@ -117,7 +163,9 @@ Prerequisites: Python 3.12+, [uv](https://docs.astral.sh/uv/). Linux GPU use nee
 
 ```bash
 git clone https://github.com/ANONYMOUS/remdm-planners.git
-cd remdm-planners/craftax
+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.
@@ -135,16 +183,17 @@ Extras: `cuda13` and `cuda12` are mutually exclusive and Linux-only. JAX ships G
 ## Repo layout
 
 ```
-craftax/
-├── Craftax_Baselines/       Vendored PPO expert training and env wrappers (upstream-derived)
+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, HF upload utilities
+├── 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/inference/       Eval JSONs from --mode inference (published, 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
@@ -162,7 +211,7 @@ python main.py --mode smoke
 
 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 (`Craftax_Baselines/`)
+### Stage 1 — Train the PPO expert (submodule)
 
 ```bash
 cd Craftax_Baselines
@@ -179,8 +228,6 @@ Rolls out the PPO agent live at each update.
 
 ```bash
 python main.py --mode offline --ppo-checkpoint /path/to/ppo_checkpoint
-python main.py --mode offline --ppo-checkpoint /path/to/ppo_checkpoint \
-    --override offline_total_timesteps=100000000
 ```
 
 ### Online DAgger
@@ -189,8 +236,6 @@ Trained from scratch. Per iteration a mixed expert/learner policy rolls out, the
 
 ```bash
 python main.py --mode online --ppo-checkpoint /path/to/ppo_checkpoint
-python main.py --mode online --ppo-checkpoint /path/to/ppo_checkpoint \
-    --override online_total_timesteps=100000000
 
 # Optional: warm-start from a pre-trained offline checkpoint
 python main.py --mode online --ppo-checkpoint /path/to/ppo_checkpoint \
@@ -212,36 +257,43 @@ python main.py --mode collect --ppo-checkpoint /path/to/ppo_checkpoint \
 ### Resuming a training run
 
 ```bash
-# Offline. --resume also accepts a wandb: artifact reference.
+# 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
-
-# Online. --resume-step / --resume-wandb-run-id override the metadata sidecar.
-python main.py --mode online --ppo-checkpoint /path/to/ppo_checkpoint \
-    --resume /path/to/completed_online_checkpoint \
-    --override online_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`, with the step counter offset so the LR resumes exactly where it stopped. With a metadata sidecar, `resume_step` and `resume_wandb_run_id` are auto-detected; without one, pass `--resume-step`.
+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 on the optimiser-state structure, and there is no compatibility path. Use `--checkpoint` instead — parameters only, warm-starting a fresh run.
+`--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
 
 ```bash
 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:
+
+| Reported | JSON key | Meaning |
+|---|---|---|
+| Mean return, completed episodes | `mean_return_completed_episodes` (with `n_completed_episodes`) | Mean over every episode that terminated inside the rollout — the `returned_episode_returns` statistic the ablation tables and the paper report |
+| Mean return, first life only | `mean_return_first_life`, and `mean_score` for backwards compatibility | Strict single-life return: the first episode of each env only. A harsher statistic |
 
-# Released checkpoints need their matching config (see Checkpoints)
+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:
+
+```bash
 python main.py --mode inference \
     --config configs/final_craftax_classic_gpu_24gb.yaml \
     --checkpoint checkpoints/online/Craftax-Classic-Symbolic-v1-Online-Diffusion-DAgger-100M
 ```
 
-Prints mean episode return, completed episodes, steps per second, and per-achievement unlock counts; `--output` also writes them as JSON. Uses historical inpainting: the first `hist_len` plan positions are locked to observed history. Evaluation length is set by the `eval_steps` / `eval_num_envs` config keys.
-
-Write eval JSONs into `results/inference/` (created for you): `scripts/hf_upload.py` publishes every JSON it finds there.
-
 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/`).
 
 ```bash
@@ -253,7 +305,7 @@ python main.py --mode inference \
 
 ### RL baselines
 
-PPO baselines (the expert family: `ppo`, `ppo_rnn`, `ppo_rnd`) train in the vendored `Craftax_Baselines` tree, see [Training](#stage-1--train-the-ppo-expert-craftax_baselines). Evaluate an expert with `scripts/eval_ppo_expert.py`:
+PPO baselines (the expert family: `ppo`, `ppo_rnn`, `ppo_rnd`) train in the `Craftax_Baselines` submodule, see [Training](#stage-1--train-the-ppo-expert-submodule). Evaluate an expert with `scripts/eval_ppo_expert.py`:
 
 ```bash
 uv run python scripts/eval_ppo_expert.py \
@@ -275,7 +327,7 @@ python main.py --mode online --ppo-checkpoint <ppo> --config configs/classic_exp
 
 ### RL fine-tuning ablation suite
 
-25 registered ablations (same names as in the minihack repo). See `experiments/README.md`.
+26 registered ablations (same names as in the minihack repo). See `experiments/README.md`.
 
 ```bash
 python experiments/rl_finetuning/run_ablations.py --list
@@ -286,26 +338,33 @@ python experiments/rl_finetuning/run_ablations.py \
     --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:
+
+```bash
+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.
 
-- **Config files** (`configs/*.yaml`): hyperparameters, model and method settings, ablation definitions. Any file passed via `--config` is merged onto `configs/defaults.yaml`. Exactly two layers: a preset never inherits from another preset.
-- **CLI flags**: per-invocation values — `--seed`, `--checkpoint`, `--ppo-checkpoint`, `--data`, `--output`, `--resume*`, `--jit/--no-jit` (disable JIT for debugging).
-- **`--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` preset < `--override` and run flags. Exactly two config layers — a preset never inherits from another preset.
 
-Precedence, lowest to highest: `configs/defaults.yaml` < `--config` file < `--override` and run flags.
+- **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, never restate a value they would inherit.** A key belongs in a preset only if its value differs from `defaults.yaml`. Restating one is not harmless duplication: it silently pins the preset when the recipe later moves. `tests/test_config.py` enforces this.
+**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.
 
 ```bash
-python main.py --mode offline --ppo-checkpoint <ppo> \
-    --override lr=1e-4 --override plan_horizon=64 --override num_minibatches=16
-python main.py --mode offline --ppo-checkpoint <ppo> --no-jit --override num_envs=4
+python main.py --mode offline --ppo-checkpoint <ppo> --no-jit \
+    --override lr=1e-4 --override plan_horizon=64 --override num_envs=4
 ```
 
 | Preset | Purpose |
@@ -318,29 +377,19 @@ python main.py --mode offline --ppo-checkpoint <ppo> --no-jit --override num_env
 | `configs/classic_exp_d_{100K,250K,850K,3M}_model.yaml` | Craftax Classic model-size scaling sweep |
 | `configs/craftax_exp_d_{500K,1M,3M,7M}_model.yaml` | Full Craftax model-size scaling sweep |
 | `configs/final_craftax_classic_{gpu_h200,gpu_24gb}.yaml` | Final Classic DAgger — `num_envs` and `seed` only; the recipe is `defaults.yaml` |
-| `configs/final_craftax_{gpu_h200,gpu_24gb}.yaml` | Final Full Craftax DAgger — the 11 keys where Full Craftax departs from the Classic recipe, plus `num_envs` and `seed` |
+| `configs/final_craftax_{gpu_h200,gpu_24gb}.yaml` | Final Full Craftax DAgger — the 8 keys where Full Craftax departs from the Classic recipe, plus `num_envs` and `seed` |
 
-Within each family the two cluster configs differ only in `num_envs` and `seed`. Nothing in the loader enforces that: the guard is `test_cluster_siblings_differ_only_in_num_envs_and_seed`. **A Full Craftax hyperparameter change must be made in both `final_craftax_*` files**, since with no inheritance those 11 keys are duplicated verbatim in each; a Classic one belongs in `defaults.yaml`.
+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`.
 
-Fairness-critical values are env-frame denominated (the six keys above) and rescaled by `resolve_scaled_hyperparams()` at load, so one recipe runs on any hardware tier. Key hyperparameters are documented inline in `configs/defaults.yaml`; the [appendix](#key-hyperparameters) tabulates them. Ablation-suite hyperparameters live in `experiments/rl_finetuning/configs/`, loaded by `run_ablations.py`, not `main.py`.
+Key hyperparameters are documented inline in `configs/defaults.yaml`; the [appendix](#key-hyperparameters) 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). With W&B on these sit under `wandb.run.dir` and are uploaded as W&B artifacts named `{env_name}-policy` and `{env_name}-policy-best`; with W&B off they go to `{checkpoint_dir}/{mode}/{run_name}/` instead, so a run never discards its weights. Diffusion checkpoints carry a `resume_metadata.json` sidecar — the authoritative record of the producing run's config, and what `--resume` reads to auto-detect `resume_step` and `resume_wandb_run_id`. PPO checkpoints carry `config.yaml` and `wandb-summary.json`.
-
-**Pass the checkpoint directory, not the step subdirectory** — `CheckpointManager` resolves the latest step itself.
+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`.
 
-Offline checkpoints save at the resolved env-frame budget, which is 99,942,400 for the Classic recipe at 512 envs (1525 updates × 512 × 128).
+**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. The layout below is the published layout.
-
-Released weights and the pre-computed ablation outputs are mirrored on the
-Hugging Face Hub at [`AnonMLuser/remdm-planner-craftax`](https://huggingface.co/AnonMLuser/remdm-planner-craftax),
-an anonymous account holding the layout below.
-
-```bash
-uv run hf download AnonMLuser/remdm-planner-craftax --include "checkpoints/**" --local-dir .
-```
+`checkpoints/` is gitignored; released weights live on the Hub at [`AnonMLuser/remdm-planner-craftax`](https://huggingface.co/AnonMLuser/remdm-planner-craftax), mirroring the layout below.
 
 | Checkpoint directory | Environment | Role | Trained for |
 |---|---|---|---|
@@ -349,41 +398,62 @@ uv run hf download AnonMLuser/remdm-planner-craftax --include "checkpoints/**" -
 | `checkpoints/ppo_agents/Craftax-Classic-Symbolic-v1-PPO_RNN-1000M` | Craftax Classic | PPO-RNN expert | 1e9 env frames |
 | `checkpoints/ppo_agents/Craftax-Symbolic-v1-PPO_RNN-1000M` | Full Craftax | PPO-RNN expert | 1e9 env frames |
 
-Full-Craftax diffusion planner checkpoints are not released: no full-Craftax training run has completed.
+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:
 
-### Experiment outputs
+```bash
+python main.py --mode online \
+    --config configs/final_craftax_gpu_24gb.yaml \
+    --ppo-checkpoint checkpoints/ppo_agents/Craftax-Symbolic-v1-PPO_RNN-1000M
+```
+
+```bash
+# 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.
 
-Ablation figures, tables and `results.json` are **not in the repository** and never
-should be: they are regenerated output, and 244 MB of them was rewritten out of the
-history. `experiments/rl_finetuning/outputs/` and `results/inference/` are gitignored.
+### Experiment outputs
 
-Regenerate them:
+Ablation figures, tables and `results.json` are regenerated output, so
+`experiments/rl_finetuning/outputs/` and `results/inference/` are gitignored. Obtain
+them either way:
 
 ```bash
-# Regenerate from a checkpoint; writes to outputs/{run_id}/
+# 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_final_results/{figures,tables}`
-from the working copy, so fetch or regenerate before running it. `demo_craftax.ipynb`
-needs no local copy — it reads them from its own `snapshot_download`.
+`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`.
 
-**Match the config to the checkpoint.** The model is built from the config, not the checkpoint, and a mismatch raises at restore. All released diffusion checkpoints are `d_model` 384, `n_heads` 8, `n_layers` 6, `d_ff` 768 — the architecture `defaults.yaml` also carries — so use the matching `final_*` config, which additionally sets the right `env_name` and recipe values:
+### Paper figures
 
-```bash
-python main.py --mode inference \
-    --config configs/final_craftax_classic_gpu_24gb.yaml \
-    --checkpoint checkpoints/online/Craftax-Classic-Symbolic-v1-Online-Diffusion-DAgger-100M
+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:
 
-# Train a new planner against the released Full Craftax PPO expert
-python main.py --mode online \
-    --config configs/final_craftax_gpu_24gb.yaml \
-    --ppo-checkpoint checkpoints/ppo_agents/Craftax-Symbolic-v1-PPO_RNN-1000M
+```bash
+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 three things, each keeping its repo-relative path: `checkpoints/`, 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.
+`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.
 
 ```bash
 HF_TOKEN=hf_xxx uv run python scripts/hf_upload.py --repo-id <ANON_HF_REPO_ID> --dry-run
@@ -391,11 +461,11 @@ HF_TOKEN=hf_xxx uv run python scripts/hf_upload.py --repo-id <ANON_HF_REPO_ID> -
 
 `--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>/` — the layout the Hub repo mirrors. 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 push already-published artefacts back up into a nested `checkpoints/hf/checkpoints/...` tree.
+**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 the accompanying paper (under submission); `demo_craftax.ipynb` reproduces the headline evaluation. Citation to be added on publication. Licence: MIT, see `LICENSE`.
+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`.
 
 ---
 
@@ -412,7 +482,7 @@ Set via the `env_name` config key.
 
 ## Remasking strategies
 
-Controlled by the `remask_strategy` key. All strategies operate on top of the three-phase loop controlled by `use_loop`, `t_on`, and `t_off`.
+Selected by `remask_strategy`, on top of the three-phase loop controlled by `use_loop`, `t_on` and `t_off`.
 
 | Strategy | Formula | Description |
 |---|---|---|
@@ -422,21 +492,18 @@ Controlled by the `remask_strategy` key. All strategies operate on top of the th
 
 ## Key hyperparameters
 
-**Environment**
+`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.
 
-| Parameter | Default | Description |
-|---|---|---|
-| `env_name` | `Craftax-Classic-Symbolic-v1` | Craftax environment ID. Use `Craftax-Symbolic-v1` for Full Craftax. |
-| `use_optimistic_resets` | `false` | Use `OptimisticResetVecEnvWrapper` instead of `AutoResetEnvWrapper` |
-| `optimistic_reset_ratio` | 16 | Fraction of envs reset per step when optimistic resets are enabled |
+**Environment.** `env_name` selects the benchmark: `Craftax-Classic-Symbolic-v1`
+(default) or `Craftax-Symbolic-v1` for Full Craftax.
 
 **Diffusion model**
 
 | Parameter | Default | Description |
 |---|---|---|
 | `plan_horizon` | 32 | Action plan length H |
-| `diffusion_steps` | 15 | Denoising steps T during training |
-| `diffusion_steps_eval` | 10 | Denoising steps T at inference |
+| `diffusion_steps` / `diffusion_steps_eval` | 15 / 10 | Denoising steps T at training and at inference |
 | `diffusion_schedule` | `cosine` | Noise schedule: `cosine` or `linear` |
 | `remask_strategy` | `rescale` | Remasking strategy: `rescale`, `cap`, or `conf` |
 | `train_sigma` | 0.0 | Per-token remasking correction during training (0 = standard MDLM) |
@@ -447,38 +514,24 @@ Controlled by the `remask_strategy` key. All strategies operate on top of the th
 | `temperature` | 0.5 | Softmax temperature for token sampling |
 | `top_p` | 0.95 | Nucleus sampling threshold |
 
-**Transformer architecture**
-
-| Parameter | Default | Description |
-|---|---|---|
-| `d_model` | 384 | Hidden dimension |
-| `n_heads` | 8 | Attention heads |
-| `n_layers` | 6 | Transformer blocks |
-| `d_ff` | 768 | FFN inner dimension |
-| `obs_encoder_layers` | 2 | MLP layers in the observation encoder |
-| `obs_encoder_width` | 768 | Observation encoder hidden width |
-| `dropout_rate` | 0.1 | Dropout rate (disabled at inference) |
+**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**
 
 | Parameter | Default | Description |
 |---|---|---|
 | `offline_total_timesteps` | 1e8 | Env-frame budget. Derives `num_updates` as `offline_total_timesteps // (num_envs * num_steps)`. |
-| `num_envs` | 1024 | Parallel environments |
-| `num_steps` | 128 | Environment steps collected per update |
-| `num_minibatches` | 8 | Gradient minibatches per epoch |
-| `update_epochs` | 8 | SGD epochs per update step |
+| `num_envs` / `num_steps` | 1024 / 128 | Parallel environments, and env steps per update; their product is `fpu` |
+| `num_minibatches` / `update_epochs` | 8 / 8 | Gradient minibatches per epoch, and epochs per update |
 | `num_repeats` | 1 | Independent training seeds (vmapped) |
 | `lr` | 3e-4 | AdamW learning rate (cosine-decayed to 10% over all gradient steps) |
 | `weight_decay` | 0.0 | Decoupled AdamW decay for core training; 0.0 is Adam exactly (the ablation suite keeps 1e-4) |
 | `lr_warmup_frames` | 1.6384e6 | Env-frame linear warm-up budget (0 = disabled). Derives `LR_WARMUP_STEPS` in gradient steps. |
 | `max_grad_norm` | 1.0 | Global gradient clipping norm |
 | `return_weight_cap` | 5.0 | Clip ceiling for per-window return weights (lower clip fixed at 0.1) |
-| `collect_temperature` | 1.0 | Softmax temperature on PPO logits during live data collection |
 | `val_interval_frames` | 1e6 | Env-frames between validation rollouts. Derives `VAL_INTERVAL` in update steps. |
-| `val_diffusion_steps` | 50 | Denoising steps during validation rollouts |
-| `val_replan_every` | 4 | Env steps executed per diffusion plan during validation |
-| `val_steps` | 256 | Total env steps per validation rollout |
 
 **Online DAgger training**
 
@@ -495,28 +548,24 @@ Controlled by the `remask_strategy` key. All strategies operate on top of the th
 
 | Parameter | Default | Description |
 |---|---|---|
-| `collect_num_steps` | 10000000 | Total environment steps to collect |
-| `collect_num_envs` | 128 | Parallel environments during collection |
+| `collect_num_steps` / `collect_num_envs` | 1e7 / 128 | Steps to collect, and envs collecting them |
 | `ppo_model_type` | `ppo_rnn` | PPO architecture: `ppo`, `ppo_rnn`, or `ppo_rnd` |
-| `layer_size` | 512 | PPO actor-critic hidden layer width |
-| `eval_steps` | 10000 | Environment steps for evaluation |
-| `eval_num_envs` | 32 | Parallel agents during evaluation (independent of `num_envs`) |
+| `eval_steps` / `eval_num_envs` | 10000 / 32 | Evaluation length and width (independent of `num_envs`) |
+| `inference_sampler` | `sample_plan` | `sample_plan` (the published protocol) or `inpainting` |
+| `eval_replan` | 8 | Env steps executed per plan under `sample_plan` |
 
-**Checkpointing / resume / logging**
+**Checkpointing / logging**
 
 | Parameter | Default | Description |
 |---|---|---|
 | `save_policy` | `true` | Save final checkpoint and upload as W&B artifact |
-| `resume_checkpoint_path` | `null` | Per-run: `--resume` (accepts `wandb:` refs) |
-| `resume_wandb_run_id` | `null` | Per-run: `--resume-wandb-run-id` (auto-read from metadata) |
-| `resume_step` | `null` | Per-run: `--resume-step` (auto-read from metadata) |
+| `checkpoint_dir` | `checkpoints` | Where checkpoints land when W&B is off |
 | `seed` | `null` | RNG seed (random if null; per-run: `--seed`) |
-| `use_wandb` | `true` | Enable Weights & Biases logging |
-| `wandb_project` | `remdm-planner-craftax` | W&B project name |
-| `wandb_entity` | `null` | W&B entity (null = your default entity) |
-| `wandb_download_dir` | `null` | Download dir for W&B artifacts; null = `./artifacts/` |
 | `jax_compilation_cache_dir` | `null` | Persistent XLA compilation cache; null = off. See below |
 
+The `resume_*`, `use_wandb` and `wandb_*` keys mirror the run flags documented under
+[Configuration](#configuration).
+
 ### Persistent compilation cache
 
 The whole training run is one `jax.jit`, so every process pays one large
@@ -533,7 +582,7 @@ python main.py --mode online --ppo-checkpoint <ppo> \
 
 ## Environment wrappers
 
-From `Craftax_Baselines/wrappers.py`:
+From `Craftax_Baselines/wrappers.py` (submodule):
 
 | Wrapper | Purpose |
 |---|---|
@@ -550,30 +599,29 @@ Stack (identical for training and inference): `env -> LogWrapper -> AutoResetEnv
 uv run pytest
 ```
 
-A CPU-only suite, 13 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.
+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.
 
 | File | Covers |
 |---|---|
-| `test_smoke_src.py`, `test_smoke_experiments.py` | that things **run**: imports, model from the real config, a gradient step, checkpoint round-trip, samplers, resolvers, every CLI entry point, and all 25 ablations' losses and optimizers |
-| `test_spec_*.py`, `test_method_spec*.py` | that things are **correct**: each canonical statement of the method/training specification pinned against the implementation |
+| `test_smoke_src.py`, `test_smoke_experiments.py` | that things **run**: imports, model from the real config, a gradient step, checkpoint round-trip, samplers, resolvers, every CLI entry point, and all 26 ablations' losses and optimizers |
+| `test_spec_*.py`, `test_method_spec*.py` | that things are **correct**: each canonical statement of the parent workspace's `the spec *.md` pinned against the implementation |
 | `test_config.py`, `test_recipe_values.py` | the preset, delta-only, cluster-sibling and poolability rules, and the shipped recipe values |
+| `test_gdelta.py`, `test_tex_macros.py` | the `--measure-gdelta` decomposition, and the `--emit-tex-macros` output: definitions only, uniquely named, letters only |
 | `test_gpu_agreement.py` | CPU/GPU agreement, skipped without a device |
 
 ## Implementation notes
 
-| Topic | Note                                                                                                                                                                                                                                                       |
-|---|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| JAX purity | `make_train_offline_diffusion` / `make_train_online_dagger` are fully JIT-compatible; env construction and checkpoint I/O sit outside `jax.jit`.                                                                                                           |
-| Offline data | `--mode offline` rolls out PPO live. `--mode collect` saves an `.npz` for inspection only — re-feeding it to `--mode offline` is unsupported; pass `--ppo-checkpoint`.                                                                                     |
-| Episode-boundary masking | 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 cumulative reward, normalised by the batch mean, clipped to `[0.1, return_weight_cap]`, and applied as per-sample multipliers before loss reduction.                                                                         |
-| LR schedule | Cosine decay `lr -> lr * 0.1` over all gradient steps. `lr_warmup_frames` prepends linear warm-up, converted to gradient steps as `(frames // fpu) * update_epochs * num_minibatches (* dagger_train_passes online)`.                                       |
-| Env-frame invariance | The six frame-denominated keys are converted to update-step form by `resolve_scaled_hyperparams()` using `fpu = num_envs * num_steps`, so one config runs on any hardware tier.                                                                            |
-| DAgger sizing | `dagger_sizing()` in `src/planners/common.py` is the single source of truth for `samples_per_update`, buffer capacity and `n_train_passes`.                                                                                                                |
-| Loss weight clipping | The MDLM SUBS weight `-alpha'(t) / (1 - alpha_t)` is clipped to 1000 for stability as `alpha_t -> 1`.                                                                                                                                                      |
-| Validation rollouts | Every `val_interval` updates, using inference sampling parameters with `val_diffusion_steps`, `val_replan_every` and `val_steps`.                                                                                                                          |
-| W&B namespaces | Centralised in `src/planners/logging.py`: `diffusion/`, `train/`, `env/`, `val/`, `dagger/`. `train/sps` only in modes with live env interaction.                                                                                                          |
-| DAgger aggregation | Ross et al. (2011). A circular buffer accumulates `(obs, expert_plan)` across iterations; windows use a sliding stride so every visited state contributes a label. The expert receives correct `done` flags so its RNN state resets at episode boundaries. |
-| Best-checkpoint tracking | Highest-validation-return parameters are kept alongside the live ones and uploaded as `{env_name}-policy-best`.                                                                                                                                            |
-| Denoising indexing | Reverse scan runs `step_idx = 0 -> T-1`, mapping to `t = (T - step_idx) / T` (high to low noise).                                                                                                                                                          |
-| PPO experts | Training lives entirely in `Craftax_Baselines/`; planner modes only consume checkpoints. Released PPO checkpoints were saved on GPU and fail to restore on a CPU-only machine.                                                                   |
+| Topic | Note |
+|---|---|
+| JAX purity | `make_train_offline_diffusion` / `make_train_online_dagger` are fully JIT-compatible; env construction and checkpoint I/O sit outside `jax.jit`. |
+| Offline data | `--mode offline` rolls out PPO live. `--mode collect` saves an `.npz` for inspection only — re-feeding it to `--mode offline` is unsupported. |
+| Episode-boundary masking | 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 cumulative reward, normalised by the batch mean, clipped to `[0.1, return_weight_cap]`, and applied as per-sample multipliers before loss reduction. |
+| LR schedule | Cosine decay `lr -> lr * 0.1` over all gradient steps. `lr_warmup_frames` prepends linear warm-up, converted as `(frames // fpu) * update_epochs * num_minibatches (* dagger_train_passes online)`. |
+| DAgger sizing | `dagger_sizing()` in `src/planners/common.py` is the single source of truth for `samples_per_update`, buffer capacity and `n_train_passes`. |
+| DAgger aggregation | Ross et al. (2011). A circular buffer accumulates `(obs, expert_plan)` across iterations, with a sliding stride so every visited state contributes a label; the expert receives correct `done` flags so its RNN state resets at episode boundaries. |
+| Loss weight clipping | The MDLM SUBS weight `-alpha'(t) / (1 - alpha_t)` is clipped to 1000 for stability as `alpha_t -> 1`. |
+| Denoising indexing | Reverse scan runs `step_idx = 0 -> T-1`, mapping to `t = (T - step_idx) / T` (high to low noise). |
+| Validation and best checkpoint | Every `val_interval` updates, at `val_diffusion_steps` / `val_replan_every` / `val_steps`. The highest-return parameters are kept alongside the live ones and uploaded as `{env_name}-policy-best`. |
+| W&B namespaces | Centralised in `src/planners/logging.py`: `diffusion/`, `train/`, `env/`, `val/`, `dagger/`. `train/sps` only in modes with live env interaction. |
+| PPO experts | Training lives entirely in `Craftax_Baselines/`; planner modes only consume checkpoints. Released PPO checkpoints were saved on GPU and fail to restore on a CPU-only machine. |

checkpoints/offline/Craftax-Classic-Symbolic-v1-Offline-Diffusion-BC-100M/99942400/commit_success.txt

@@ -1 +1 @@
-Checkpoint commit was successful to <path omitted for anonymous review>
+Checkpoint commit was successful to policies/99942400

checkpoints/offline/Craftax-Classic-Symbolic-v1-Offline-Diffusion-BC-100M/99942400/default/commit_success.txt

@@ -1 +1 @@
-Checkpoint commit was successful to <path omitted for anonymous review>
+Checkpoint commit was successful to 99942400.orbax-checkpoint-tmp/default

checkpoints/offline/Craftax-Classic-Symbolic-v1-Offline-Diffusion-BC-100M/resume_metadata.json

@@ -2,7 +2,6 @@
   "mode": "offline",
   "update_step": 1525,
   "total_gradient_steps_completed": 97600,
-  "wandb_run_id": null,
   "config_snapshot": {
     "ENV_NAME": "Craftax-Classic-Symbolic-v1",
     "USE_OPTIMISTIC_RESETS": false,
@@ -58,14 +57,9 @@
     "SAVE_POLICY": true,
     "CHECKPOINT_DIR": "checkpoints",
     "RESUME_CHECKPOINT_PATH": null,
-    "RESUME_WANDB_RUN_ID": null,
     "RESUME_STEP": null,
-    "JAX_COMPILATION_CACHE_DIR": "anonymous/jax-cache",
+    "JAX_COMPILATION_CACHE_DIR": null,
     "SEED": 42,
-    "USE_WANDB": true,
-    "WANDB_PROJECT": "remdm-planner-craftax",
-    "WANDB_ENTITY": "null",
-    "WANDB_DOWNLOAD_DIR": null,
     "MODE": "offline",
     "JIT": true,
     "PPO_CHECKPOINT_PATH": "checkpoints/ppo_agents/Craftax-Classic-Symbolic-v1-PPO_RNN-1000M",

checkpoints/online/Craftax-Classic-Symbolic-v1-Online-Diffusion-DAgger-100M/40370176/commit_success.txt

@@ -1 +1 @@
-Checkpoint commit was successful to <path omitted for anonymous review>
+Checkpoint commit was successful to policies_best/40370176

checkpoints/online/Craftax-Classic-Symbolic-v1-Online-Diffusion-DAgger-100M/40370176/default/commit_success.txt

@@ -1 +1 @@
-Checkpoint commit was successful to <path omitted for anonymous review>
+Checkpoint commit was successful to 40370176.orbax-checkpoint-tmp/default

checkpoints/online/Craftax-Classic-Symbolic-v1-Online-Diffusion-DAgger-100M/resume_metadata.json

@@ -2,7 +2,6 @@
   "mode": "online",
   "update_step": 1525,
   "total_gradient_steps_completed": 97600,
-  "wandb_run_id": null,
   "config_snapshot": {
     "ENV_NAME": "Craftax-Classic-Symbolic-v1",
     "USE_OPTIMISTIC_RESETS": false,
@@ -58,14 +57,9 @@
     "SAVE_POLICY": true,
     "CHECKPOINT_DIR": "checkpoints",
     "RESUME_CHECKPOINT_PATH": null,
-    "RESUME_WANDB_RUN_ID": null,
     "RESUME_STEP": null,
     "JAX_COMPILATION_CACHE_DIR": null,
     "SEED": 42,
-    "USE_WANDB": true,
-    "WANDB_PROJECT": "remdm-planner-craftax",
-    "WANDB_ENTITY": "null",
-    "WANDB_DOWNLOAD_DIR": null,
     "MODE": "online",
     "JIT": true,
     "PPO_CHECKPOINT_PATH": "checkpoints/ppo_agents/Craftax-Classic-Symbolic-v1-PPO_RNN-1000M",

checkpoints/ppo_agents/Craftax-Classic-Symbolic-v1-PPO_RNN-1000M/wandb-summary.json

@@ -1 +1,31 @@
-{"_runtime":5855,"Achievements/collect_stone":99.93519592285156,"achievements":19.780298233032227,"Achievements/collect_diamond":18.0168514251709,"Achievements/defeat_zombie":99.22229766845703,"Achievements/collect_wood":100.00000762939453,"Achievements/collect_coal":99.1574935913086,"Achievements/defeat_skeleton":93.77835845947266,"Achievements/make_wood_sword":99.74076843261719,"sps":173326.93437804468,"Achievements/place_table":100.00000762939453,"Achievements/place_stone":99.8703842163086,"Achievements/place_plant":99.74076843261719,"_wandb":{"runtime":5855},"_timestamp":1.7731733404075222e+09,"Achievements/make_iron_sword":97.40765380859375,"_step":1906,"episode_return":19.358198165893555,"Achievements/eat_cow":99.54634094238281,"Achievements/collect_sapling":99.74076843261719,"Achievements/make_wood_pickaxe":100.00000762939453,"Achievements/place_furnace":99.8703842163086,"Achievements/collect_drink":98.57421112060547,"episode_length":341.3253479003906,"Achievements/make_stone_sword":99.54634094238281,"Achievements/make_iron_pickaxe":92.22294616699219,"Achievements/wake_up":83.79779815673828,"Achievements/collect_iron":98.18535614013672,"Achievements/eat_plant":0.3240440785884857,"Achievements/make_stone_pickaxe":99.35191345214844}
+{
+  "_runtime": 5855,
+  "Achievements/collect_stone": 99.93519592285156,
+  "achievements": 19.780298233032227,
+  "Achievements/collect_diamond": 18.0168514251709,
+  "Achievements/defeat_zombie": 99.22229766845703,
+  "Achievements/collect_wood": 100.00000762939453,
+  "Achievements/collect_coal": 99.1574935913086,
+  "Achievements/defeat_skeleton": 93.77835845947266,
+  "Achievements/make_wood_sword": 99.74076843261719,
+  "sps": 173326.93437804468,
+  "Achievements/place_table": 100.00000762939453,
+  "Achievements/place_stone": 99.8703842163086,
+  "Achievements/place_plant": 99.74076843261719,
+  "_timestamp": 1773173340.4075222,
+  "Achievements/make_iron_sword": 97.40765380859375,
+  "_step": 1906,
+  "episode_return": 19.358198165893555,
+  "Achievements/eat_cow": 99.54634094238281,
+  "Achievements/collect_sapling": 99.74076843261719,
+  "Achievements/make_wood_pickaxe": 100.00000762939453,
+  "Achievements/place_furnace": 99.8703842163086,
+  "Achievements/collect_drink": 98.57421112060547,
+  "episode_length": 341.3253479003906,
+  "Achievements/make_stone_sword": 99.54634094238281,
+  "Achievements/make_iron_pickaxe": 92.22294616699219,
+  "Achievements/wake_up": 83.79779815673828,
+  "Achievements/collect_iron": 98.18535614013672,
+  "Achievements/eat_plant": 0.3240440785884857,
+  "Achievements/make_stone_pickaxe": 99.35191345214844
+}

checkpoints/ppo_agents/Craftax-Symbolic-v1-PPO_RNN-1000M/config.yaml

@@ -42,11 +42,5 @@ UPDATE_EPOCHS:
   value: 4
 USE_OPTIMISTIC_RESETS:
   value: true
-USE_WANDB:
-  value: true
 VF_COEF:
   value: 0.5
-WANDB_ENTITY:
-  value: null
-WANDB_PROJECT:
-  value: null

checkpoints/ppo_agents/Craftax-Symbolic-v1-PPO_RNN-1000M/wandb-summary.json

@@ -1 +1,76 @@
-{"Achievements/place_plant":97.02970123291016,"Achievements/collect_drink":88.11881256103516,"Achievements/place_torch":92.07920837402344,"Achievements/defeat_pigman":0,"Achievements/enter_vault":0,"Achievements/defeat_skeleton":34.158416748046875,"Achievements/collect_ruby":14.356435775756836,"Achievements/make_wood_sword":95.04950714111328,"_timestamp":1.77421704354534e+09,"Achievements/defeat_troll":0,"Achievements/defeat_kobold":0,"Achievements/make_stone_sword":89.1089096069336,"Achievements/defeat_necromancer":0,"Achievements/fire_bow":68.81188201904297,"Achievements/collect_stone":98.51484680175781,"Achievements/defeat_gnome_archer":0,"Achievements/cast_iceball":0,"Achievements/enchant_armour":0,"Achievements/make_diamond_sword":1.9801980257034302,"Achievements/make_iron_pickaxe":6.93069314956665,"_wandb":{"runtime":10138},"Achievements/eat_plant":0.9900990128517151,"Achievements/enter_gnomish_mines":0,"Achievements/damage_necromancer":0,"Achievements/find_bow":82.67326354980469,"Achievements/defeat_gnome_warrior":0,"_runtime":10138,"Achievements/enter_graveyard":0,"Achievements/eat_cow":70.79207611083984,"Achievements/drink_potion":58.41584014892578,"Achievements/enter_dungeon":87.62376403808594,"Achievements/eat_bat":0,"Achievements/defeat_deep_thing":0,"Achievements/make_iron_sword":3.9603960514068604,"achievements":23.94554901123047,"Achievements/make_wood_pickaxe":99.00990295410156,"Achievements/eat_snail":75.7425765991211,"Achievements/wake_up":53.96039581298828,"Achievements/cast_fireball":0,"Achievements/defeat_zombie":44.05940628051758,"Achievements/defeat_lizard":0,"Achievements/make_arrow":84.15841674804688,"Achievements/defeat_frost_troll":0,"Achievements/make_diamond_armour":0,"Achievements/enter_fire_realm":0,"Achievements/make_torch":92.07920837402344,"Achievements/enchant_sword":0,"episode_length":331.6633605957031,"Achievements/place_table":99.50495147705078,"Achievements/enter_ice_realm":0,"Achievements/learn_fireball":0,"Achievements/collect_wood":99.50495147705078,"Achievements/collect_sapphire":12.87128734588623,"Achievements/open_chest":82.67326354980469,"Achievements/collect_iron":65.3465347290039,"Achievements/defeat_archer":0,"Achievements/defeat_knight":0,"Achievements/make_iron_armour":0,"Achievements/place_furnace":98.51484680175781,"Achievements/defeat_orc_solider":60.39603805541992,"Achievements/collect_sapling":97.02970123291016,"_step":15257,"Achievements/make_diamond_pickaxe":3.9603960514068604,"Achievements/defeat_orc_mage":50.99009704589844,"Achievements/place_stone":98.51484680175781,"episode_return":35.05544662475586,"Achievements/make_stone_pickaxe":83.1683120727539,"Achievements/enter_troll_mines":0,"Achievements/defeat_fire_elemental":0,"Achievements/learn_iceball":0,"Achievements/collect_diamond":15.841584205627441,"Achievements/collect_coal":86.63365936279297,"Achievements/defeat_ice_elemental":0,"Achievements/enter_sewers":0,"sps":100008.66166595295}
+{
+  "Achievements/place_plant": 97.02970123291016,
+  "Achievements/collect_drink": 88.11881256103516,
+  "Achievements/place_torch": 92.07920837402344,
+  "Achievements/defeat_pigman": 0,
+  "Achievements/enter_vault": 0,
+  "Achievements/defeat_skeleton": 34.158416748046875,
+  "Achievements/collect_ruby": 14.356435775756836,
+  "Achievements/make_wood_sword": 95.04950714111328,
+  "_timestamp": 1774217043.54534,
+  "Achievements/defeat_troll": 0,
+  "Achievements/defeat_kobold": 0,
+  "Achievements/make_stone_sword": 89.1089096069336,
+  "Achievements/defeat_necromancer": 0,
+  "Achievements/fire_bow": 68.81188201904297,
+  "Achievements/collect_stone": 98.51484680175781,
+  "Achievements/defeat_gnome_archer": 0,
+  "Achievements/cast_iceball": 0,
+  "Achievements/enchant_armour": 0,
+  "Achievements/make_diamond_sword": 1.9801980257034302,
+  "Achievements/make_iron_pickaxe": 6.93069314956665,
+  "Achievements/eat_plant": 0.9900990128517151,
+  "Achievements/enter_gnomish_mines": 0,
+  "Achievements/damage_necromancer": 0,
+  "Achievements/find_bow": 82.67326354980469,
+  "Achievements/defeat_gnome_warrior": 0,
+  "_runtime": 10138,
+  "Achievements/enter_graveyard": 0,
+  "Achievements/eat_cow": 70.79207611083984,
+  "Achievements/drink_potion": 58.41584014892578,
+  "Achievements/enter_dungeon": 87.62376403808594,
+  "Achievements/eat_bat": 0,
+  "Achievements/defeat_deep_thing": 0,
+  "Achievements/make_iron_sword": 3.9603960514068604,
+  "achievements": 23.94554901123047,
+  "Achievements/make_wood_pickaxe": 99.00990295410156,
+  "Achievements/eat_snail": 75.7425765991211,
+  "Achievements/wake_up": 53.96039581298828,
+  "Achievements/cast_fireball": 0,
+  "Achievements/defeat_zombie": 44.05940628051758,
+  "Achievements/defeat_lizard": 0,
+  "Achievements/make_arrow": 84.15841674804688,
+  "Achievements/defeat_frost_troll": 0,
+  "Achievements/make_diamond_armour": 0,
+  "Achievements/enter_fire_realm": 0,
+  "Achievements/make_torch": 92.07920837402344,
+  "Achievements/enchant_sword": 0,
+  "episode_length": 331.6633605957031,
+  "Achievements/place_table": 99.50495147705078,
+  "Achievements/enter_ice_realm": 0,
+  "Achievements/learn_fireball": 0,
+  "Achievements/collect_wood": 99.50495147705078,
+  "Achievements/collect_sapphire": 12.87128734588623,
+  "Achievements/open_chest": 82.67326354980469,
+  "Achievements/collect_iron": 65.3465347290039,
+  "Achievements/defeat_archer": 0,
+  "Achievements/defeat_knight": 0,
+  "Achievements/make_iron_armour": 0,
+  "Achievements/place_furnace": 98.51484680175781,
+  "Achievements/defeat_orc_solider": 60.39603805541992,
+  "Achievements/collect_sapling": 97.02970123291016,
+  "_step": 15257,
+  "Achievements/make_diamond_pickaxe": 3.9603960514068604,
+  "Achievements/defeat_orc_mage": 50.99009704589844,
+  "Achievements/place_stone": 98.51484680175781,
+  "episode_return": 35.05544662475586,
+  "Achievements/make_stone_pickaxe": 83.1683120727539,
+  "Achievements/enter_troll_mines": 0,
+  "Achievements/defeat_fire_elemental": 0,
+  "Achievements/learn_iceball": 0,
+  "Achievements/collect_diamond": 15.841584205627441,
+  "Achievements/collect_coal": 86.63365936279297,
+  "Achievements/defeat_ice_elemental": 0,
+  "Achievements/enter_sewers": 0,
+  "sps": 100008.66166595295
+}

configs/defaults.yaml

@@ -136,6 +136,16 @@ layer_size: 512
 # ── Inference (--mode inference) ─────────────────────────────────────────────
 eval_steps: 10000
 eval_num_envs: 32
+# Sampler used by --mode inference:
+#   sample_plan  replan from scratch every eval_replan steps, conditioned only
+#                on the current observation. What build_eval_fn does, and the
+#                protocol behind every published number.
+#   inpainting   replan every step with each executed action locked as an
+#                inpainting prefix. An ablation on the planning-as-inpainting
+#                design choice, not the source of any published number; it
+#                scores far lower on the same weights.
+inference_sampler: sample_plan
+eval_replan: 8                      # env steps executed per plan (sample_plan)
 
 # ── Checkpointing ────────────────────────────────────────────────────────────
 save_policy: true

demo_craftax.ipynb

@@ -61,10 +61,9 @@
     "# =============================================================================\n",
     "\n",
     "# Public HuggingFace repo holding the released artefacts: source, checkpoints\n",
-    "# and the pre-computed ablation outputs. This is an anonymous mirror published\n",
-    "# for double-blind review, so it carries no account name or URL that would\n",
-    "# identify the authors. Point it at your own Hub repo holding the same layout,\n",
-    "# or train from source (see README.md).\n",
+    "# and the pre-computed ablation outputs. No authentication required. Point it\n",
+    "# at your own Hub repo holding the same layout, or train from source\n",
+    "# (see README.md).\n",
     "HF_REPO_ID = \"AnonMLuser/remdm-planner-craftax\"\n",
     "LOCAL_DIR = \"remdm-planner-craftax\"\n",
     "\n",
@@ -264,7 +263,7 @@
     "    \"craftax>=1.5.0\",\n",
     "    \"flax>=0.12.6\",\n",
     "    \"optax>=0.2.8\",\n",
-    "    \"orbax-checkpoint>=0.5\",\n",
+    "    \"orbax-checkpoint>=0.12\",\n",
     "    \"distrax>=0.1.7\",\n",
     "    \"chex>=0.1.91\",\n",
     "    \"polars>=1.39.3\",\n",
@@ -293,9 +292,9 @@
     "\n",
     "if HF_REPO_ID == \"UNSET_HF_REPO_ID\":\n",
     "    raise RuntimeError(\n",
-    "        \"HF_REPO_ID is unset: checkpoints are withheld for anonymous review. \"\n",
-    "        \"Edit the constant in Cell 1 to point at a Hub repo holding the \"\n",
-    "        \"published layout, or train from source (see README.md).\"\n",
+    "        \"HF_REPO_ID is unset. Edit the constant in Cell 1 to point at a Hub \"\n",
+    "        \"repo holding the published layout, or train from source \"\n",
+    "        \"(see README.md).\"\n",
     "    )\n",
     "\n",
     "snapshot_path = snapshot_download(repo_id=HF_REPO_ID, local_dir=LOCAL_DIR)\n",
@@ -537,6 +536,7 @@
     "            \"EVAL_STEPS\": EVAL_STEPS,\n",
     "            \"EVAL_NUM_ENVS\": EVAL_NUM_ENVS,\n",
     "            \"DIFFUSION_STEPS_EVAL\": INPAINTING_DIFFUSION_STEPS,\n",
+    "            \"INFERENCE_SAMPLER\": \"inpainting\",\n",
     "            \"USE_WANDB\": False,\n",
     "            \"CHECKPOINT_PATH\": DIFFUSION_ONLINE_CKPT,\n",
     "        }\n",
@@ -981,12 +981,16 @@
     "   **1.16** on Craftax Classic (effective sample size 437 of a 1024-window batch),\n",
     "   so the returns *do* rank windows within a batch and the familiar explanation\n",
     "   that sparse returns fail to discriminate is unavailable. Measured directly at\n",
-    "   the pretrained parameters, `‖g_δ‖/‖∇L_BC‖ = 0.49 ± 0.01` at cosine\n",
-    "   **0.02 ± 0.05**. A cosine needs a reference: in `D = 9.33 × 10⁶` dimensions two\n",
-    "   independent directions have `cos ≈ 3.3 × 10⁻⁴`, so 0.02 is about sixty standard\n",
-    "   deviations from random — `g_δ` carries structure. But two independent noise\n",
-    "   draws of `∇L_BC` itself agree at **0.893 ± 0.010**, so `g_δ` points somewhere\n",
-    "   imitation does not.\n",
+    "   the pretrained parameters, `‖g_δ‖/‖∇L_BC‖ = 0.48 ± 0.01` at cosine\n",
+    "   **0.02 ± 0.05**. A cosine needs a reference. In `D = 9.33 × 10⁶` dimensions two\n",
+    "   independent directions have `cos ~ N(0, 1/D)`, a standard deviation of\n",
+    "   `3.3 × 10⁻⁴` — but the seed-level dispersion of the measured cosine is two\n",
+    "   orders of magnitude larger than that, so **we do not read this column as a test\n",
+    "   against isotropy**: 0.02 ± 0.05 is consistent with zero at the precision of the\n",
+    "   seed mean. The informative reference is how well imitation agrees with itself:\n",
+    "   two independent noise draws of `∇L_BC` give **0.893 ± 0.010**. Against that,\n",
+    "   `g_δ` is half the imitation gradient in norm and points somewhere imitation\n",
+    "   does not.\n",
     "3. **Shrinking the return term makes things worse — the central negative control.**\n",
     "   If the return term drove the degradation, suppressing it should help. Advantage\n",
     "   clipping confines the weights to [0.8, 1.2] and does exactly what it should:\n",
@@ -994,10 +998,19 @@
     "   **0.097 ± 0.003**, a fivefold reduction. Its score is **5.06** — 3.16 below\n",
     "   baseline RL and second worst in the suite. The binary win mask sits between the\n",
     "   two on both axes (`CV_A` 0.77, ratio 0.417 ± 0.025) and scores 7.28, also below\n",
-    "   baseline. Across these three, the ordering of the return term's magnitude is the\n",
-    "   **reverse** of the ordering of final score. A condition whose return term is\n",
+    "   baseline. Across these three the return term's magnitude and final score are\n",
+    "   ordered **the same way**, largest term scoring highest — which is the reverse of\n",
+    "   what a return-driven degradation predicts. A condition whose return term is\n",
     "   nearly absent degrades further than one whose return term is large, so the\n",
     "   degradation cannot be attributed to the return weighting.\n",
+    "4. **Most of the ratio is weight dispersion, not the return signal.** A\n",
+    "   shuffled-`δ` null permutes `δ` across the batch, preserving the weight multiset\n",
+    "   and `CV_A` while destroying every association between a window's weight and its\n",
+    "   own gradient. On Craftax Classic it retains roughly three quarters of the\n",
+    "   measured ratio at every transform, at a near-constant **0.36 of `CV_A`**. So on\n",
+    "   this benchmark the ratio is largely a function of dispersion and the ordering\n",
+    "   above tracks dispersion rather than the return signal specifically. MiniHack\n",
+    "   separates the two much further — see its notebook.\n",
     "\n",
     "### What is left is the data\n",
     "\n",
@@ -1147,12 +1160,235 @@
     "    print(fh.read())"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. The return term, measured directly\n",
+    "\n",
+    "Section 2's conditions vary *how* the return enters the loss. This section measures\n",
+    "*how much* it contributes, at a single parameter point, with no training involved.\n",
+    "\n",
+    "The return-weighted ELBO gradient factors exactly, with no approximation:\n",
+    "\n",
+    "```\n",
+    "∇L_RW  =  Ā · [ ∇L_BC  +  (1/B) Σ_i δ_i ∇ℓ_i ]      δ_i = A_i/Ā − 1\n",
+    "                          └────── g_δ ──────┘\n",
+    "```\n",
+    "\n",
+    "`Ā` rescales the step size and leaves the direction unchanged, so everything the\n",
+    "return does enters through `g_δ`. Both terms are computable at one parameter point,\n",
+    "so we compute them: load the checkpoint every ablation starts from, collect one\n",
+    "on-policy batch under the suite's own rollout function, and take `∇L_BC`, `g_δ` and\n",
+    "`∇L_RW` on that batch under a **shared** `(z_t, t)` draw, so the only difference is\n",
+    "the weight vector. 3 rollout seeds × 8 noise draws, 1,024-window batches. It runs in\n",
+    "minutes on a laptop CPU — `run_ablations.py --measure-gdelta`.\n",
+    "\n",
+    "**Correctness check.** Reconstructing `∇L_RW` from the two right-hand terms and\n",
+    "comparing against the directly computed one gives a relative residual below\n",
+    "**5 × 10⁻⁷** on every draw — float32 accumulation over 9.33M parameters, not a\n",
+    "modelling gap.\n",
+    "\n",
+    "The `Shuffled δ` column permutes `δ` across the batch. That preserves the weight\n",
+    "multiset and `CV_A` exactly while destroying every association between a window's\n",
+    "weight and its own gradient, so it bounds how much of the ratio needs that\n",
+    "association at all.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# =============================================================================\n",
+    "# CELL 12 — The measured return term (paper Table 2, appendix \"Measuring the\n",
+    "#           Return Term\"). Read from the shipped gdelta artefacts, not retyped.\n",
+    "# =============================================================================\n",
+    "\n",
+    "import json\n",
+    "import os\n",
+    "\n",
+    "import polars as pl\n",
+    "\n",
+    "RUN_DIR = \"experiments/rl_finetuning/outputs/craftax_classic_ablations\"\n",
+    "GDELTA_DIR = os.path.join(RUN_DIR, \"gdelta\")\n",
+    "\n",
+    "gdelta_df = pl.read_csv(os.path.join(RUN_DIR, \"tables\", \"gdelta.csv\"))\n",
+    "print(\"Return term at the pretrained checkpoint, 3 seeds x 8 draws:\")\n",
+    "print(\n",
+    "    gdelta_df.select(\n",
+    "        [\"Transform\", \"CV_A\", \"Ratio\", \"Ratio_Sd\", \"Ratio_Shuffled\", \"Cos\", \"Cos_Sd\",\n",
+    "         \"A1_Violated\"]\n",
+    "    )\n",
+    ")\n",
+    "\n",
+    "with open(os.path.join(GDELTA_DIR, \"gdelta_aggregate.json\")) as fh:\n",
+    "    agg = json.load(fh)\n",
+    "\n",
+    "print()\n",
+    "_draws = agg[\"n_draws_per_seed\"]\n",
+    "_draws = _draws[0] if isinstance(_draws, list) and len(set(_draws)) == 1 else _draws\n",
+    "print(f\"seeds x draws            : {agg['n_seeds']} x {_draws}\")\n",
+    "print(f\"batch (windows)          : {agg['batch']}\")\n",
+    "print(f\"parameters               : {agg['n_params']:,}\")\n",
+    "print(f\"Eq.-4 max rel. residual  : {agg['eq4_residual_max']:.3e}\")\n",
+    "print(f\"cos(random directions) sd: {agg['random_cos_sd']:.3e}\")\n",
+    "print(\n",
+    "    f\"cos(dL_BC, dL_BC') mean  : {agg['bc_self_cos_mean']:.3f} \"\n",
+    "    f\"+/- {agg['bc_self_cos_std_seeds']:.3f}   <- what alignment looks like when a \"\n",
+    "    f\"direction IS imitation\"\n",
+    ")\n",
+    "\n",
+    "print()\n",
+    "print(\n",
+    "    \"Reading it: the ratio and the final score are ordered the SAME way across the\\n\"\n",
+    "    \"three transforms where the decomposition is defined (baseline 0.485/8.22,\\n\"\n",
+    "    \"bc_wins 0.417/7.28, advantage_clip 0.097/5.06). A return-driven degradation\\n\"\n",
+    "    \"predicts the opposite. normalized_adv is excluded: it mean-centres the weights,\\n\"\n",
+    "    \"so Abar ~ 0, delta is undefined and (A1) fails -- the objective is no longer a\\n\"\n",
+    "    \"bound there, which is the point, not a measurement artefact.\\n\"\n",
+    "    \"\\n\"\n",
+    "    \"Note the Shuffled column: on Craftax Classic it retains roughly three quarters\\n\"\n",
+    "    \"of each measured ratio, at a near-constant 0.36 of CV_A. Most of the ratio is\\n\"\n",
+    "    \"weight dispersion rather than the return signal specifically. MiniHack does not\\n\"\n",
+    "    \"reproduce that -- there the null costs about five sixths -- so the redundancy is\\n\"\n",
+    "    \"a Craftax-side property, not a law.\"\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Review arms: the unweighted control and the learning-rate sweep\n",
+    "\n",
+    "The 25-condition suite leaves two accounts of the degradation open, and each has a\n",
+    "run that closes it. These sit **outside** the suite's tables — they are separate\n",
+    "runs on a different host, at three seeds each, shipped under\n",
+    "`experiments/rl_finetuning/outputs/review_*/`.\n",
+    "\n",
+    "**Read them against each other, not against Section 2.** The suite does not\n",
+    "reproduce across hosts: the same checkpoint and configuration score the pretrained\n",
+    "planner at **11.9754** here against **11.8084** in Section 2, and baseline RL at\n",
+    "**8.3395** against **8.2233** — a consistent offset of 1.41%. So `baseline_rl` was\n",
+    "re-run alongside the new arms as a same-stack anchor, and every comparison below is\n",
+    "within the review stack.\n",
+    "\n",
+    "| Arm | What it asks |\n",
+    "|---|---|\n",
+    "| `review_anchor_baseline_rl` | The same-stack anchor every other row is quoted against |\n",
+    "| `review_run1_bc_all` | **Is the weighting the problem, or the data?** Every weight fixed at one, over the same rollout windows — the unweighted control |\n",
+    "| `review_run2_advclip_lr_matched` | **Was the clipping control just a smaller step?** `advantage_clip` re-run with its learning rate scaled by the measured `Ā_base/Ā_clip = 1.0744`, giving 3.223e-4 |\n",
+    "| `review_run4_baseline_lr1e-4` / `_lr1e-5` | **Is the degradation just too large a step?** The suite fine-tunes at the pretraining learning rate of 3e-4 |\n",
+    "\n",
+    "The MiniHack counterparts of all three questions were run too, and the paper reports\n",
+    "them, but **their run directories are not part of this release** — only the Craftax\n",
+    "Classic ones are shipped here. The MiniHack numbers are quoted in its notebook.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# =============================================================================\n",
+    "# CELL 13 — Review arms, read from the shipped run directories\n",
+    "# =============================================================================\n",
+    "\n",
+    "import os\n",
+    "\n",
+    "import polars as pl\n",
+    "\n",
+    "OUT = \"experiments/rl_finetuning/outputs\"\n",
+    "REVIEW_RUNS = [\n",
+    "    (\"review_anchor_baseline_rl\",      \"anchor: baseline_rl @3e-4 (same stack)\"),\n",
+    "    (\"review_run1_bc_all\",             \"bc_all: uniform weights, all windows\"),\n",
+    "    (\"review_run2_advclip_lr_matched\", \"advantage_clip @3.223e-4 (matched step)\"),\n",
+    "    (\"review_run4_baseline_lr1e-4\",    \"baseline_rl @1e-4\"),\n",
+    "    (\"review_run4_baseline_lr1e-5\",    \"baseline_rl @1e-5\"),\n",
+    "]\n",
+    "\n",
+    "rows = []\n",
+    "for run, label in REVIEW_RUNS:\n",
+    "    csv = os.path.join(OUT, run, \"tables\", \"main_results.csv\")\n",
+    "    if not os.path.exists(csv):\n",
+    "        print(f\"missing: {csv}\")\n",
+    "        continue\n",
+    "    df = pl.read_csv(csv)\n",
+    "    r = df.row(0, named=True)\n",
+    "    rows.append(\n",
+    "        {\n",
+    "            \"Arm\": label,\n",
+    "            \"Condition\": r[\"Method\"],\n",
+    "            \"Score\": r[\"Final_Score\"],\n",
+    "            \"Sd\": r[\"Score_Std\"],\n",
+    "            \"vs_Pretrained\": r[\"Delta_vs_Pretrained\"],\n",
+    "        }\n",
+    "    )\n",
+    "\n",
+    "review_df = pl.DataFrame(rows)\n",
+    "print(\"Review arms (same-stack pretrained checkpoint = 11.9754):\")\n",
+    "print(review_df)\n",
+    "\n",
+    "anchor = next(r[\"Score\"] for r in rows if r[\"Condition\"] == \"baseline_rl\" and \"anchor\" in r[\"Arm\"])\n",
+    "bc_all = next((r[\"Score\"] for r in rows if r[\"Condition\"] == \"bc_all\"), None)\n",
+    "lr5 = next((r[\"Score\"] for r in rows if \"1e-5\" in r[\"Arm\"]), None)\n",
+    "PRETRAINED_SAME_STACK = 11.9754\n",
+    "\n",
+    "print()\n",
+    "if bc_all is not None:\n",
+    "    print(\n",
+    "        f\"Unweighted control : bc_all {bc_all:.4f} vs anchor {anchor:.4f} \"\n",
+    "        f\"= {bc_all - anchor:+.4f}\"\n",
+    "    )\n",
+    "    print(\n",
+    "        \"                     Removing the weighting while keeping the rollouts\\n\"\n",
+    "        \"                     degrades the model FURTHER than the weighted objective\\n\"\n",
+    "        \"                     does. The seeds do not overlap. That is the protective\\n\"\n",
+    "        \"                     reading -- the inverse of a return-driven degradation.\\n\"\n",
+    "        \"                     The host offset works against this result, not for it:\\n\"\n",
+    "        \"                     on the published stack bc_all would sit lower still.\\n\"\n",
+    "        \"                     Exact two-sided permutation test: p = 0.10, which is the\\n\"\n",
+    "        \"                     design floor at 3 seeds, not a significant result.\"\n",
+    "    )\n",
+    "if lr5 is not None:\n",
+    "    shortfall = PRETRAINED_SAME_STACK - lr5\n",
+    "    recovered = (lr5 - anchor) / (PRETRAINED_SAME_STACK - anchor)\n",
+    "    print()\n",
+    "    print(\n",
+    "        f\"Learning-rate sweep: 3e-4 {anchor:.4f} -> 1e-4 -> 1e-5 {lr5:.4f}; still \"\n",
+    "        f\"{shortfall:.4f} below the same-stack checkpoint,\"\n",
+    "    )\n",
+    "    print(\n",
+    "        f\"                     having recovered {recovered:.0%} of the default's \"\n",
+    "        f\"shortfall. That residual is more than ten times the\"\n",
+    "    )\n",
+    "    print(\n",
+    "        \"                     seed sd at that rate, and the sd does not fall as the \"\n",
+    "        \"step shrinks. Lowering the step by a\"\n",
+    "    )\n",
+    "    print(\n",
+    "        \"                     factor of thirty buys back two thirds of the \"\n",
+    "        \"degradation WITHOUT removing it. 3e-5 was not run.\"\n",
+    "    )\n",
+    "print()\n",
+    "print(\n",
+    "    \"Matched effective step: advantage_clip at 3.223e-4 reaches 4.9246, leaving it\\n\"\n",
+    "    \"                     3.4149 below the same-stack anchor, where the unmatched arm\\n\"\n",
+    "    \"                     sat 3.1601 below the published baseline. Equalising the\\n\"\n",
+    "    \"                     effective step does NOT close the clipping gap, so Section 3's\\n\"\n",
+    "    \"                     central negative control was not measuring step size.\"\n",
+    ")\n"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "d2dbe7e3e2ab",
    "metadata": {},
    "source": [
-    "## 3. Conclusions\n",
+    "## 5. Conclusions\n",
     "\n",
     "1. **Return-weighted ELBO fine-tuning degrades the checkpoint.** From **11.81**,\n",
     "   baseline RL falls to **8.22 ± 0.14** and none of the 25 conditions ends above\n",
@@ -1167,21 +1403,27 @@
     "   3.73** over the same iterations in which eval score falls from **12.06 to\n",
     "   8.49**. The objective is failing on its own terms too.\n",
     "3. **The reward is not what does the damage.** The return term is real and large:\n",
-    "   `‖g_δ‖/‖∇L_BC‖ = 0.49 ± 0.01` at cosine `0.02 ± 0.05`, against `0.893 ± 0.010`\n",
-    "   for `∇L_BC` against itself — around sixty standard deviations from a random\n",
-    "   direction, so it is structure and not noise. But the central control does not\n",
-    "   go our way. Advantage clipping cuts that ratio fivefold to `0.097 ± 0.003` and\n",
-    "   scores **3.16 below** baseline RL, second worst in the suite. Across the three\n",
-    "   transforms where the decomposition is defined, the ordering of the return\n",
-    "   term's magnitude is the reverse of the ordering of final score.\n",
+    "   `‖g_δ‖/‖∇L_BC‖ = 0.48 ± 0.01`, half the imitation gradient in norm, against\n",
+    "   `0.893 ± 0.010` for `∇L_BC` against itself. Its cosine with imitation is\n",
+    "   `0.02 ± 0.05`, which we read as consistent with zero rather than as a departure\n",
+    "   from isotropy. But the central control does not go our way. Advantage clipping\n",
+    "   cuts that ratio fivefold to `0.097 ± 0.003` and scores **3.16 below** baseline\n",
+    "   RL, second worst in the suite. Across the three transforms where the\n",
+    "   decomposition is defined, the return term's magnitude and final score are\n",
+    "   ordered the same way — the reverse of what a return-driven degradation\n",
+    "   predicts.\n",
     "4. **What is left is the data.** The suite is ordered by plasticity — conditions\n",
     "   that restrict the update keep more of the checkpoint (Spearman **−0.79**\n",
-    "   between final KL and final score) — and the condition closest to unweighted\n",
-    "   training on self-generated rollouts sits near the bottom. Together these point\n",
-    "   at fine-tuning on the model's own rollouts as the damaging ingredient, with the\n",
-    "   return weighting a large but incidental passenger. The paper does not dress\n",
-    "   this up as more than it is: without the unweighted arm it is an inference from\n",
-    "   a near-substitute.\n",
+    "   between final KL and final score) — and the unweighted control now exists.\n",
+    "   `bc_all` fixes every weight at one over the same rollout windows the weighted\n",
+    "   baseline trains on, and it reaches **4.73 ± 0.19** against a same-stack anchor\n",
+    "   of **8.34 ± 0.15** (Section 4 below). Removing the weighting while keeping the\n",
+    "   rollouts costs *more* than the weighting itself, on both benchmarks. That fits\n",
+    "   a **protective** reading of the return weighting rather than a damaging one —\n",
+    "   the inverse of the reading this study set out to test. It does not show whether\n",
+    "   the rollouts themselves contribute, which would need an arm that does not train\n",
+    "   on them, and the pairwise permutation test reaches only `p = 0.10`, the design\n",
+    "   floor at three seeds.\n",
     "\n",
     "**Why the sign constraint is not the whole story.** Eq. 3 is a bound only for\n",
     "`A_i ≥ 0`, so the objective can re-rank sampled behaviour but never push mass away\n",
@@ -1221,7 +1463,9 @@
     "objective: `CV_A` comes free from an effective-sample-size counter and says whether\n",
     "the returns rank anything at all. The ratio `‖g_δ‖/‖∇L_BC‖` costs two backward\n",
     "passes and says how much of the update the reward is responsible for. Ours said\n",
-    "0.49, which looked like a mechanism until the clipping control said otherwise.\n",
+    "0.48, which looked like a mechanism until the clipping control said otherwise.\n",
+    "Take the shuffled-`δ` null alongside it — on Craftax Classic three quarters of\n",
+    "that 0.48 survives destroying the weight/gradient association entirely.\n",
     "\n",
     "**Open problem.** Per-step formulations over the denoising chain (d1, DiffPO) admit\n",
     "**signed** advantages, and so change the object being optimised rather than the\n",

experiments/README.md

@@ -8,7 +8,7 @@ These scripts are **standalone research code** — they import from `src/` but d
 ## `rl_finetuning/` — RL Fine-Tuning Ablation Suite
 
 Diagnoses why RL fine-tuning of the diffusion model collapses and which interventions fix it.
-Implements **25 ablations**: a baseline plus four groups (A: Regularisation, B: Training Signal, C: Architecture, D: Data Quality), with a comprehensive diagnostic and analysis pipeline.
+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.
 
@@ -20,27 +20,30 @@ rl_finetuning/
 ├── ablations/
 │   ├── losses.py             # All loss/objective variants as factory functions
 │   ├── optimizers.py         # LLRD, LoRA, gradient surgery, param masking
-│   ├── registry.py           # AblationSpec dataclass + REGISTRY (25 ablations)
+│   ├── 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/
-│   ├── action_distribution.py # Pre- vs post-finetuning action dist divergences + plots
-│   ├── plots.py              # All matplotlib figure generators
+│   ├── 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
+│   ├── report.py             # diagnosis.md + decision tree figure
+│   └── action_distribution.py  # Pre/post-RL action distribution analysis
 └── configs/
-    ├── ablations_default.yaml              # Base hyperparameters for every ablation run
-    ├── ablations_fast.yaml                 # Smoke-test overlay (50 iterations, 16 envs)
-    ├── ablations_final_craftax_classic_gpu_24gb.yaml    # Matches configs/final_craftax_classic_gpu_24gb.yaml   (RTX 3090 Ti, seed 42)
-    ├── ablations_final_craftax_classic_gpu_h200.yaml   # Matches configs/final_craftax_classic_gpu_h200.yaml  (H200, seed 43)
-    ├── ablations_final_craftax_gpu_24gb.yaml    # Matches configs/final_craftax_gpu_24gb.yaml   (24 GB reference machine (GPU model unrecorded),   seed 42)
-    └── ablations_final_craftax_gpu_h200.yaml   # Matches configs/final_craftax_gpu_h200.yaml  (H200, seed 43)
+    ├── 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
 ```
 
-`ablations_default.yaml` carries the transformer architecture of the released DAgger checkpoints — 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.
+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
 
@@ -69,9 +72,12 @@ jax_compilation_cache_dir: /var/tmp/your-user/jax-cache
 
 ### Usage
 
-The pretrained diffusion checkpoint can come from either offline training (`--mode offline`) or DAgger online training (`--mode online`). Use `--checkpoint` to point to it. For DAgger runs, either the final (`{env}-policy`) or best-validation (`{env}-policy-best`) artifact produced by `--mode online` can be consumed directly.
+`--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.
 
-Checkpoint paths accept `wandb:` prefixed artifact references (e.g., `wandb:team/project/artifact:latest`), which are downloaded automatically before training begins.
+**List all ablations:**
+```bash
+python experiments/rl_finetuning/run_ablations.py --list
+```
 
 **Smoke test (2 ablations, fast config):**
 ```bash
@@ -81,7 +87,7 @@ python experiments/rl_finetuning/run_ablations.py \
     --checkpoint $PRETRAINED_CKPT
 ```
 
-**Full suite (all 25 ablations):**
+**Full suite (all 26 ablations):**
 ```bash
 python experiments/rl_finetuning/run_ablations.py \
     --config configs/defaults.yaml \
@@ -94,7 +100,7 @@ python experiments/rl_finetuning/run_ablations.py \
 
 **Full suite against a pinned `final_*` checkpoint:**
 ```bash
-# Craftax Classic, 24 GB hardware (seed 42 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 \
@@ -109,13 +115,6 @@ python experiments/rl_finetuning/run_ablations.py \
     --checkpoint $PRETRAINED_CKPT
 ```
 
-**From a W&B artifact:**
-```bash
-python experiments/rl_finetuning/run_ablations.py \
-    --ablations baseline_rl \
-    --checkpoint wandb:my-team/remdm-planner-craftax/Craftax-Classic-Symbolic-v1-policy:latest
-```
-
 **Re-plot from saved results (no training):**
 ```bash
 python experiments/rl_finetuning/run_ablations.py \
@@ -143,7 +142,8 @@ poolable** with it:
 | `eval_steps` | 1024 | 512 | 1024 | 512 | noisier score |
 | `mixed_replay_buffer_size` | 20000 | 10000 | 10000 | 10000 | replay horizon |
 
-Differences in diagnostic cadence (`eval_every`, `cka_every`,
+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.
 
@@ -155,11 +155,6 @@ 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.
 
-**List all ablations:**
-```bash
-python experiments/rl_finetuning/run_ablations.py --list
-```
-
 ### Ablations
 
 | Group | Name | Tests |
@@ -177,6 +172,7 @@ python experiments/rl_finetuning/run_ablations.py --list
 | | `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 |
@@ -209,36 +205,62 @@ experiments/rl_finetuning/outputs/{run_id}/
 │   ├── 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/
+│   └── 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
+    ├── 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:**
 ```json
 {
   "pretrained_score": 0.1234,
   "pretrained_ach_rates": {"achievement_collect_wood": 0.42, ...},
-  "config": {"MAX_ITER": 1000, ...},
+  "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
     }
   }
@@ -246,14 +268,99 @@ experiments/rl_finetuning/outputs/{run_id}/
 ```
 
 `results.json` is written incrementally after each ablation completes — a partial file with
-N of 25 ablations is fully valid and loadable by `--analyze-only --results-path`.
+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):**
+```bash
+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`:
+```bash
+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
 
-All metrics are logged under `ablations/{method_name}/{metric}`, e.g.:
-- `ablations/kl_penalty/eval_score`
-- `ablations/gradient_surgery/grad_align`
-- `ablations/ewc/repr_drift_kl`
+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
 

experiments/rl_finetuning/ablations/losses.py

@@ -380,6 +380,35 @@ def make_loss_bc_wins(ctx: LossContext) -> LossFn:
     return loss_fn
 
 
+def make_loss_bc_all(ctx: LossContext) -> LossFn:
+    """Uniform ELBO over every rollout window (BC on self-generated data).
+
+    The return weighting is dropped outright: ``advantages`` is ignored and
+    ``compute_loss`` takes its plain batch mean, which is exactly a weight of
+    1.0 on every window. The batch is the one ``baseline_rl`` trains on --
+    same collection, same windows -- so the pair separates the on-policy data
+    distribution from the weighting applied to it, which no other arm does.
+
+    ``bc_wins`` is the nearest neighbour and does not answer this: it keeps
+    only the winning windows, so it varies the data as well as the weights.
+
+    Hypothesis: if this degrades like ``baseline_rl``, fine-tuning on
+    self-generated rollouts is the cause and the weighting is incidental.
+
+    Args:
+        ctx: Shared loss context.
+
+    Returns:
+        ``LossFn`` averaging uniformly over the whole batch.
+    """
+
+    def loss_fn(params, acts, obs, valid, rng, advantages):
+        del advantages  # uniform weights: the point of the arm
+        return _core_loss(ctx, params, rng, acts, obs, valid, None)
+
+    return loss_fn
+
+
 def make_loss_low_t(ctx: LossContext) -> LossFn:
     """Return-weighted ELBO restricted to t ∈ [ε, t_max_low].
 

experiments/rl_finetuning/ablations/registry.py

@@ -16,6 +16,7 @@ from experiments.rl_finetuning.ablations.losses import (
     LossFn,
     make_loss_advantage_clip,
     make_loss_baseline,
+    make_loss_bc_all,
     make_loss_bc_wins,
     make_loss_entropy_bonus,
     make_loss_ewc,
@@ -303,6 +304,15 @@ REGISTRY: dict[str, AblationSpec] = {
         optimizer_factory=_std_opt,
         wins_only=True,
     ),
+    "bc_all": AblationSpec(
+        name="bc_all",
+        group="B",
+        description="Uniform ELBO on all rollout windows (no advantage weighting)",
+        hypothesis="If it degrades like baseline_rl: the self-generated rollouts, "
+        "not the return weighting, are the cause",
+        loss_factory=make_loss_bc_all,
+        optimizer_factory=_std_opt,
+    ),
     "low_t": AblationSpec(
         name="low_t",
         group="B",

experiments/rl_finetuning/ablations/training.py

@@ -1549,7 +1549,7 @@ def run_ablation(
     rng: jax.Array,
     wandb_run: Any = None,
     output_dir: Any = None,
-) -> tuple[AblationHistory, float, Any]:
+) -> tuple[AblationHistory, float, Any, dict[str, float]]:
     """Run one complete ablation (Python-level wrapper around make_run_ablation).
 
     Handles EWC Fisher estimation (which requires rollouts before JIT),
@@ -1572,7 +1572,9 @@ def run_ablation(
         output_dir:        Optional Path for per-iteration checkpoint saving.
 
     Returns:
-        Tuple of ``(history, final_score, final_params)``.
+        Tuple of ``(history, final_score, final_params, final_ach)``, where
+        ``final_ach`` is the per-achievement unlock rate of the same post-loop
+        evaluation that produced ``final_score``.
     """
     logger.info("=" * 60)
     logger.info("ABLATION: %s  [Group %s]", spec.name, spec.group)
@@ -1684,13 +1686,14 @@ def run_ablation(
         final_params = jax.device_get(final_carry.ema_params)
     final_score = float(final_info.get("returned_episode_returns", jnp.array(0.0)))
 
-    # Extract per-achievement unlock rates from final eval
+    # Per-achievement detail of the *same* evaluation that produced
+    # `final_score`. It is returned separately rather than written over the
+    # last entry of `history.per_achievement_rates`: that entry belongs to the
+    # last in-loop evaluation, a different draw, and overwriting it made the
+    # achievement tables and the headline score disagree by construction.
     final_ach = {
         k: float(v) / 100.0 for k, v in final_info.items() if "achievement" in k.lower()
     }
-    # Overwrite the empty dicts in history with final eval achievements
-    if history.per_achievement_rates:
-        history.per_achievement_rates[-1] = final_ach
 
     logger.info("  [%s] FINAL score: %.4f", spec.name, final_score)
 
@@ -1713,4 +1716,4 @@ def run_ablation(
                 }
             )
 
-    return history, final_score, final_params
+    return history, final_score, final_params, final_ach

experiments/rl_finetuning/analysis/gdelta.py

@@ -0,0 +1,548 @@
+"""Measure the return term g_delta of the return-weighted ELBO decomposition.
+
+Implements the measurement behind Eq. 8 and Table 6 of the paper. Writing the
+per-window weight as A_i and the batch mean as Abar, the training gradient
+decomposes exactly as
+
+    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,
+
+so g_delta carries the entire directional contribution of the return and Abar
+is the scalar step-size rescaling. :func:`measure` loads a 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 every weight transform the ablation suite uses, and reports
+three references the cosine column needs:
+
+  * the random-direction null, cos ~ N(0, 1/sqrt(D)) for D parameters;
+  * cos(grad L_BC, grad L_BC) across two independent noise draws, which is the
+    value a direction attains when it *is* the imitation direction;
+  * the shuffled-delta null, which permutes delta across the batch. This keeps
+    the multiset of weights and hence CV_A, and destroys the association
+    between a window's weight and that window's gradient. Any part of the
+    measured ratio and cosine that survives the shuffle is batch
+    heterogeneity, not return signal.
+
+No training and no optimiser step occur. Runs on CPU in a few minutes.
+
+The per-draw standard deviations a single seed reports are *within* one rollout
+seed. The figure Table 6 prints is the standard deviation across rollout seeds,
+which :func:`aggregate` computes over the per-seed records.
+
+Driven by ``run_ablations.py --measure-gdelta``; see :func:`run_gdelta`.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import numpy as np
+import orjson
+
+logger = logging.getLogger(__name__)
+
+# The variant -> ablation mapping this module assumes. Each entry names the
+# registry key, the loss factory that key must still use, and its wins_only
+# flag. verify_registry() fails loudly if the suite drifts away from it, so a
+# registry edit cannot silently desynchronise the measurement.
+REGISTRY_RULES = {
+    "baseline_clipped_ratio": ("baseline_rl", "make_loss_baseline", False),
+    "advantage_clip": ("advantage_clip", "make_loss_advantage_clip", False),
+    "normalized_adv": ("normalized_adv", "make_loss_normalized_adv", False),
+    "bc_wins": ("bc_wins", "make_loss_bc_wins", True),
+}
+
+BASELINE = "baseline_clipped_ratio"
+
+GDELTA_DIRNAME = "gdelta"
+AGGREGATE_FILENAME = "gdelta_aggregate.json"
+
+
+class RegistryDriftError(RuntimeError):
+    """The ablation registry no longer matches the variants measured here."""
+
+
+def verify_registry() -> None:
+    """Fail if the ablation registry no longer matches the assumed variants."""
+    from experiments.rl_finetuning.ablations.registry import REGISTRY
+
+    for variant, (name, factory, wins_only) in REGISTRY_RULES.items():
+        spec = REGISTRY.get(name)
+        if spec is None:
+            raise RegistryDriftError(
+                f"registry has no ablation {name!r}; variant {variant!r} is stale"
+            )
+        if spec.loss_factory.__name__ != factory:
+            raise RegistryDriftError(
+                f"ablation {name!r} now uses {spec.loss_factory.__name__}, "
+                f"not {factory}; variant {variant!r} measures a weighting the "
+                "trainer no longer applies"
+            )
+        if bool(spec.wins_only) != wins_only:
+            raise RegistryDriftError(
+                f"ablation {name!r} has wins_only={spec.wins_only}, expected "
+                f"{wins_only}; variant {variant!r} is stale"
+            )
+
+
+def build_variants(adv, returns, cfg, batch):
+    """The weight vectors the suite's four weighting ablations apply.
+
+    ``bc_wins`` is the binary win mask rescaled by ``B / n_wins``, which is the
+    vector ``make_loss_bc_wins`` builds internally from the mask that
+    ``_compute_advantages(wins_only=True)`` hands it. The rescaling is what
+    turns ``compute_loss``'s batch mean into a uniform mean over the winners,
+    so it is part of the weight the trainer applies, not a normalisation added
+    here.
+    """
+    import jax.numpy as jnp
+
+    eps = cfg["ADV_CLIP_EPS"]
+    win = (returns > cfg["WIN_THRESHOLD"]).astype(jnp.float32)
+    n_win = win.sum()
+    return {
+        BASELINE: adv,
+        "advantage_clip": jnp.clip(adv, 1.0 - eps, 1.0 + eps),
+        "normalized_adv": (adv - adv.mean()) / (adv.std() + 1e-8),
+        "bc_wins": win * jnp.where(n_win > 0, batch / jnp.maximum(n_win, 1.0), 0.0),
+    }
+
+
+def centred_delta(weights):
+    """Return ``(delta, Abar, a1_holds)`` for a weight vector.
+
+    (A1) of the paper asks for non-negative weights with a strictly positive
+    mean. Where it holds, ``delta_i = A_i/Abar - 1`` and ``mean(delta) == 0``.
+    Where it does not -- ``normalized_adv`` mean-centres, so its weights are
+    signed and its mean vanishes -- the ratio is meaningless, and the caller is
+    told rather than handed a number that divides by a value near zero.
+    """
+    import jax.numpy as jnp
+
+    wbar = float(jnp.mean(weights))
+    scale = float(jnp.max(jnp.abs(weights))) + 1e-12
+    a1_holds = bool(jnp.all(weights >= 0.0)) and wbar > 1e-6 * scale
+    if not a1_holds:
+        return weights - wbar, wbar, False
+    return weights / wbar - 1.0, wbar, True
+
+
+def against_bc(g, g_bc, norm_bc):
+    """``(||g|| / ||g_BC||, cos(g, g_BC))`` for one gradient against imitation."""
+    import jax.numpy as jnp
+
+    norm = float(jnp.linalg.norm(g))
+    return norm / norm_bc, float(jnp.dot(g, g_bc) / (norm * norm_bc + 1e-12))
+
+
+def effective_sample_size(weights) -> float:
+    """ESS as a fraction of the batch: (sum A)^2 / (B sum A^2)."""
+    import jax.numpy as jnp
+
+    total = float(jnp.sum(weights))
+    sq = float(jnp.sum(weights ** 2))
+    if sq <= 0.0:
+        return float("nan")
+    return total ** 2 / (sq * weights.shape[0])
+
+
+def restore_params(net, rng, obs_dim, plan_horizon, ckpt: str):
+    """Restore Orbax parameters with an explicit single-device sharding.
+
+    ``model.load_checkpoint`` builds its restore target with ``jax.eval_shape``,
+    whose leaves carry no sharding. Orbax >= 0.12 rejects that when the
+    checkpoint was written under a different device topology, which is the
+    common case for a checkpoint trained on GPU and inspected on CPU. Attaching
+    a concrete sharding here is the only difference from the library path.
+    """
+    import jax
+    import orbax.checkpoint as ocp
+
+    from src.planners.model import _validate_restored_tree, abstract_params
+
+    sharding = jax.sharding.SingleDeviceSharding(jax.devices()[0])
+    abstract = abstract_params(net, rng, obs_dim, plan_horizon)
+    restore_args = jax.tree.map(
+        lambda x: ocp.ArrayRestoreArgs(sharding=sharding, dtype=x.dtype), abstract
+    )
+    with ocp.CheckpointManager(str(Path(ckpt).resolve())) as mgr:
+        step = mgr.latest_step()
+        restored = mgr.restore(
+            step,
+            args=ocp.args.PyTreeRestore(
+                item={"params": abstract},
+                restore_args={"params": restore_args},
+                partial_restore=True,
+            ),
+        )
+    params = restored["params"]
+    _validate_restored_tree(params, abstract, ckpt)
+    return params, step
+
+
+def aggregate(blobs: list[dict], inputs: list[str] | None = None) -> dict:
+    """Combine per-seed records, reporting the standard deviation across seeds.
+
+    Each input contributes one number per variant per column -- its own mean
+    over the ``(z_t, t)`` draws. The dispersion reported here is over those
+    per-seed means, which is the quantity the paper's table claims.
+
+    Args:
+        blobs:  Per-seed records, as returned by :func:`measure`.
+        inputs: Optional provenance paths recorded in the output.
+
+    Returns:
+        The aggregate record.
+    """
+    labels = inputs if inputs is not None else [
+        f"seed{b.get('seed', i)}" for i, b in enumerate(blobs)
+    ]
+    for blob, label in zip(blobs, labels, strict=True):
+        if blob.get("aggregate"):
+            raise ValueError(f"{label} is already an aggregate")
+
+    names = list(blobs[0]["variants"])
+    for blob, label in zip(blobs, labels, strict=True):
+        if list(blob["variants"]) != names:
+            raise ValueError(f"{label} has a different variant set")
+
+    def across(values):
+        arr = np.array(values, dtype=float)
+        return float(arr.mean()), float(arr.std())
+
+    out = {
+        "aggregate": True,
+        "inputs": [str(p) for p in labels],
+        "seeds": [int(b["seed"]) for b in blobs],
+        "n_seeds": len(blobs),
+        "n_draws_per_seed": [int(b["n_draws"]) for b in blobs],
+        "n_params": int(blobs[0]["n_params"]),
+        "random_cos_sd": float(blobs[0]["random_cos_sd"]),
+        "batch": int(blobs[0]["batch"]),
+        "eq4_residual_max": float(max(b["eq4_residual_max"] for b in blobs)),
+        "variants": {},
+    }
+    out["bc_self_cos_mean"], out["bc_self_cos_std_seeds"] = across(
+        [b["bc_self_cos_mean"] for b in blobs]
+    )
+
+    columns = [
+        "cv_a", "abar", "abar_ratio_to_baseline", "ess_fraction",
+        "ratio_mean", "cos_mean", "ratio_shuffled_mean", "cos_shuffled_mean",
+    ]
+    for name in names:
+        rec = {"a1_violated": any(b["variants"][name]["a1_violated"] for b in blobs)}
+        for col in columns:
+            mean, std = across([b["variants"][name][col] for b in blobs])
+            stem = col[:-5] if col.endswith("_mean") else col
+            rec[f"{stem}_mean"] = mean
+            rec[f"{stem}_std_seeds"] = std
+        out["variants"][name] = rec
+    return out
+
+
+def print_aggregate(agg: dict) -> None:
+    print(f"\naggregate over {agg['n_seeds']} rollout seeds {agg['seeds']}, "
+          f"{agg['n_draws_per_seed']} draws each; +/- is ACROSS SEEDS")
+    print(f"cos(grad L_BC, grad L_BC) = {agg['bc_self_cos_mean']:.3f} "
+          f"+/- {agg['bc_self_cos_std_seeds']:.3f}")
+    print(f"Eq. 4 identity, max relative residual = {agg['eq4_residual_max']:.2e}\n")
+    header = (f"{'weight transform':26s} {'CV_A':>7s} {'Abar':>8s} {'Abar/base':>10s} "
+              f"{'ESS/B':>7s} {'ratio':>16s} {'cos':>16s} "
+              f"{'ratio(shuf)':>16s} {'cos(shuf)':>16s}")
+    print(header)
+    for name, rec in agg["variants"].items():
+        flag = "  [(A1) violated]" if rec["a1_violated"] else ""
+        print(f"{name:26s} {rec['cv_a_mean']:7.3f} {rec['abar_mean']:8.3f} "
+              f"{rec['abar_ratio_to_baseline_mean']:10.3f} "
+              f"{rec['ess_fraction_mean']:7.3f} "
+              f"{rec['ratio_mean']:9.3f} +/-{rec['ratio_std_seeds']:.3f} "
+              f"{rec['cos_mean']:+9.3f} +/-{rec['cos_std_seeds']:.3f} "
+              f"{rec['ratio_shuffled_mean']:9.3f} +/-{rec['ratio_shuffled_std_seeds']:.3f} "
+              f"{rec['cos_shuffled_mean']:+9.3f} +/-{rec['cos_shuffled_std_seeds']:.3f}"
+              f"{flag}")
+
+
+def measure(
+    config: dict,
+    ckpt: str,
+    *,
+    seed: int = 0,
+    n_draws: int = 8,
+    num_envs: int | None = None,
+    batch_size: int | None = None,
+) -> dict:
+    """Measure the decomposition at one rollout seed.
+
+    Args:
+        config:     UPPERCASE config dict; the run's own, so the weight
+                    transforms match the ones it trained under.
+        ckpt:       Orbax checkpoint directory of the pretrained planner.
+        seed:       Rollout seed. Also fixes the ``(z_t, t)`` draws.
+        n_draws:    Independent ``(z_t, t)`` draws to average over.
+        num_envs:   Override ``NUM_ENVS``; ``None`` keeps the config value.
+        batch_size: Override ``BATCH_SIZE``; ``None`` keeps the config value.
+
+    Returns:
+        The per-seed record.
+    """
+    import jax
+    import jax.numpy as jnp
+
+    from experiments.rl_finetuning.ablations.training import (
+        _compute_advantages,
+        build_rollout_fn,
+    )
+    from src.diffusion.loss import compute_loss
+    from src.diffusion.schedules import SCHEDULE_MAP
+    from src.planners.env import make_env
+    from src.planners.model import build_model, make_apply_fns
+
+    verify_registry()
+
+    cfg = dict(config)
+    if num_envs is not None:
+        cfg["NUM_ENVS"] = num_envs
+    if batch_size is not None:
+        cfg["BATCH_SIZE"] = batch_size
+
+    env, env_params = make_env(cfg, cfg["NUM_ENVS"])
+    num_actions = env.action_space(env_params).n
+    obs_dim = env.observation_space(env_params).shape[0]
+    cfg["NUM_ACTIONS"] = num_actions
+    schedule_fn, schedule_deriv_fn = SCHEDULE_MAP[cfg.get("DIFFUSION_SCHEDULE", "cosine")]
+
+    net = build_model(cfg, num_actions)
+    apply_eval, apply_train = make_apply_fns(net)
+
+    rng = jax.random.PRNGKey(seed)
+    # A separate stream for the shuffled-delta null, so adding the control
+    # leaves every draw of the real measurement bit-for-bit unchanged.
+    perm_rng = jax.random.PRNGKey(seed + 10_000)
+    rng, key = jax.random.split(rng)
+    params, step = restore_params(net, key, obs_dim, cfg["PLAN_HORIZON"], ckpt)
+    n_params = sum(int(np.prod(x.shape)) for x in jax.tree.leaves(params))
+    random_cos_sd = 1.0 / np.sqrt(n_params)
+    print(f"checkpoint step {step}, D = {n_params/1e6:.2f}M, "
+          f"random-cosine null sd = {random_cos_sd:.2e}", flush=True)
+
+    # ---- one on-policy batch from the pretrained policy ----
+    collect = build_rollout_fn(env, env_params, apply_eval, cfg, obs_dim)
+    rng, key = jax.random.split(rng)
+    obs0, state0 = env.reset(key, env_params)
+    done0 = jnp.zeros(cfg["NUM_ENVS"], dtype=bool)
+    rng, key = jax.random.split(rng)
+    _, _, _, _, f_obs, f_acts, f_valid, f_ret, _ = collect(
+        params, state0, obs0, done0, key
+    )
+    adv, _, _ = _compute_advantages(
+        f_ret,
+        cfg["RETURN_WEIGHT_FLOOR"],
+        cfg["RETURN_WEIGHT_CAP"],
+        wins_only=False,
+        win_thresh=cfg["WIN_THRESHOLD"],
+        use_running_stats=False,
+        ema_decay=0.99,
+        running_mean=jnp.array(0.0),
+        running_std=jnp.array(1.0),
+    )
+
+    batch = min(cfg["BATCH_SIZE"], f_obs.shape[0])
+    rng, key = jax.random.split(rng)
+    idx = jax.random.permutation(key, f_obs.shape[0])[:batch]
+    obs_b, act_b, val_b = f_obs[idx], f_acts[idx], f_valid[idx]
+    adv_b, ret_b = adv[idx], f_ret[idx]
+    print(f"batch {batch} windows, win rate {float((ret_b > cfg['WIN_THRESHOLD']).mean()):.3f}",
+          flush=True)
+
+    variants = build_variants(adv_b, ret_b, cfg, batch)
+
+    def flatten(tree):
+        return jnp.concatenate([x.ravel() for x in jax.tree.leaves(tree)])
+
+    def gradient(weights, key):
+        def loss(p):
+            value, _ = compute_loss(
+                apply_train, p, key, act_b, obs_b, val_b, num_actions,
+                schedule_fn, schedule_deriv_fn,
+                sigma_t=cfg.get("TRAIN_SIGMA", 0.0),
+                label_smoothing=cfg.get("LABEL_SMOOTHING", 0.0),
+                advantages=weights,
+            )
+            return value
+        return flatten(jax.grad(loss)(params))
+
+    deltas = {}
+    stats = {}
+    for name, weights in variants.items():
+        delta, wbar, a1_holds = centred_delta(weights)
+        deltas[name] = delta
+        stats[name] = {
+            "cv_a": float(jnp.sqrt(jnp.mean(delta ** 2))),
+            "abar": wbar,
+            "ess_fraction": effective_sample_size(weights),
+            "a1_violated": not a1_holds,
+        }
+    base_abar = stats[BASELINE]["abar"]
+    for rec in stats.values():
+        rec["abar_ratio_to_baseline"] = rec["abar"] / base_abar
+
+    acc = {name: {"ratio": [], "cos": [], "ratio_shuf": [], "cos_shuf": []}
+           for name in variants}
+    bc_self, residuals = [], []
+    for draw in range(n_draws):
+        rng, key = jax.random.split(rng)
+        g_bc = gradient(None, key)
+        norm_bc = float(jnp.linalg.norm(g_bc))
+
+        rng, key2 = jax.random.split(rng)
+        g_bc2 = gradient(None, key2)  # same objective, independent noise draw
+        bc_self.append(
+            float(jnp.dot(g_bc, g_bc2) / (norm_bc * jnp.linalg.norm(g_bc2) + 1e-12))
+        )
+
+        for name, weights in variants.items():
+            delta = deltas[name]
+            g_delta = gradient(delta, key)
+            ratio, cos = against_bc(g_delta, g_bc, norm_bc)
+            acc[name]["ratio"].append(ratio)
+            acc[name]["cos"].append(cos)
+
+            # Shuffled-delta null: same multiset of weights, same CV_A, the
+            # association with each window's own return destroyed. The (z_t, t)
+            # key is the one the real measurement used, so the two differ only
+            # in which window carries which weight.
+            perm_rng, perm_key = jax.random.split(perm_rng)
+            shuffled = jax.random.permutation(perm_key, delta)
+            ratio_s, cos_s = against_bc(gradient(shuffled, key), g_bc, norm_bc)
+            acc[name]["ratio_shuf"].append(ratio_s)
+            acc[name]["cos_shuf"].append(cos_s)
+
+            if name == BASELINE:
+                # Eq. 4 identity check: grad L_RW == Abar * (grad L_BC + g_delta)
+                g_rw = gradient(weights, key)
+                residuals.append(float(
+                    jnp.linalg.norm(g_rw - stats[name]["abar"] * (g_bc + g_delta))
+                    / (jnp.linalg.norm(g_rw) + 1e-12)
+                ))
+        print(f"  draw {draw + 1}/{n_draws}", flush=True)
+
+    out = {
+        "aggregate": False,
+        "checkpoint_step": int(step),
+        "n_params": int(n_params),
+        "random_cos_sd": float(random_cos_sd),
+        "batch": int(batch),
+        "seed": seed,
+        "n_draws": n_draws,
+        "bc_self_cos_mean": float(np.mean(bc_self)),
+        "bc_self_cos_std": float(np.std(bc_self)),
+        "eq4_residual_max": float(np.max(residuals)),
+        "variants": {},
+    }
+    print(f"\ncos(grad L_BC, grad L_BC) across draws = "
+          f"{np.mean(bc_self):.3f} +/- {np.std(bc_self):.3f}   [same-objective reference]")
+    print(f"random-direction null: cos ~ N(0, {random_cos_sd:.2e})")
+    print(f"Eq. 4 identity, max relative residual = {np.max(residuals):.2e}")
+    print(f"+/- below is across the {n_draws} (z_t, t) draws of this one seed, "
+          "not across seeds\n")
+    print(f"{'weight transform':26s} {'CV_A':>7s} {'Abar':>8s} {'Abar/base':>10s} "
+          f"{'ESS/B':>7s} {'ratio':>16s} {'cos':>16s} "
+          f"{'ratio(shuf)':>16s} {'cos(shuf)':>16s}")
+    for name, rec in acc.items():
+        ratio, cos = np.array(rec["ratio"]), np.array(rec["cos"])
+        ratio_s, cos_s = np.array(rec["ratio_shuf"]), np.array(rec["cos_shuf"])
+        out["variants"][name] = {
+            **stats[name],
+            "ratio_mean": float(ratio.mean()), "ratio_std_draws": float(ratio.std()),
+            "cos_mean": float(cos.mean()), "cos_std_draws": float(cos.std()),
+            "ratio_shuffled_mean": float(ratio_s.mean()),
+            "ratio_shuffled_std": float(ratio_s.std()),
+            "cos_shuffled_mean": float(cos_s.mean()),
+            "cos_shuffled_std": float(cos_s.std()),
+        }
+        flag = "  [(A1) violated]" if stats[name]["a1_violated"] else ""
+        print(f"{name:26s} {stats[name]['cv_a']:7.3f} {stats[name]['abar']:8.3f} "
+              f"{stats[name]['abar_ratio_to_baseline']:10.3f} "
+              f"{stats[name]['ess_fraction']:7.3f} "
+              f"{ratio.mean():9.3f} +/-{ratio.std():.3f} "
+              f"{cos.mean():+9.3f} +/-{cos.std():.3f} "
+              f"{ratio_s.mean():9.3f} +/-{ratio_s.std():.3f} "
+              f"{cos_s.mean():+9.3f} +/-{cos_s.std():.3f}{flag}")
+    return out
+
+
+def _write(path: Path, blob: dict) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_bytes(orjson.dumps(blob, option=orjson.OPT_INDENT_2))
+    logger.info("Wrote %s", path)
+
+
+def load_gdelta_aggregate(output_dir: Path) -> dict | None:
+    """The aggregate record for a run, or ``None`` if it was never measured."""
+    path = Path(output_dir) / GDELTA_DIRNAME / AGGREGATE_FILENAME
+    if not path.is_file():
+        return None
+    return orjson.loads(path.read_bytes())
+
+
+def run_gdelta(
+    config: dict,
+    ckpt: str,
+    output_dir: Path,
+    *,
+    seeds: list[int],
+    n_draws: int = 8,
+    num_envs: int | None = None,
+    batch_size: int | None = None,
+) -> dict:
+    """Measure every seed and write the per-seed and aggregate records.
+
+    Artifacts land in ``output_dir/gdelta/``, beside the run's ``results.json``,
+    so one run directory holds the suite's scores and the gradient measurement
+    taken at the checkpoint they started from.
+
+    Args:
+        config:     UPPERCASE config dict.
+        ckpt:       Orbax checkpoint directory of the pretrained planner.
+        output_dir: The run's root output directory.
+        seeds:      Rollout seeds to measure.
+        n_draws:    Independent ``(z_t, t)`` draws per seed.
+        num_envs:   Override ``NUM_ENVS``; ``None`` keeps the config value.
+        batch_size: Override ``BATCH_SIZE``; ``None`` keeps the config value.
+
+    Returns:
+        The aggregate record.
+    """
+    gdelta_dir = Path(output_dir) / GDELTA_DIRNAME
+    blobs, labels = [], []
+    for seed in seeds:
+        logger.info("Measuring g_delta at rollout seed %d", seed)
+        blob = measure(
+            config,
+            ckpt,
+            seed=seed,
+            n_draws=n_draws,
+            num_envs=num_envs,
+            batch_size=batch_size,
+        )
+        path = gdelta_dir / f"gdelta_seed{seed}.json"
+        _write(path, blob)
+        blobs.append(blob)
+        labels.append(str(path))
+
+    agg = aggregate(blobs, labels)
+    _write(gdelta_dir / AGGREGATE_FILENAME, agg)
+    print_aggregate(agg)
+    return agg
+
+
+def aggregate_files(paths: list[str], output_dir: Path) -> dict:
+    """Aggregate per-seed records written elsewhere, e.g. by another machine."""
+    blobs = [orjson.loads(Path(p).read_bytes()) for p in paths]
+    agg = aggregate(blobs, [str(p) for p in paths])
+    _write(Path(output_dir) / GDELTA_DIRNAME / AGGREGATE_FILENAME, agg)
+    print_aggregate(agg)
+    return agg

experiments/rl_finetuning/analysis/plots.py

@@ -791,10 +791,12 @@ def plot_achievement_breakdown(
                               names, values in [0, 1]).
         output_dir:           Output directory for the figure.
     """
-    # Collect ablations that have at least two eval checkpoints.
+    # Collect ablations with at least two *populated* eval checkpoints. Empty
+    # entries are checkpoints where the scan could not return variable-key
+    # dicts, not evaluations that unlocked nothing.
     valid: list[tuple[str, dict[str, float], dict[str, float]]] = []
     for name, res in results.items():
-        rates = res["history"].per_achievement_rates
+        rates = [r for r in res["history"].per_achievement_rates if r]
         if len(rates) >= 2:
             valid.append((name, rates[0], rates[-1]))
 

experiments/rl_finetuning/analysis/tables.py

@@ -15,6 +15,10 @@ import polars as pl
 
 from experiments.rl_finetuning.ablations.registry import REGISTRY
 from experiments.rl_finetuning.ablations.training import AblationHistory
+from experiments.rl_finetuning.analysis.gdelta import (
+    REGISTRY_RULES,
+    load_gdelta_aggregate,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -100,6 +104,26 @@ def baseline_rl_score_of(results: dict[str, dict], pretrained_score: float) -> f
     return results.get("baseline_rl", {}).get("score", pretrained_score)
 
 
+def final_ach_rates(res: dict) -> dict[str, float]:
+    """Per-achievement rates of the evaluation that produced ``score``.
+
+    ``final_ach_rates`` is the detail of the post-loop evaluation the headline
+    score comes from, so tables built on it reconcile with the score by
+    construction. Suite runs from before that field existed wrote the same
+    evaluation over the last in-loop history entry; those are read from there.
+
+    Args:
+        res: One entry of a results dict.
+
+    Returns:
+        ``{achievement: rate}``, empty when neither source has any.
+    """
+    if res.get("final_ach_rates"):
+        return res["final_ach_rates"]
+    rates = res["history"].per_achievement_rates
+    return rates[-1] if rates else {}
+
+
 def _latex_escape(text: str) -> str:
     """Escape LaTeX special characters in a string.
 
@@ -267,13 +291,13 @@ def write_significance_test(results: dict[str, dict], out_dir: Path) -> None:
         f"bootstrap 95% CI of the difference (10000 resamples, seed 0): "
         f"[{lo:.4f}, {hi:.4f}]\n"
         + (
-            f"arms excluded from the max (seed count != {n_b}): "
-            f"{', '.join(dropped)}\n"
+            f"arms excluded from the max (seed count != {n_b}): {', '.join(dropped)}\n"
             if dropped
             else ""
         )
     )
 
+
 def make_main_results_table(
     results: dict[str, dict],
     pretrained_score: float,
@@ -369,6 +393,47 @@ def make_gradient_analysis_table(
     return pl.DataFrame(rows)
 
 
+def make_gdelta_table(agg: dict) -> pl.DataFrame:
+    """The gradient-decomposition table: one row per weight transform.
+
+    Columns follow the paper's ``tab:gdelta``. Every ``+/-`` is the standard
+    deviation ACROSS ROLLOUT SEEDS, which is what :func:`gdelta.aggregate`
+    computes; a single seed's dispersion over its own ``(z_t, t)`` draws is a
+    different and smaller quantity and is not reported here.
+
+    ``normalized_adv`` mean-centres its weights, so (A1) fails and its ratio
+    divides by a mean near zero. The row is kept with an explicit flag rather
+    than dropped, because its absence would read as an omission.
+
+    Args:
+        agg: The aggregate record from ``gdelta/gdelta_aggregate.json``.
+
+    Returns:
+        Polars DataFrame.
+    """
+    rows = []
+    for name, rec in agg["variants"].items():
+        rows.append(
+            {
+                "Transform": name,
+                "CV_A": round(rec["cv_a_mean"], 3),
+                "Abar": round(rec["abar_mean"], 3),
+                "Abar_Ratio_Baseline": round(rec["abar_ratio_to_baseline_mean"], 3),
+                "ESS_Fraction": round(rec["ess_fraction_mean"], 3),
+                "Ratio": round(rec["ratio_mean"], 3),
+                "Ratio_Sd": round(rec["ratio_std_seeds"], 3),
+                "Cos": round(rec["cos_mean"], 3),
+                "Cos_Sd": round(rec["cos_std_seeds"], 3),
+                "Ratio_Shuffled": round(rec["ratio_shuffled_mean"], 3),
+                "Ratio_Shuffled_Sd": round(rec["ratio_shuffled_std_seeds"], 3),
+                "Cos_Shuffled": round(rec["cos_shuffled_mean"], 3),
+                "Cos_Shuffled_Sd": round(rec["cos_shuffled_std_seeds"], 3),
+                "A1_Violated": bool(rec["a1_violated"]),
+            }
+        )
+    return pl.DataFrame(rows)
+
+
 def make_t_distribution_table(
     results: dict[str, dict],
 ) -> pl.DataFrame:
@@ -592,8 +657,7 @@ def make_per_env_table(
     """
     ablation_finals: dict[str, dict[str, float]] = {}
     for name, res in results.items():
-        rates = res["history"].per_achievement_rates
-        ablation_finals[name] = rates[-1] if rates else {}
+        ablation_finals[name] = final_ach_rates(res)
 
     all_keys: list[str] = sorted(
         set(pretrained_ach_rates) | {k for d in ablation_finals.values() for k in d}
@@ -645,9 +709,7 @@ def make_hypothesis_verdict_table(
         result = verdict(score, baseline_rl_score, pretrained_score)
         conclusion = {
             "IMPROVEMENT": "Hypothesis SUPPORTED — this intervention helps",
-            "COLLAPSE": (
-                "Hypothesis REFUTED — intervention did not prevent collapse"
-            ),
+            "COLLAPSE": ("Hypothesis REFUTED — intervention did not prevent collapse"),
             "NEUTRAL": "Inconclusive — no significant change",
         }[result]
 
@@ -685,8 +747,7 @@ def make_achievement_table(
     # Collect final achievement rates for every ablation.
     ablation_finals: dict[str, dict[str, float]] = {}
     for name, res in results.items():
-        rates = res["history"].per_achievement_rates
-        ablation_finals[name] = rates[-1] if rates else {}
+        ablation_finals[name] = final_ach_rates(res)
 
     # Union of all achievement keys.
     all_keys: list[str] = sorted(
@@ -709,11 +770,262 @@ def make_achievement_table(
     return pl.DataFrame(rows)
 
 
+_DIGIT_WORDS = (
+    "Zero",
+    "One",
+    "Two",
+    "Three",
+    "Four",
+    "Five",
+    "Six",
+    "Seven",
+    "Eight",
+    "Nine",
+)
+
+
+def _macro_name(*parts: str) -> str:
+    """A TeX-legal control sequence name built from arbitrary identifiers.
+
+    ``\\newcommand`` names may contain letters only, so underscores and
+    hyphens become word boundaries and digits are spelled out:
+    ``layer_ablation_top1`` -> ``LayerAblationTopOne``.
+    """
+    out = []
+    for part in parts:
+        for word in str(part).replace("-", "_").split("_"):
+            if not word:
+                continue
+            chars = [_DIGIT_WORDS[int(c)] if c.isdigit() else c for c in word]
+            joined = "".join(chars)
+            out.append(joined[0].upper() + joined[1:])
+    return "".join(out)
+
+
+def _cv_a(ess: float, batch: int) -> float:
+    """Weight dispersion recovered from an effective sample size."""
+    if not batch or ess <= 0:
+        return float("nan")
+    return float(np.sqrt(max(batch / ess - 1.0, 0.0)))
+
+
+def write_gdelta_table(agg: dict, output_dir: Path) -> pl.DataFrame:
+    """Write ``tables/gdelta.{csv,tex}`` for one aggregate record."""
+    df = make_gdelta_table(agg)
+    _save_table(
+        df,
+        Path(output_dir) / "tables" / "gdelta",
+        caption="Return-term decomposition per weight transform.",
+        label="tab:gdelta",
+    )
+    return df
+
+
+def _define_gdelta_macros(agg: dict, define) -> None:
+    """Emit the measured decomposition quantities as ``\\newcommand``s.
+
+    Each variant is tagged by its ABLATION name, not its gdelta variant name,
+    so ``baseline_clipped_ratio`` becomes ``BaselineRl`` and sits beside the
+    ``Score`` and ``Ess`` macros for the same condition. The ``Gdelta`` infix
+    keeps these clear of the ESS-derived ``CvA`` macros, which measure a
+    different thing on different batches.
+    """
+    define("GdeltaNSeeds", value=str(int(agg["n_seeds"])))
+    define("GdeltaNParams", value=f"{agg['n_params'] / 1e6:.2f}")
+    define("GdeltaRandomCosSd", value=f"{agg['random_cos_sd']:.1e}")
+    define("GdeltaBatch", value=str(int(agg["batch"])))
+    define("GdeltaBcSelfCos", value=f"{agg['bc_self_cos_mean']:.3f}")
+    define("GdeltaBcSelfCosSd", value=f"{agg['bc_self_cos_std_seeds']:.3f}")
+    define("GdeltaEqFourResidual", value=f"{agg['eq4_residual_max']:.1e}")
+
+    for variant, rec in agg["variants"].items():
+        rule = REGISTRY_RULES.get(variant)
+        tag = _macro_name(rule[0] if rule else variant)
+        define("GdeltaCvA", tag=tag, value=f"{rec['cv_a_mean']:.2f}")
+        define("GdeltaAbar", tag=tag, value=f"{rec['abar_mean']:.3f}")
+        define(
+            "GdeltaAbarRatio", tag=tag,
+            value=f"{rec['abar_ratio_to_baseline_mean']:.3f}",
+        )
+        define("GdeltaEss", tag=tag, value=f"{rec['ess_fraction_mean']:.2f}")
+        define("GdeltaRatio", tag=tag, value=f"{rec['ratio_mean']:.3f}")
+        define("GdeltaRatioSd", tag=tag, value=f"{rec['ratio_std_seeds']:.3f}")
+        define("GdeltaCos", tag=tag, value=f"{rec['cos_mean']:.3f}")
+        define("GdeltaCosSd", tag=tag, value=f"{rec['cos_std_seeds']:.3f}")
+        define(
+            "GdeltaRatioShuf", tag=tag,
+            value=f"{rec['ratio_shuffled_mean']:.3f}",
+        )
+        define(
+            "GdeltaRatioShufSd", tag=tag,
+            value=f"{rec['ratio_shuffled_std_seeds']:.3f}",
+        )
+        define("GdeltaCosShuf", tag=tag, value=f"{rec['cos_shuffled_mean']:.3f}")
+        define(
+            "GdeltaCosShufSd", tag=tag,
+            value=f"{rec['cos_shuffled_std_seeds']:.3f}",
+        )
+
+
+def write_tex_macros(
+    results: dict[str, dict],
+    pretrained_score: float,
+    out_path: Path,
+    config: dict | None = None,
+    prefix: str = "rw",
+    scale: float = 1.0,
+    gdelta: dict | None = None,
+) -> Path:
+    """Write ``\\newcommand`` definitions for every headline quantity.
+
+    The paper workspace expects generated numbers to reach the manuscript
+    through macros so each is traceable to the run that produced it, rather
+    than being retyped as a literal.
+
+    Each quantity is emitted at the number of decimals the manuscript prints
+    it with, not at a uniform significant-figure count. A macro is a literal
+    substitution: ``%.4g`` would render ``\\rwGroupMeanA`` as ``9.744`` where
+    the paper prints ``9.74``, silently changing the displayed precision at
+    every site. Scores, standard deviations and group statistics take two
+    decimals; effective sample size takes none, which is how Section 6.4
+    quotes it.
+
+    Args:
+        results:          Dict mapping ablation_name -> result entry.
+        pretrained_score: Pretrained model eval score.
+        out_path:         File to write, conventionally ``results.tex``.
+        config:           The run's config, read for ``BATCH_SIZE`` so
+                          :math:`\\mathrm{CV}_A` and ESS can be emitted.
+        prefix:           Macro-name prefix, keeping these out of the way of
+                          LaTeX's own names. Two suites sharing one
+                          ``results.tex`` need different prefixes.
+        gdelta:           Optional aggregate record from
+                          ``gdelta/gdelta_aggregate.json``. When present, the
+                          measured decomposition quantities are emitted under a
+                          ``Gdelta`` infix. They are deliberately NOT merged
+                          into the ``CvA`` macros above: those recover
+                          :math:`\\mathrm{CV}_A` from the ESS logged during
+                          training, whereas these are measured on the
+                          measurement batches at the pretrained checkpoint. The
+                          two differ, and the manuscript quotes both.
+        scale:            Multiplier on score-like quantities before
+                          formatting, for a suite whose metric is a fraction
+                          the paper reports as a percentage. Never applied to
+                          ESS or :math:`\\mathrm{CV}_A`, which are unitless.
+
+    Returns:
+        The path written.
+    """
+    cfg = config or {}
+    batch = int(cfg.get("BATCH_SIZE") or cfg.get("batch_size") or 0)
+
+    def score(value: float, decimals: int = 2) -> str:
+        return f"{float(value) * scale:.{decimals}f}"
+
+    lines = [
+        "% Generated by experiments/rl_finetuning/analysis/tables.py.",
+        "% Do not hand-edit: regenerate with --emit-tex-macros.",
+        "",
+    ]
+    defined: dict[str, str] = {}
+
+    def define(*parts: str, value: str, tag: str = "") -> None:
+        """Emit one ``\\newcommand``, refusing to redefine a name.
+
+        The mangling in :func:`_macro_name` is lossy -- ``top1`` and
+        ``top_one`` both become ``TopOne`` -- so two conditions can collide.
+        Silently overwriting would put a number under a name that reads as
+        another condition's, which is exactly the failure the macro
+        convention exists to prevent.
+        """
+        macro = f"{prefix}{''.join(parts)}{tag}"
+        if macro in defined:
+            raise ValueError(
+                f"macro name collision: \\{macro} would be defined twice "
+                f"(existing value {defined[macro]!r}, new value {value!r}). "
+                f"_macro_name() mangles distinct identifiers onto one TeX "
+                f"control sequence; rename the condition or extend the rule."
+            )
+        defined[macro] = value
+        lines.append(f"\\newcommand{{\\{macro}}}{{{value}}}")
+
+    define("PretrainedScore", value=score(pretrained_score))
+    if batch:
+        define("BatchSize", value=str(batch))
+
+    within_seed_var: list[float] = []
+    baseline = (
+        float(results["baseline_rl"]["score"]) if "baseline_rl" in results else None
+    )
+    for name in sorted(results):
+        res = results[name]
+        tag = _macro_name(name)
+        define("Score", tag=tag, value=score(res["score"]))
+        # Magnitude only: the manuscript carries the sign as $-$ / $+$.
+        define(
+            "DeltaPretrained",
+            tag=tag,
+            value=score(abs(float(res["score"]) - pretrained_score)),
+        )
+        if baseline is not None:
+            define(
+                "DeltaBaseline",
+                tag=tag,
+                value=score(abs(float(res["score"]) - baseline)),
+            )
+        define("ScoreSd", tag=tag, value=score(res.get("score_std", 0.0)))
+        scores = [float(s) for s in res.get("all_scores", [])]
+        if len(scores) >= 2:
+            within_seed_var.append(float(np.var(scores, ddof=1)))
+
+        ess_series = res["history"].effective_batch_size
+        if ess_series:
+            ess = float(np.mean(ess_series))
+            define("Ess", tag=tag, value=f"{ess:.0f}")
+            if batch:
+                cv = float(np.mean([_cv_a(e, batch) for e in ess_series if e > 0]))
+                define("CvA", tag=tag, value=f"{cv:.2f}")
+
+    # Pooled within-condition sd across seeds: the seed noise a single
+    # condition's score carries, not the spread between conditions.
+    if within_seed_var:
+        pooled = float(np.sqrt(np.mean(within_seed_var)))
+        define("PooledSeedSd", value=score(pooled))
+
+    group_df = make_group_summary_table(results)
+    for row in group_df.iter_rows(named=True):
+        tag = _macro_name(row["Group"])
+        for col, suffix in (
+            ("Mean", "GroupMean"),
+            ("Best", "GroupBest"),
+            ("Worst", "GroupWorst"),
+            ("StdDev", "GroupSd"),
+        ):
+            define(suffix, tag=tag, value=score(row[col]))
+        define("GroupN", tag=tag, value=str(int(row["N"])))
+        if baseline is not None:
+            define(
+                "GroupDelta",
+                tag=tag,
+                value=score(abs(float(row["Mean"]) - baseline)),
+            )
+
+    if gdelta:
+        _define_gdelta_macros(gdelta, define)
+
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    out_path.write_text("\n".join(lines) + "\n")
+    logger.info("Saved %s (%d macros)", out_path, len(defined))
+    return out_path
+
+
 def generate_summary_tables(
     results: dict[str, dict],
     pretrained_score: float,
     output_dir: Path,
     pretrained_ach_rates: dict[str, float] | None = None,
+    emit_tex_macros: bool = False,
+    config: dict | None = None,
 ) -> dict[str, pl.DataFrame]:
     """Generate all summary tables and save to output_dir/tables/.
 
@@ -727,6 +1039,15 @@ def generate_summary_tables(
         pretrained_ach_rates: Optional per-achievement unlock rates for the pretrained
                               baseline (keys = achievement name, values in [0, 1]).
                               When provided, a per-achievement summary table is generated.
+        emit_tex_macros:      Also write ``tables/results.tex``, one
+                              ``\\newcommand`` per headline quantity, for the
+                              manuscript to cite instead of a literal.
+        config:               The run's config, needed by the macro file for
+                              ``BATCH_SIZE``.
+
+    A ``gdelta/gdelta_aggregate.json`` in ``output_dir``, written by
+    ``run_ablations.py --measure-gdelta``, additionally produces
+    ``tables/gdelta.{csv,tex}`` and the measured decomposition macros.
 
     Returns:
         Dict mapping table name -> polars DataFrame.
@@ -818,5 +1139,20 @@ def generate_summary_tables(
                 label="tab:per_env",
             )
 
+    # Present only once the gradient measurement has been run for this run
+    # directory; its absence is the normal state and not an error.
+    gdelta_agg = load_gdelta_aggregate(output_dir)
+    if gdelta_agg is not None:
+        tables["gdelta"] = write_gdelta_table(gdelta_agg, output_dir)
+
+    if emit_tex_macros:
+        write_tex_macros(
+            results,
+            pretrained_score,
+            tables_dir / "results.tex",
+            config=config,
+            gdelta=gdelta_agg,
+        )
+
     logger.info("All tables saved to %s", tables_dir)
     return tables

experiments/rl_finetuning/configs/ablations_final_craftax_classic_gpu_24gb.yaml

@@ -1,8 +1,8 @@
 # =============================================================================
-# RL Fine-Tuning Ablation Suite — Craftax Classic, 24 GB hardware
+# RL Fine-Tuning Ablation Suite — Craftax Classic, GPU-24GB hardware
 # Matches checkpoint produced by: configs/final_craftax_classic_gpu_24gb.yaml
 #   (d_model=384, n_heads=8, n_layers=6, d_ff=768, plan_horizon=32, seed 42)
-# Target hardware: RTX 3090 Ti (24 GB VRAM)
+# Target hardware: GPU-24GB RTX 3090 Ti (24 GB VRAM)
 # That shape is the ablations_default.yaml architecture, so it is inherited
 # rather than restated here. This file holds only its own deltas; every
 # other key comes from the base, which is layered on automatically.

experiments/rl_finetuning/configs/ablations_final_craftax_classic_gpu_h200.yaml

@@ -1,5 +1,5 @@
 # =============================================================================
-# RL Fine-Tuning Ablation Suite — Craftax Classic, H200 hardware
+# RL Fine-Tuning Ablation Suite — Craftax Classic, GPU-H200 hardware
 # Matches checkpoint produced by: configs/final_craftax_classic_gpu_h200.yaml
 #   (d_model=384, n_heads=8, n_layers=6, d_ff=768, plan_horizon=32, seed 43)
 # Target hardware: H200 8 GB partition

experiments/rl_finetuning/configs/ablations_final_craftax_gpu_24gb.yaml

@@ -1,8 +1,8 @@
 # =============================================================================
-# RL Fine-Tuning Ablation Suite — Full Craftax, 24 GB hardware
+# RL Fine-Tuning Ablation Suite — Full Craftax, GPU-24GB hardware
 # Matches checkpoints produced by: configs/final_craftax_gpu_24gb.yaml
 #   (d_model=384, n_heads=8, n_layers=6, d_ff=768, plan_horizon=32, seed 42)
-# Target hardware: 24 GB reference machine (GPU model unrecorded) (24 GB VRAM)
+# Target hardware: GPU-24GB reference machine (GPU model unrecorded) (24 GB VRAM)
 # That shape is the ablations_default.yaml architecture, so it is inherited
 # rather than restated here. This file holds only its own deltas; every
 # other key comes from the base, which is layered on automatically.

experiments/rl_finetuning/configs/ablations_final_craftax_gpu_h200.yaml

@@ -1,5 +1,5 @@
 # =============================================================================
-# RL Fine-Tuning Ablation Suite — Full Craftax, H200 hardware
+# RL Fine-Tuning Ablation Suite — Full Craftax, GPU-H200 hardware
 # Matches checkpoints produced by: configs/final_craftax_gpu_h200.yaml
 #   (d_model=384, n_heads=8, n_layers=6, d_ff=768, plan_horizon=32, seed 43)
 # Target hardware: H200 8 GB partition

experiments/rl_finetuning/outputs/craftax_classic_ablations/gdelta/gdelta_aggregate.json

@@ -0,0 +1,103 @@
+{
+  "aggregate": true,
+  "inputs": [
+    "experiments/rl_finetuning/outputs/craftax_classic_ablations/gdelta/gdelta_seed0.json",
+    "experiments/rl_finetuning/outputs/craftax_classic_ablations/gdelta/gdelta_seed1.json",
+    "experiments/rl_finetuning/outputs/craftax_classic_ablations/gdelta/gdelta_seed2.json"
+  ],
+  "seeds": [
+    0,
+    1,
+    2
+  ],
+  "n_seeds": 3,
+  "n_draws_per_seed": [
+    8,
+    8,
+    8
+  ],
+  "n_params": 9334289,
+  "random_cos_sd": 0.000327310078675855,
+  "batch": 1024,
+  "eq4_residual_max": 4.80532605706685e-7,
+  "variants": {
+    "baseline_clipped_ratio": {
+      "a1_violated": false,
+      "cv_a_mean": 0.9854568839073181,
+      "cv_a_std_seeds": 0.018878114308338592,
+      "abar_mean": 1.0235402981440227,
+      "abar_std_seeds": 0.010061440877455132,
+      "abar_ratio_to_baseline_mean": 1.0,
+      "abar_ratio_to_baseline_std_seeds": 0.0,
+      "ess_fraction_mean": 0.507413487651479,
+      "ess_fraction_std_seeds": 0.009639155208801934,
+      "ratio_mean": 0.48452565418889765,
+      "ratio_std_seeds": 0.006983698107912942,
+      "cos_mean": 0.023888170331095655,
+      "cos_std_seeds": 0.055179809801403155,
+      "ratio_shuffled_mean": 0.3560985313002038,
+      "ratio_shuffled_std_seeds": 0.01911878846489955,
+      "cos_shuffled_mean": -0.05353597580688074,
+      "cos_shuffled_std_seeds": 0.025540160726829325
+    },
+    "advantage_clip": {
+      "a1_violated": false,
+      "cv_a_mean": 0.2014569491147995,
+      "cv_a_std_seeds": 0.0004557291260598919,
+      "abar_mean": 0.9526590506235758,
+      "abar_std_seeds": 0.009015712817429673,
+      "abar_ratio_to_baseline_mean": 0.9307592739286493,
+      "abar_ratio_to_baseline_std_seeds": 0.003565955741888492,
+      "ess_fraction_mean": 0.9609977946615301,
+      "ess_fraction_std_seeds": 0.00016941930355432837,
+      "ratio_mean": 0.09726429819029546,
+      "ratio_std_seeds": 0.0032068849286760316,
+      "cos_mean": -0.002687329212979724,
+      "cos_std_seeds": 0.06825983346564395,
+      "ratio_shuffled_mean": 0.07327280714136783,
+      "ratio_shuffled_std_seeds": 0.002498156230349314,
+      "cos_shuffled_mean": -0.03329448491179695,
+      "cos_shuffled_std_seeds": 0.052956390835973914
+    },
+    "normalized_adv": {
+      "a1_violated": true,
+      "cv_a_mean": 0.9999999205271403,
+      "cv_a_std_seeds": 7.434005313662571e-8,
+      "abar_mean": -1.0322158535321554e-8,
+      "abar_std_seeds": 3.960053117034971e-8,
+      "abar_ratio_to_baseline_mean": -1.0236111820250936e-8,
+      "abar_ratio_to_baseline_std_seeds": 3.8875732176752924e-8,
+      "ess_fraction_mean": 1.6747491186726309e-15,
+      "ess_fraction_std_seeds": 9.26030088934509e-16,
+      "ratio_mean": 0.4919759005782294,
+      "ratio_std_seeds": 0.016081065492987393,
+      "cos_mean": 0.023888198743710138,
+      "cos_std_seeds": 0.05517983637811973,
+      "ratio_shuffled_mean": 0.357488476116716,
+      "ratio_shuffled_std_seeds": 0.020395976400501593,
+      "cos_shuffled_mean": 0.05980315369864305,
+      "cos_shuffled_std_seeds": 0.04091382822120188
+    },
+    "bc_wins": {
+      "a1_violated": false,
+      "cv_a_mean": 0.7747447292009989,
+      "cv_a_std_seeds": 0.021514876051124374,
+      "abar_mean": 1.0,
+      "abar_std_seeds": 1.7547132305955582e-7,
+      "abar_ratio_to_baseline_mean": 0.9770956797109048,
+      "abar_ratio_to_baseline_std_seeds": 0.009621438085986775,
+      "ess_fraction_mean": 0.624999973036458,
+      "ess_fraction_std_seeds": 0.013053179236684295,
+      "ratio_mean": 0.4170324476063991,
+      "ratio_std_seeds": 0.024545794473638916,
+      "cos_mean": 0.1420223995228298,
+      "cos_std_seeds": 0.09300934862546119,
+      "ratio_shuffled_mean": 0.2839805669591658,
+      "ratio_shuffled_std_seeds": 0.007360032818748407,
+      "cos_shuffled_mean": 0.014593212322021524,
+      "cos_shuffled_std_seeds": 0.056382858532646296
+    }
+  },
+  "bc_self_cos_mean": 0.892872134844462,
+  "bc_self_cos_std_seeds": 0.009867953837605352
+}

experiments/rl_finetuning/outputs/craftax_classic_ablations/gdelta/gdelta_seed0.json

@@ -0,0 +1,74 @@
+{
+  "aggregate": false,
+  "checkpoint_step": 40370176,
+  "n_params": 9334289,
+  "random_cos_sd": 0.000327310078675855,
+  "batch": 1024,
+  "seed": 0,
+  "n_draws": 8,
+  "bc_self_cos_mean": 0.8886517286300659,
+  "bc_self_cos_std": 0.013864377043282333,
+  "eq4_residual_max": 4.80532605706685e-7,
+  "variants": {
+    "baseline_clipped_ratio": {
+      "cv_a": 0.9587592482566833,
+      "abar": 1.0352622270584106,
+      "ess_fraction": 0.5210452873754741,
+      "a1_violated": false,
+      "abar_ratio_to_baseline": 1.0,
+      "ratio_mean": 0.4929142524131889,
+      "ratio_std_draws": 0.062149263585443426,
+      "cos_mean": -0.05232078512199223,
+      "cos_std_draws": 0.12823055031892652,
+      "ratio_shuffled_mean": 0.3702785550029438,
+      "ratio_shuffled_std": 0.04985673247875491,
+      "cos_shuffled_mean": -0.07264717968791956,
+      "cos_shuffled_std": 0.1375778261666584
+    },
+    "advantage_clip": {
+      "cv_a": 0.2019623965024948,
+      "abar": 0.9651404619216919,
+      "ess_fraction": 0.9608097961720935,
+      "a1_violated": false,
+      "abar_ratio_to_baseline": 0.9322666631661406,
+      "ratio_mean": 0.10178494489910378,
+      "ratio_std_draws": 0.010863613185205834,
+      "cos_mean": -0.08856321198982187,
+      "cos_std_draws": 0.09348871621494285,
+      "ratio_shuffled_mean": 0.07580454134313921,
+      "ratio_shuffled_std": 0.006146976066602656,
+      "cos_shuffled_mean": -0.08668288914486766,
+      "cos_shuffled_std": 0.10151970823488679
+    },
+    "normalized_adv": {
+      "cv_a": 0.9999999403953552,
+      "abar": -2.1653249859809875e-8,
+      "ess_fraction": 4.688632853842156e-16,
+      "a1_violated": true,
+      "abar_ratio_to_baseline": -2.0915715162655282e-8,
+      "ratio_mean": 0.5141167874551692,
+      "ratio_std_draws": 0.06482259643725917,
+      "cos_mean": -0.05232076484389836,
+      "cos_std_draws": 0.12823054295881925,
+      "ratio_shuffled_mean": 0.38073488366371927,
+      "ratio_shuffled_std": 0.024812852770472734,
+      "cos_shuffled_mean": 0.06606726971222088,
+      "cos_shuffled_std": 0.16606576384474356
+    },
+    "bc_wins": {
+      "cv_a": 0.7458090782165527,
+      "abar": 0.9999998211860657,
+      "ess_fraction": 0.6425777397495268,
+      "a1_violated": false,
+      "abar_ratio_to_baseline": 0.9659386723955539,
+      "ratio_mean": 0.3823743698913481,
+      "ratio_std_draws": 0.04196826406011548,
+      "cos_mean": 0.02126681269146502,
+      "cos_std_draws": 0.12931455515444024,
+      "ratio_shuffled_mean": 0.2905509104827731,
+      "ratio_shuffled_std": 0.0079822831305291,
+      "cos_shuffled_mean": 0.0010593719780445099,
+      "cos_shuffled_std": 0.22469821243454166
+    }
+  }
+}

experiments/rl_finetuning/outputs/craftax_classic_ablations/gdelta/gdelta_seed1.json

@@ -0,0 +1,74 @@
+{
+  "aggregate": false,
+  "checkpoint_step": 40370176,
+  "n_params": 9334289,
+  "random_cos_sd": 0.000327310078675855,
+  "batch": 1024,
+  "seed": 1,
+  "n_draws": 8,
+  "bc_self_cos_mean": 0.8834625333547592,
+  "bc_self_cos_std": 0.028890878138939563,
+  "eq4_residual_max": 4.458028683984594e-7,
+  "variants": {
+    "baseline_clipped_ratio": {
+      "cv_a": 0.9988502860069275,
+      "abar": 1.0246647596359253,
+      "ess_fraction": 0.5005752960258663,
+      "a1_violated": false,
+      "abar_ratio_to_baseline": 1.0,
+      "ratio_mean": 0.48484595625152527,
+      "ratio_std_draws": 0.03739208589573637,
+      "cos_mean": 0.047454304061830044,
+      "cos_std_draws": 0.06307783640351931,
+      "ratio_shuffled_mean": 0.36894560309691693,
+      "ratio_shuffled_std": 0.022733560664731905,
+      "cos_shuffled_mean": -0.017437556954973843,
+      "cos_shuffled_std": 0.13464438960797637
+    },
+    "advantage_clip": {
+      "cv_a": 0.20085792243480682,
+      "abar": 0.94867342710495,
+      "ess_fraction": 0.9612204281808505,
+      "a1_violated": false,
+      "abar_ratio_to_baseline": 0.9258378588544649,
+      "ratio_mean": 0.0953185879622219,
+      "ratio_std_draws": 0.011684367167521479,
+      "cos_mean": 0.002065446664346382,
+      "cos_std_draws": 0.07558667393060121,
+      "ratio_shuffled_mean": 0.0741409312015546,
+      "ratio_shuffled_std": 0.0062286893856330055,
+      "cos_shuffled_mean": 0.03888408502098173,
+      "cos_shuffled_std": 0.14181727110818865
+    },
+    "normalized_adv": {
+      "cv_a": 0.9999998211860657,
+      "abar": 4.284083843231201e-8,
+      "ess_fraction": 1.835337984556805e-15,
+      "a1_violated": true,
+      "abar_ratio_to_baseline": 4.18096143440454e-8,
+      "ratio_mean": 0.4854039954925514,
+      "ratio_std_draws": 0.037435128863097536,
+      "cos_mean": 0.04745421721599996,
+      "cos_std_draws": 0.06307784333867371,
+      "ratio_shuffled_mean": 0.3606536861227143,
+      "ratio_shuffled_std": 0.0351022186659039,
+      "cos_shuffled_mean": 0.10648557799868286,
+      "cos_shuffled_std": 0.06018093926588145
+    },
+    "bc_wins": {
+      "cv_a": 0.7973600029945374,
+      "abar": 0.9999999403953552,
+      "ess_fraction": 0.6113279738392736,
+      "a1_violated": false,
+      "abar_ratio_to_baseline": 0.9759288889281859,
+      "ratio_mean": 0.4326711384712623,
+      "ratio_std_draws": 0.03118598407923515,
+      "cos_mean": 0.1572374363313429,
+      "cos_std_draws": 0.14007506866888625,
+      "ratio_shuffled_mean": 0.287686680066821,
+      "ratio_shuffled_std": 0.014448778378290157,
+      "cos_shuffled_mean": -0.04669254180043936,
+      "cos_shuffled_std": 0.2434607700636855
+    }
+  }
+}

experiments/rl_finetuning/outputs/craftax_classic_ablations/gdelta/gdelta_seed2.json

@@ -0,0 +1,74 @@
+{
+  "aggregate": false,
+  "checkpoint_step": 40370176,
+  "n_params": 9334289,
+  "random_cos_sd": 0.000327310078675855,
+  "batch": 1024,
+  "seed": 2,
+  "n_draws": 8,
+  "bc_self_cos_mean": 0.9065021425485611,
+  "bc_self_cos_std": 0.0077027198757505366,
+  "eq4_residual_max": 4.4017934897055966e-7,
+  "variants": {
+    "baseline_clipped_ratio": {
+      "cv_a": 0.9987611174583435,
+      "abar": 1.010693907737732,
+      "ess_fraction": 0.5006198795530965,
+      "a1_violated": false,
+      "abar_ratio_to_baseline": 1.0,
+      "ratio_mean": 0.47581675390197886,
+      "ratio_std_draws": 0.050013706533038056,
+      "cos_mean": 0.07653099205344915,
+      "cos_std_draws": 0.1371696455986107,
+      "ratio_shuffled_mean": 0.3290714358007507,
+      "ratio_shuffled_std": 0.017993567005007666,
+      "cos_shuffled_mean": -0.07052319077774882,
+      "cos_shuffled_std": 0.14557688155890267
+    },
+    "advantage_clip": {
+      "cv_a": 0.20155052840709686,
+      "abar": 0.9441632628440857,
+      "ess_fraction": 0.9609631596316461,
+      "a1_violated": false,
+      "abar_ratio_to_baseline": 0.9341732997653425,
+      "ratio_mean": 0.09468936170956067,
+      "ratio_std_draws": 0.0089181304684684,
+      "cos_mean": 0.07843577768653631,
+      "cos_std_draws": 0.15751443556107037,
+      "ratio_shuffled_mean": 0.06987294887940967,
+      "ratio_shuffled_std": 0.005245661765416722,
+      "cos_shuffled_mean": -0.05208465061150491,
+      "cos_shuffled_std": 0.1682371637969602
+    },
+    "normalized_adv": {
+      "cv_a": 1.0,
+      "abar": -5.21540641784668e-8,
+      "ess_fraction": 2.720046086076872e-15,
+      "a1_violated": true,
+      "abar_ratio_to_baseline": -5.1602234642142926e-8,
+      "ratio_mean": 0.4764069187869675,
+      "ratio_std_draws": 0.05007575024715252,
+      "cos_mean": 0.07653114385902882,
+      "cos_std_draws": 0.13716963854788167,
+      "ratio_shuffled_mean": 0.3310768585637146,
+      "ratio_shuffled_std": 0.02154253356764448,
+      "cos_shuffled_mean": 0.006856613385025412,
+      "cos_shuffled_std": 0.19543885350694304
+    },
+    "bc_wins": {
+      "cv_a": 0.7810651063919067,
+      "abar": 1.000000238418579,
+      "ess_fraction": 0.6210942055205736,
+      "a1_violated": false,
+      "abar_ratio_to_baseline": 0.989419477808975,
+      "ratio_mean": 0.4360518344565869,
+      "ratio_std_draws": 0.025181532407221242,
+      "cos_mean": 0.24756294954568148,
+      "cos_std_draws": 0.08224887485750225,
+      "ratio_shuffled_mean": 0.2737041103279033,
+      "ratio_shuffled_std": 0.02088922719350182,
+      "cos_shuffled_mean": 0.08941280678845942,
+      "cos_shuffled_std": 0.19656693288770055
+    }
+  }
+}

experiments/rl_finetuning/outputs/craftax_classic_ablations/results.json

@@ -26,11 +26,11 @@
   },
   "merge_provenance": {
     "inputs": [
-      "<path omitted for anonymous review>",
-      "<path omitted for anonymous review>",
-      "<path omitted for anonymous review>"
+      "gpu_24gb_0/results.json",
+      "gpu_24gb_1/results.json",
+      "gpu_24gb_2/results.json"
     ],
-    "config_from": "<path omitted for anonymous review>"
+    "config_from": "gpu_24gb_0/results.json"
   },
   "config": {
     "ENV_NAME": "Craftax-Classic-Symbolic-v1",
@@ -87,14 +87,9 @@
     "SAVE_POLICY": true,
     "CHECKPOINT_DIR": "checkpoints",
     "RESUME_CHECKPOINT_PATH": null,
-    "RESUME_WANDB_RUN_ID": null,
     "RESUME_STEP": null,
-    "JAX_COMPILATION_CACHE_DIR": "<path omitted for anonymous review>",
+    "JAX_COMPILATION_CACHE_DIR": null,
     "SEED": null,
-    "USE_WANDB": true,
-    "WANDB_PROJECT": "remdm-planner-craftax-ablations",
-    "WANDB_ENTITY": "null",
-    "WANDB_DOWNLOAD_DIR": null,
     "MAX_ITER": 500,
     "BATCH_SIZE": 1024,
     "DIFFUSION_STEPS_COLLECT": 5,
@@ -132,7 +127,7 @@
     "REWARD_MODEL_LR": 0.001,
     "REWARD_MODEL_TRAIN_STEPS": 50,
     "NUM_SEEDS": 3,
-    "CHECKPOINT_PATH": "<path omitted for anonymous review>",
+    "CHECKPOINT_PATH": "ablation_src/policy-best-v2",
     "NUM_ACTIONS": 17
   },
   "ablations": {