Upload MiniHack ReMDM planner checkpoints

README.md

@@ -12,39 +12,32 @@ tags:
 - pytorch
 ---
 
-# ReMDM Planner — MiniHack checkpoints
+# ReMDM Planner: MiniHack checkpoints
 
 Trained weights accompanying *The Double Intractability of Reinforcement Learning for Discrete Diffusion Planners*: a remasking discrete diffusion model
 (ReMDM) used as an action-sequence planner in
-[MiniHack](https://github.com/facebookresearch/minihack).
+[MiniHack](https://github.com/facebookresearch/minihack), together with the BFS
+oracle rollouts that supervise it.
 
 Code, configs and evaluation harness: https://github.com/mathisweil/minihack-ReMDM-planner
 
 ## Contents
 
-| Path | Role | Architecture | Params | Trained to | Full state |
-|---|---|---|---|---|---|
-| `checkpoints/offline/Minihack-OfflineDiffusion-BC-82M` | Diffusion planner (offline BC) | 4L, d_model 256, 4 heads, horizon 64 | 5.2M | 40,000 gradient steps, 81,920,000 sample-equivalents | 80 MB |
-| `checkpoints/online/Minihack-Online-Diffusion-DAgger-100M` | Diffusion planner (online DAgger) | 4L, d_model 256, 4 heads, horizon 64 | 5.2M | DAgger iteration 563, 5,654,965 env steps | 80 MB |
+| Path | Role | Environment | Architecture | Selected at | Training | Size |
+|---|---|---|---|---|---|---|
+| `checkpoints` | minihack-ReMDM-planner | `MiniHack` | 4L, d_model 256, 4 heads, horizon 64, 5M params | iteration 5 | 2,835 env steps | 260 MB |
+| `checkpoints/offline/Minihack-OfflineDiffusion-BC-82M` | Diffusion planner (offline BC) | `MiniHack` | 4L, d_model 256, 4 heads, horizon 64, 5M params | gradient step 40,000 | 81,920,000 sample-equivalents | 100 MB |
+| `checkpoints/online/Minihack-Online-Diffusion-DAgger-100M` | Diffusion planner (online DAgger) | `MiniHack` | 4L, d_model 256, 4 heads, horizon 64, 5M params | iteration 563 | 5,654,965 env steps | 100 MB |
 
-Each directory holds three things: the original `.pth` training state (weights,
+Each checkpoint ships the `.pth` training state it was published from (weights,
 EMA shadow, optimiser, scheduler, and for the DAgger run the curriculum and RNG
-state, so training can be resumed exactly), a `model.safetensors` export of the
-EMA weights for inference, and the YAML config snapshot the run was trained
-under. Paths mirror the source repository, so a snapshot can be dropped
-straight into a working copy.
+state, so training resumes exactly), a `model.safetensors` export of the EMA
+weights for inference, the YAML config snapshot it was trained under, and a
+`selection.json` recording how it was chosen.
 
-Both files are best-checkpoint selections rather than final-step dumps: each
-trainer evaluates every periodic checkpoint on 50 episodes per environment
-using EMA weights, and the highest-scoring one is published (the metric behind that selection is not recorded in this release).
-Each directory's `selection.json` records the selected step, the candidate
-cadence and the eval protocol. Directory suffixes
-are the sample-equivalents the published model consumed (gradient steps x batch
-size, rounded); file names carry each trainer's own counter, DAgger iterations
-online and gradient steps offline. The offline run was given the
-DAgger-matched budget of 60,000 gradient steps and its best checkpoint fell at
-40,000, so the two published models sit at different points on a matched
-budget.
+Weights are PyTorch training states with a `safetensors` export of the EMA
+weights alongside, and the paths above mirror the source repository so a
+snapshot can be dropped straight into a working copy.
 
 `results/` holds the evaluation and ablation tables reported in the paper, as
 produced by `experiments/rl_finetuning`. Figures and raw logs stay in the code
@@ -55,14 +48,14 @@ repository.
 ```python
 from huggingface_hub import snapshot_download
 
-# everything
+# everything (~460 MB)
 snapshot_download(repo_id="mathisweil/remdm-minihack-checkpoints", local_dir=".")
 
-# inference weights only
+# a single model
 snapshot_download(
     repo_id="mathisweil/remdm-minihack-checkpoints",
     local_dir=".",
-    allow_patterns=["**/model.safetensors", "**/config*.yaml"],
+    allow_patterns="checkpoints/online/Minihack-*/**",
 )
 ```
 
@@ -76,7 +69,8 @@ uv run python main.py --mode inference \
     --config $DIR/config_iter600.yaml --checkpoint $DIR/iter600.pth
 ```
 
-Programmatic loading, using the safetensors export:
+Programmatic loading uses `src.models.denoiser.make_model` with the checkpoint's
+own config, then the safetensors export:
 
 ```python
 from safetensors.torch import load_file
@@ -89,18 +83,29 @@ model.load_state_dict(load_file(f"{DIR}/model.safetensors"))
 model.eval()
 ```
 
-Architecture arguments must come from the checkpoint's own config snapshot
+Architecture arguments should be read from the checkpoint's own config snapshot
 rather than from `configs/defaults.yaml`, which tracks the current code.
 
 ## Training
 
 The planners are bidirectional transformers that denoise a masked action plan
 conditioned on a cropped MiniHack glyph observation, trained either by offline
-behaviour cloning on oracle rollouts or by online DAgger against the oracle
-with a dynamic environment curriculum. In-distribution and out-of-distribution
-environment sets, remasking strategy, sampling settings and every
-hyperparameter are recorded in the per-checkpoint config snapshots, which are
-the authoritative record.
+behaviour cloning on oracle rollouts or by online DAgger against the BFS oracle
+under a dynamic environment curriculum. Model size and horizon differ per run
+(see the table). Exact hyperparameters for every run, including the
+in-distribution and out-of-distribution environment sets, the remasking
+strategy, schedule and sampling settings, are in the per-checkpoint config
+snapshots listed above, which are the authoritative record.
+
+Both models are best-checkpoint selections rather than final-step dumps: each
+trainer evaluates every periodic checkpoint on 50 episodes per environment
+using EMA weights, and the highest-scoring one is published (the metric behind that selection is not recorded in this release).
+Directory names encode the sample-equivalents the published model consumed
+(gradient steps x batch size, rounded); file names carry each trainer's own
+counter, DAgger iterations online and gradient steps offline. The offline run
+was given the DAgger-matched budget of 60,000 gradient steps and its best
+checkpoint fell at 40,000, so the two published models sit at different points
+on a matched budget.
 
 ## Limitations
 

checkpoints/config_iter10.yaml

@@ -0,0 +1,92 @@
+action_dim: 12
+aux_loss_weight: 0.5
+baselines_bc_batch_size: 256
+baselines_bc_epochs: 50
+baselines_bc_lr: 0.0003
+baselines_bc_oracle_episodes_per_env: 5000
+baselines_dqn_buffer_size: 100000
+baselines_dt_batch_size: 256
+baselines_dt_context_len: 64
+baselines_dt_embed_dim: 256
+baselines_dt_epochs: 50
+baselines_dt_eval_max_steps: 200
+baselines_dt_lr: 0.0003
+baselines_dt_max_ep_len: 200
+baselines_dt_n_heads: 4
+baselines_dt_n_layers: 4
+baselines_dt_oracle_episodes_per_env: 5000
+baselines_eval_episodes_per_env: null
+baselines_eval_freq_env_steps: 10000
+baselines_n_envs_per_id: 2
+baselines_output_dir: outputs/baselines
+baselines_wandb_project: minihack-ReMDM-baselines
+buffer_capacity: 50
+checkpoint_dir: checkpoints
+checkpoint_eval_episodes: 2
+checkpoint_every_timesteps: 2500
+collect_episodes_per_env: 5000
+collect_num_workers: 8
+collect_output: data/dataset.pt
+crop_size: 9
+curriculum_preseed: true
+curriculum_queue_size: 100
+dagger_batch_size: 256
+dagger_grad_clip: 1.0
+dagger_lr: 3.0e-05
+device: cpu
+diffusion_steps_collect: 5
+diffusion_steps_eval: 10
+dropout: 0.0
+efficiency_multiplier: 1.5
+ema_decay: 0.999
+episodes_per_iteration: 2
+eta: 0.15
+eval_episodes_per_env: 2
+global_gate_init: -3.0
+grad_steps_per_iteration: 5
+id_envs:
+- MiniHack-Room-Random-5x5-v0
+- MiniHack-Room-Random-15x15-v0
+- MiniHack-Corridor-R2-v0
+- MiniHack-MazeWalk-9x9-v0
+id_eval_every_timesteps: 2500
+label_smoothing: 0.0
+loss_weight_clip: 1000.0
+map_h: 21
+map_w: 79
+mask_token: 12
+n_embd: 256
+n_global_tokens: 8
+n_head: 4
+n_layer: 4
+noise_schedule: linear
+num_collection_workers: 0
+num_diffusion_steps: 100
+offline_batch_size: 256
+offline_buffer_capacity: null
+offline_checkpoint_every_grad_steps: null
+offline_eval_every_grad_steps: null
+offline_grad_clip: 1.0
+offline_log_every: 10
+offline_lr: 0.0003
+offline_total_grad_steps: null
+ood_envs:
+- MiniHack-Room-Dark-15x15-v0
+- MiniHack-Corridor-R5-v0
+- MiniHack-MazeWalk-45x19-v0
+ood_eval_every_timesteps: 2500
+pad_token: 13
+physics_aware_sampling: false
+remask_strategy: conf
+replan_every: 16
+save_policy: true
+seed: null
+seq_len: 64
+temperature: 0.5
+top_p: 0.9
+torch_compile: false
+total_timesteps: 5000
+use_amp: false
+use_global_stream: true
+use_wandb: false
+weight_decay: 0.0001

checkpoints/config_iter11.yaml

@@ -0,0 +1,92 @@
+action_dim: 12
+aux_loss_weight: 0.5
+baselines_bc_batch_size: 256
+baselines_bc_epochs: 50
+baselines_bc_lr: 0.0003
+baselines_bc_oracle_episodes_per_env: 5000
+baselines_dqn_buffer_size: 100000
+baselines_dt_batch_size: 256
+baselines_dt_context_len: 64
+baselines_dt_embed_dim: 256
+baselines_dt_epochs: 50
+baselines_dt_eval_max_steps: 200
+baselines_dt_lr: 0.0003
+baselines_dt_max_ep_len: 200
+baselines_dt_n_heads: 4
+baselines_dt_n_layers: 4
+baselines_dt_oracle_episodes_per_env: 5000
+baselines_eval_episodes_per_env: null
+baselines_eval_freq_env_steps: 10000
+baselines_n_envs_per_id: 2
+baselines_output_dir: outputs/baselines
+baselines_wandb_project: minihack-ReMDM-baselines
+buffer_capacity: 50
+checkpoint_dir: checkpoints
+checkpoint_eval_episodes: 2
+checkpoint_every_timesteps: 2500
+collect_episodes_per_env: 5000
+collect_num_workers: 8
+collect_output: data/dataset.pt
+crop_size: 9
+curriculum_preseed: true
+curriculum_queue_size: 100
+dagger_batch_size: 256
+dagger_grad_clip: 1.0
+dagger_lr: 3.0e-05
+device: cpu
+diffusion_steps_collect: 5
+diffusion_steps_eval: 10
+dropout: 0.0
+efficiency_multiplier: 1.5
+ema_decay: 0.999
+episodes_per_iteration: 2
+eta: 0.15
+eval_episodes_per_env: 2
+global_gate_init: -3.0
+grad_steps_per_iteration: 5
+id_envs:
+- MiniHack-Room-Random-5x5-v0
+- MiniHack-Room-Random-15x15-v0
+- MiniHack-Corridor-R2-v0
+- MiniHack-MazeWalk-9x9-v0
+id_eval_every_timesteps: 2500
+label_smoothing: 0.0
+loss_weight_clip: 1000.0
+map_h: 21
+map_w: 79
+mask_token: 12
+n_embd: 256
+n_global_tokens: 8
+n_head: 4
+n_layer: 4
+noise_schedule: linear
+num_collection_workers: 0
+num_diffusion_steps: 100
+offline_batch_size: 256
+offline_buffer_capacity: null
+offline_checkpoint_every_grad_steps: null
+offline_eval_every_grad_steps: null
+offline_grad_clip: 1.0
+offline_log_every: 10
+offline_lr: 0.0003
+offline_total_grad_steps: null
+ood_envs:
+- MiniHack-Room-Dark-15x15-v0
+- MiniHack-Corridor-R5-v0
+- MiniHack-MazeWalk-45x19-v0
+ood_eval_every_timesteps: 2500
+pad_token: 13
+physics_aware_sampling: false
+remask_strategy: conf
+replan_every: 16
+save_policy: true
+seed: null
+seq_len: 64
+temperature: 0.5
+top_p: 0.9
+torch_compile: false
+total_timesteps: 5000
+use_amp: false
+use_global_stream: true
+use_wandb: false
+weight_decay: 0.0001

checkpoints/config_iter5.yaml

@@ -0,0 +1,92 @@
+action_dim: 12
+aux_loss_weight: 0.5
+baselines_bc_batch_size: 256
+baselines_bc_epochs: 50
+baselines_bc_lr: 0.0003
+baselines_bc_oracle_episodes_per_env: 5000
+baselines_dqn_buffer_size: 100000
+baselines_dt_batch_size: 256
+baselines_dt_context_len: 64
+baselines_dt_embed_dim: 256
+baselines_dt_epochs: 50
+baselines_dt_eval_max_steps: 200
+baselines_dt_lr: 0.0003
+baselines_dt_max_ep_len: 200
+baselines_dt_n_heads: 4
+baselines_dt_n_layers: 4
+baselines_dt_oracle_episodes_per_env: 5000
+baselines_eval_episodes_per_env: null
+baselines_eval_freq_env_steps: 10000
+baselines_n_envs_per_id: 2
+baselines_output_dir: outputs/baselines
+baselines_wandb_project: minihack-ReMDM-baselines
+buffer_capacity: 50
+checkpoint_dir: checkpoints
+checkpoint_eval_episodes: 2
+checkpoint_every_timesteps: 2500
+collect_episodes_per_env: 5000
+collect_num_workers: 8
+collect_output: data/dataset.pt
+crop_size: 9
+curriculum_preseed: true
+curriculum_queue_size: 100
+dagger_batch_size: 256
+dagger_grad_clip: 1.0
+dagger_lr: 3.0e-05
+device: cpu
+diffusion_steps_collect: 5
+diffusion_steps_eval: 10
+dropout: 0.0
+efficiency_multiplier: 1.5
+ema_decay: 0.999
+episodes_per_iteration: 2
+eta: 0.15
+eval_episodes_per_env: 2
+global_gate_init: -3.0
+grad_steps_per_iteration: 5
+id_envs:
+- MiniHack-Room-Random-5x5-v0
+- MiniHack-Room-Random-15x15-v0
+- MiniHack-Corridor-R2-v0
+- MiniHack-MazeWalk-9x9-v0
+id_eval_every_timesteps: 2500
+label_smoothing: 0.0
+loss_weight_clip: 1000.0
+map_h: 21
+map_w: 79
+mask_token: 12
+n_embd: 256
+n_global_tokens: 8
+n_head: 4
+n_layer: 4
+noise_schedule: linear
+num_collection_workers: 0
+num_diffusion_steps: 100
+offline_batch_size: 256
+offline_buffer_capacity: null
+offline_checkpoint_every_grad_steps: null
+offline_eval_every_grad_steps: null
+offline_grad_clip: 1.0
+offline_log_every: 10
+offline_lr: 0.0003
+offline_total_grad_steps: null
+ood_envs:
+- MiniHack-Room-Dark-15x15-v0
+- MiniHack-Corridor-R5-v0
+- MiniHack-MazeWalk-45x19-v0
+ood_eval_every_timesteps: 2500
+pad_token: 13
+physics_aware_sampling: false
+remask_strategy: conf
+replan_every: 16
+save_policy: true
+seed: null
+seq_len: 64
+temperature: 0.5
+top_p: 0.9
+torch_compile: false
+total_timesteps: 5000
+use_amp: false
+use_global_stream: true
+use_wandb: false
+weight_decay: 0.0001

checkpoints/iter10.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:684c6efd17adb8158e436b67f5cba5b66f58d81305658d92f13001cac943e61e
+size 83984510

checkpoints/iter11.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de57e58957564c62bc48e3aa9c31844b30b0098de069efab3d307a937002053e
+size 83984510

checkpoints/iter5.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c933ca32ad857e3731bb9dddeddb1c013db732cf4a824e8452924a2c8b97785a
+size 83972495

checkpoints/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e4a91c4aed49e5ab182816a8bcaa714387f89b26f170d67bd1902b9eb56425fa
+size 20974876

checkpoints/online/Minihack-Online-Diffusion-DAgger-100M/model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1a25616b8ebd7673982fea78ba33030a552cebd0356cbe9a6ff55c42380a3709
+oid sha256:f023044d7a9d3b56396a0c71e2a4389a7fac941044e214a9fc05cc52af152b91
 size 20974876

checkpoints/selection.json

@@ -0,0 +1,27 @@
+{
+  "policy": "best-of-N over periodic checkpoints",
+  "selected": {
+    "dagger_iteration": 5
+  },
+  "selection_metric": null,
+  "candidates": {
+    "unit": "dagger_iterations",
+    "every": null,
+    "configured_max": null
+  },
+  "eval_protocol": {
+    "episodes_per_env": 2,
+    "weights": "ema",
+    "id_envs": [
+      "MiniHack-Room-Random-5x5-v0",
+      "MiniHack-Room-Random-15x15-v0",
+      "MiniHack-Corridor-R2-v0",
+      "MiniHack-MazeWalk-9x9-v0"
+    ],
+    "ood_envs": [
+      "MiniHack-Room-Dark-15x15-v0",
+      "MiniHack-Corridor-R5-v0",
+      "MiniHack-MazeWalk-45x19-v0"
+    ]
+  }
+}