Skip to content

Commit cbab5dc

Browse files
Idate96Idate96
committed
Add larger ResNet distillation training setup
1 parent a7ed731 commit cbab5dc

7 files changed

Lines changed: 919 additions & 37 deletions

File tree

‎docs/EXPERIMENTS_LOG.md‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,65 @@
11
# Experiment Log
22

3+
## 2026-05-18 Larger ResNet Distillation Setup
4+
5+
Goal: launch larger delayed-ResNet policies and test whether teacher imitation
6+
from the trained ResMap64 policy stabilizes larger PPO.
7+
8+
Changes:
9+
10+
- Added explicit delayed-ResNet capacity scaling to `model_size=medium` and
11+
`model_size=large`; the previous presets widened heads/trunks but left the
12+
delayed ResNet itself fixed.
13+
- Added optional teacher-policy imitation warm start to `train_mixed.py`.
14+
The teacher rolls out actions, the student matches teacher logits with KL and
15+
teacher values with an auxiliary MSE, then the Adam optimizer is reset before
16+
PPO starts.
17+
- Oracle review required tightening the experiment before launch:
18+
- PPO now restores the pre-imitation RNG/env/previous-action state so
19+
distill-vs-scratch tests parameter warm start only.
20+
- The final imitation metrics are logged with the first PPO row, and
21+
`_POST_DISTILL.pkl` is saved.
22+
- Teacher/student `num_prev_actions` and parameter-shape compatibility are
23+
checked before distillation.
24+
- Unknown CLI args now fail, and PPO minibatch divisibility is validated.
25+
- Slurm smoke now covers tiny eval and checkpoint paths before online runs.
26+
- Added `scripts/euler/terra_train_larger_resnet_4gpu.sbatch`, parameterized by
27+
`RUN_KIND=medium_distill`, `medium_scratch`, `large_distill`, or
28+
`large_scratch`.
29+
- Added `docs/LARGER_RESNET_DISTILLATION_PLAN.md` with the experiment matrix,
30+
local memory findings, and decision criteria.
31+
32+
Local validation:
33+
34+
- `python -m py_compile train_mixed.py utils/models.py utils/utils_ppo.py`
35+
passed.
36+
- `bash -n scripts/euler/terra_train_larger_resnet_4gpu.sbatch` passed.
37+
- `git diff --check` passed.
38+
- Tiny CPU imitation+PPO smoke passed before and after the Oracle fixes; the
39+
post-fix run saved both `_POST_DISTILL.pkl` and `_FINAL.pkl`.
40+
- RTX 4090 medium smoke:
41+
- `model_size=medium`, `map_feature_dim=192`, `1024` envs/GPU,
42+
`num_steps=32`.
43+
- Failed with `num_minibatches=32` due to ResNet activation OOM.
44+
- Passed with `num_minibatches=64`; one imitation update and one PPO update
45+
completed.
46+
- RTX 4090 large smoke:
47+
- `model_size=large`, `map_feature_dim=256`, `1024` envs/GPU,
48+
`num_steps=32`, `num_minibatches=128`.
49+
- Failed without autotune mitigation with `CUDNN_STATUS_INTERNAL_ERROR`.
50+
- Passed with `XLA_FLAGS=--xla_gpu_autotune_level=0`; one imitation update
51+
and one PPO update completed.
52+
53+
Planned Euler jobs:
54+
55+
- `medium_distill`: medium ResNet, `64` minibatches, `200` imitation updates.
56+
- `medium_scratch`: same medium ResNet and minibatches, no teacher; this is the
57+
warm-start control.
58+
- `large_distill`: large ResNet, `128` minibatches, `100` imitation updates,
59+
cuDNN autotune disabled.
60+
- `large_scratch`: same large ResNet/autotune setup, no teacher; this is the
61+
large-model warm-start control.
62+
363
## 2026-05-18 Default-Unmasked PPO Cleanup
464

565
Goal: make the ResMap path the default unmasked-actor path and remove

‎docs/EXPERIMENTS_RUNNING.md‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,33 @@
11
# Running Experiments
22

3+
## 2026-05-18 Larger ResNet Distillation Jobs
4+
5+
Pending launch from paired `codex/mask-speedup-wip` worktrees:
6+
7+
- Script: `scripts/euler/terra_train_larger_resnet_4gpu.sbatch`
8+
- Teacher checkpoint:
9+
`terra-solo-resmap64-r1r2-terminalfix-mb32-unmasked-4gpu-50B-20260516-euler-4gpu-2026-05-16-12-18-34.pkl`
10+
- Shared semantics: unmasked PPO actor, ResMap derived channels, separate
11+
actor/critic trunks, critic-only edge/progress affordances, timeout bootstrap
12+
and initial timeout-phase staggering from the current branch.
13+
- Local gates on `2026-05-18`:
14+
- CPU imitation+PPO smoke passed before and after the Oracle fixes; the
15+
post-fix run saved `_POST_DISTILL.pkl` and `_FINAL.pkl`.
16+
- Medium RTX 4090 smoke passed with `1024` envs/GPU,
17+
`num_minibatches=64`.
18+
- Large RTX 4090 smoke passed with `1024` envs/GPU,
19+
`num_minibatches=128`, and `XLA_FLAGS=--xla_gpu_autotune_level=0`.
20+
- Full CPU validation sweep passed after the Oracle fixes.
21+
- Planned submissions:
22+
- `RUN_KIND=medium_distill`: `model_size=medium`, `map_feature_dim=192`,
23+
`num_minibatches=64`, `imitation_updates=200`.
24+
- `RUN_KIND=medium_scratch`: same medium architecture, no teacher; control
25+
for the warm-start hypothesis.
26+
- `RUN_KIND=large_distill`: `model_size=large`, `map_feature_dim=256`,
27+
`num_minibatches=128`, `imitation_updates=100`, autotune disabled.
28+
- `RUN_KIND=large_scratch`: same large architecture and autotune mitigation,
29+
no teacher; control for the large warm-start hypothesis.
30+
331
## 2026-05-15 ResMap64 Four-GPU Architecture Run
432

533
Active Slurm jobs:

‎docs/LARGER_RESNET_DISTILLATION_PLAN.md‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
# Larger ResNet Distillation Plan
2+
3+
Date: 2026-05-18
4+
5+
## Objective
6+
7+
Train larger delayed-downsample ResNet policies for the solo-excavator
8+
ringmaps task without changing the primitive action interface. Test the
9+
hypothesis that larger PPO policies are harder to optimize from scratch and
10+
benefit from an imitation warm start from the smaller trained ResMap policy.
11+
12+
## Research Takeaways
13+
14+
- Larger RL networks can help, but the optimization is less forgiving than in
15+
supervised learning; capacity should be paired with stable initialization,
16+
conservative learning-rate/minibatch choices, and stronger auxiliary signal.
17+
- Policy distillation is the lowest-friction warm-start mechanism here because
18+
it does not require matching parameter shapes between small and large
19+
networks. The larger student matches the trained teacher's action
20+
distribution on live environment observations, then PPO takes over.
21+
- A scratch larger-network control is needed. Otherwise a better or worse
22+
larger run cannot be attributed to the imitation stage.
23+
24+
Reference threads:
25+
26+
- Nauman and Raffel, 2024, "The Case for Bigger Models in Reinforcement
27+
Learning": https://arxiv.org/abs/2407.15134
28+
- Rusu et al., 2015, "Policy Distillation": https://arxiv.org/abs/1511.06295
29+
- Oh et al., 2018, "Self-Imitation Learning": https://arxiv.org/abs/1806.05635
30+
31+
## Implemented Setup
32+
33+
Teacher:
34+
35+
- Checkpoint:
36+
`/home/lorenzo/moleworks/checkpoints/terra-solo-resmap64-r1r2-terminalfix-mb32-unmasked-4gpu-50B-20260516-euler-4gpu-2026-05-16-12-18-34.pkl`
37+
- Architecture: base delayed ResMap64, `537,689` parameters.
38+
- Actor remains unmasked; `action_mask` is not used to mask PPO logits.
39+
40+
Student warm start:
41+
42+
- `train_mixed.py` accepts `--teacher_checkpoint` and `--imitation_updates`.
43+
- The teacher rolls out actions in the live env.
44+
- The student minimizes KL from teacher logits plus a weighted teacher-value
45+
regression loss.
46+
- The PPO start RNG, env state, previous actions, and previous reward are saved
47+
before imitation and restored afterward. This makes the distill-vs-scratch
48+
comparison a policy-parameter warm-start test rather than a teacher-driven
49+
curriculum/env-state warm start.
50+
- The imitation optimizer uses its own learning rate, then Adam is reset before
51+
PPO starts. This keeps PPO optimizer state from inheriting distillation
52+
transients.
53+
- The final imitation metrics are logged on W&B step `0`, and the post-distill
54+
student checkpoint is saved as `checkpoints/<name>_POST_DISTILL.pkl`.
55+
56+
Architecture presets:
57+
58+
- `model_size=medium`, `map_feature_dim=192`:
59+
- ResNet channels `(24, 48, 64)`, two blocks/stage.
60+
- Local measured size: `951,081` parameters.
61+
- Local 4090 smoke passes at `1024` envs/GPU, `32` steps,
62+
`64` minibatches.
63+
- `model_size=large`, `map_feature_dim=256`:
64+
- ResNet channels `(32, 64, 96, 128)`, three blocks/stage.
65+
- Local measured size: `3,157,145` parameters.
66+
- Local 4090 smoke passes at `1024` envs/GPU, `32` steps,
67+
`128` minibatches, with `XLA_FLAGS=--xla_gpu_autotune_level=0`.
68+
69+
## Local Validation
70+
71+
- `python -m py_compile train_mixed.py utils/models.py utils/utils_ppo.py`
72+
passed.
73+
- `bash -n scripts/euler/terra_train_larger_resnet_4gpu.sbatch` passed.
74+
- `git diff --check` passed.
75+
- Tiny CPU imitation+PPO smoke passed after the PPO-start reset fix:
76+
- `model_size=medium`, `4` envs, `2` steps, one imitation update, one PPO
77+
update.
78+
- Imitation loss: `2.3258`; PPO completed and saved both
79+
`_POST_DISTILL.pkl` and `_FINAL.pkl`.
80+
- Local RTX 4090 medium smoke:
81+
- Failed with `1024` envs/GPU and `32` minibatches due to ResNet activation
82+
OOM.
83+
- Passed with `64` minibatches.
84+
- Imitation loss: `2.0055`; one PPO update completed in `60.25s`,
85+
`545.87` steps/s.
86+
- Local RTX 4090 large smoke:
87+
- Failed without autotune mitigation with `CUDNN_STATUS_INTERNAL_ERROR`.
88+
- Passed with `XLA_FLAGS=--xla_gpu_autotune_level=0` and `128` minibatches.
89+
- Imitation loss: `1.8941`; one PPO update completed in `63.94s`,
90+
`513.88` steps/s.
91+
92+
## Euler Experiment Matrix
93+
94+
Script:
95+
96+
- `scripts/euler/terra_train_larger_resnet_4gpu.sbatch`
97+
98+
Runs to launch:
99+
100+
- `RUN_KIND=medium_distill`
101+
- `model_size=medium`, `map_feature_dim=192`, `num_minibatches=64`.
102+
- `imitation_updates=200`.
103+
- Tests whether a moderately larger ResNet benefits from the teacher warm
104+
start.
105+
- `RUN_KIND=medium_scratch`
106+
- Same medium architecture and minibatches, no teacher.
107+
- Control for the warm-start hypothesis.
108+
- `RUN_KIND=large_distill`
109+
- `model_size=large`, `map_feature_dim=256`, `num_minibatches=128`.
110+
- `imitation_updates=100`.
111+
- Uses `--xla_gpu_autotune_level=0` because the local 4090 smoke otherwise
112+
hit a cuDNN internal error.
113+
- `RUN_KIND=large_scratch`
114+
- Same large architecture and autotune mitigation, no teacher.
115+
- Control for whether the large model specifically benefits from imitation.
116+
117+
Every Slurm job runs:
118+
119+
- hard GPU guard for exactly four RTX 3090/4090 GPUs;
120+
- `check_jax_runtime.py --min-devices 4`;
121+
- one W&B-disabled full-shape smoke before the online W&B run;
122+
- online PPO with `50B` global env steps, eval/checkpoint every `100`
123+
updates.
124+
125+
## Decision Criteria
126+
127+
Early health:
128+
129+
- The W&B-disabled full-shape smoke must complete update 1.
130+
- Online row 0 should have entropy close to `ln(8)` unless imitation has made
131+
the actor deliberately sharper.
132+
- No synchronized timeout value-loss spike pattern should reappear.
133+
- The smoke now exercises tiny eval and checkpoint paths before starting the
134+
online run.
135+
136+
Learning comparison:
137+
138+
- Compare medium-distill against medium-scratch by W&B update, not wall clock.
139+
- Compare large-distill against large-scratch by W&B update, not wall clock.
140+
- The warm-start path should show earlier nonzero completion/success or better
141+
productive terrain-change/action without collapsing entropy.
142+
- Large-distill is not a pure warm-start ablation; it tests whether the much
143+
larger spatial encoder can exploit the same teacher signal without becoming
144+
unstable.

0 commit comments

Comments
 (0)