feat(vla): add SmolVLA conditioning and experiment artifacts
This commit is contained in:
Logic
@@ -0,0 +1,242 @@
|
||||
# VLA Experiment Sweep Design
|
||||
|
||||
## Goal
|
||||
|
||||
Run a matched-configuration training sweep on `/home/droid/project/diana_sim/sim_transfer` for the aligned Transformer diffusion policy, comparing:
|
||||
|
||||
1. ResNet ImageNet pretraining vs no pretraining under matched settings.
|
||||
2. Transformer width/depth choices across a 3×3 grid.
|
||||
|
||||
The sweep should use the full dataset, headless rollout validation every 20 epochs, SwanLab logging, and a total budget of about 200 epochs per run.
|
||||
|
||||
## Current Context
|
||||
|
||||
- Dataset: `100` episodes, `70000` frame-level samples.
|
||||
- Cameras used as image conditions: `r_vis`, `top`, `front`.
|
||||
- Current policy family: ResNet visual encoder + Transformer DDPM head.
|
||||
- Current training path already supports:
|
||||
- full-dataset training (`val_split=0.0`)
|
||||
- checkpoint-time and epoch-time rollout hooks
|
||||
- headless eval path
|
||||
- SwanLab logging
|
||||
- best-checkpoint selection by rollout reward once rollout metrics exist
|
||||
- Current trainer behavior that matters for execution:
|
||||
- step checkpoints are written to `checkpoints/` relative to the Hydra job working directory
|
||||
- epoch-time rollout validation is controlled by `train.rollout_val_freq_epochs`
|
||||
- checkpoint-time rollout validation is separately controlled by `train.rollout_validate_on_checkpoint`
|
||||
- per-run log files, pid files, and run directories must be created by the external launcher, not by the trainer itself
|
||||
|
||||
## Derived Schedule
|
||||
|
||||
For a chosen global batch size `B`, define:
|
||||
|
||||
- `steps_per_epoch = floor(70000 / B)` because training currently uses `drop_last=True` when the train set is larger than the batch.
|
||||
- `rollout_freq_steps = 20 * steps_per_epoch`
|
||||
- `max_steps = 200 * steps_per_epoch`
|
||||
|
||||
Example at `B = 32`:
|
||||
|
||||
- `steps_per_epoch = 2187`
|
||||
- `rollout_freq_steps = 43740`
|
||||
- `max_steps = 437400`
|
||||
|
||||
This schedule must be recomputed after the batch-size probe chooses the final shared batch size.
|
||||
|
||||
## Experiment Strategy
|
||||
|
||||
### Recommended approach
|
||||
|
||||
Use one **shared batch size across all runs**, chosen by probing the largest model that will appear in the sweep. This preserves fairness while still using as much GPU memory as safely available.
|
||||
|
||||
### Why not per-run max batch size
|
||||
|
||||
Allowing each model to use a different batch size would add another changing variable, making the 11-run comparison harder to interpret.
|
||||
|
||||
### Unique executions vs logical comparisons
|
||||
|
||||
There are `11` logical comparison points in the user request (`2` for the pretraining comparison and `9` for the emb/layer grid), but only `10` **unique** training configurations because:
|
||||
|
||||
- the pretrained baseline in Section A (`pretrained_backbone_weights=IMAGENET1K_V1`, `n_emb=128`, `n_layer=4`) is also one cell of the Section B sweep.
|
||||
|
||||
This design will therefore run **10 unique training jobs**, and the pretrained-on baseline will be reused as the `(128, 4)` architecture-sweep point. This avoids wasting one full 200-epoch run on an identical configuration.
|
||||
|
||||
The reused baseline has one canonical identity throughout execution:
|
||||
|
||||
- run name stem: `sim-transfer-baseline-pretrain-on-emb128-layer4`
|
||||
- this same run is counted both as:
|
||||
- the pretrained-on arm of the 2-point pretraining comparison
|
||||
- the `(n_emb=128, n_layer=4)` point of the 9-point architecture sweep
|
||||
|
||||
## Batch-Size Probe
|
||||
|
||||
Before launching the full sweep:
|
||||
|
||||
1. Probe the largest planned model:
|
||||
- `n_emb=384`
|
||||
- `n_layer=12`
|
||||
- ResNet pretrained enabled
|
||||
2. Try candidate batch sizes in ascending order:
|
||||
- `32`
|
||||
- `48`
|
||||
- `64`
|
||||
- optionally `80` if `64` is clearly safe
|
||||
3. For each candidate, run a short GPU smoke training with:
|
||||
- real dataset
|
||||
- three cameras
|
||||
- no SwanLab
|
||||
- no rollout validation
|
||||
- `max_steps=4`
|
||||
4. A candidate **passes** if all 4 steps finish without:
|
||||
- CUDA OOM
|
||||
- process crash / abort
|
||||
- NaN / Inf loss
|
||||
5. Choose the **highest passing** candidate and use it for all sweep runs.
|
||||
|
||||
Learning rate should scale linearly using the current training rule:
|
||||
|
||||
- config baseline: `lr=1e-4` at `batch_size=16`
|
||||
- equivalent rule: `lr = 1e-4 * (batch_size / 16)`
|
||||
- equivalently: `lr = 2e-4 * (batch_size / 32)`
|
||||
|
||||
## Experiment Matrix
|
||||
|
||||
Total logical comparison points: `11`
|
||||
|
||||
Total unique training executions: `10`
|
||||
|
||||
### A. ResNet pretraining comparison (2 runs)
|
||||
|
||||
Fixed Transformer settings:
|
||||
|
||||
- `n_emb=128`
|
||||
- `n_layer=4`
|
||||
- `n_head=4`
|
||||
- `freeze_backbone=true`
|
||||
- `use_separate_rgb_encoder_per_camera=true`
|
||||
|
||||
Interpretation note:
|
||||
|
||||
- This comparison keeps the current frozen-backbone recipe unchanged.
|
||||
- Therefore, the `pretrained_backbone_weights=null` run measures **random frozen visual features vs ImageNet frozen visual features**.
|
||||
- It does **not** test end-to-end training of the ResNet from scratch; that would be a separate experiment.
|
||||
|
||||
Runs:
|
||||
|
||||
1. `pretrained_backbone_weights=IMAGENET1K_V1`
|
||||
2. `pretrained_backbone_weights=null`
|
||||
|
||||
### B. Transformer architecture sweep (9 runs)
|
||||
|
||||
Fixed visual backbone setting:
|
||||
|
||||
- `pretrained_backbone_weights=IMAGENET1K_V1`
|
||||
- `freeze_backbone=true`
|
||||
- `use_separate_rgb_encoder_per_camera=true`
|
||||
- `n_head=4`
|
||||
|
||||
Rationale:
|
||||
|
||||
- The 9-point emb/layer sweep fixes `pretrained_backbone_weights=IMAGENET1K_V1`.
|
||||
- The separate 2-run comparison in Section A isolates the effect of ResNet pretraining at the baseline architecture.
|
||||
|
||||
Grid:
|
||||
|
||||
- `n_emb ∈ {128, 256, 384}`
|
||||
- `n_layer ∈ {4, 8, 12}`
|
||||
|
||||
## Shared Training Configuration
|
||||
|
||||
All full runs should use:
|
||||
|
||||
- dataset: `/home/droid/project/diana_sim/sim_transfer`
|
||||
- cameras: `[r_vis, top, front]`
|
||||
- device: `cuda`
|
||||
- `val_split=0.0`
|
||||
- `num_workers=12`
|
||||
- `seed=42`
|
||||
- headless rollout validation
|
||||
- `rollout_val_freq_epochs=20`
|
||||
- `rollout_validate_on_checkpoint=false`
|
||||
- `rollout_num_episodes=3`
|
||||
- `use_swanlab=true`
|
||||
- shared batch size from probe
|
||||
- shared learning-rate rule from probe
|
||||
- `max_steps` recomputed to represent about `200` epochs
|
||||
|
||||
Reproducibility note:
|
||||
|
||||
- `train.seed=42` will be fixed and recorded for all runs.
|
||||
- Under the current `train_vla.py`, this does **not** make the runs bitwise deterministic when `val_split=0.0`, because model initialization and shuffled training order are not fully seeded by the trainer today.
|
||||
- The sweep is therefore treated as a **matched-config comparison**, not a strict deterministic benchmark.
|
||||
|
||||
## Rollout / Validation Policy
|
||||
|
||||
- No held-out val split will be used.
|
||||
- Loss remains a fallback best-model metric until the first rollout reward exists.
|
||||
- After rollout metrics begin, the best checkpoint should be chosen by **higher average rollout reward**.
|
||||
- Rollout should run with:
|
||||
- no GUI
|
||||
- CPU eval device
|
||||
- current evaluation helper path
|
||||
|
||||
## Execution Model
|
||||
|
||||
Runs should be launched **serially**, not concurrently, because:
|
||||
|
||||
- the user currently has one GPU available
|
||||
- concurrent jobs would reduce throughput and complicate memory management
|
||||
- serial execution preserves clean logs and easier resume semantics
|
||||
|
||||
Each run gets:
|
||||
|
||||
- a dedicated run directory under `runs/`
|
||||
- its own log file
|
||||
- its own pid / launcher metadata if backgrounded
|
||||
- a unique SwanLab run name
|
||||
- `hydra.job.chdir=true`
|
||||
- `hydra.run.dir=<run_dir>` so the trainer-local `checkpoints/` land inside the run directory
|
||||
|
||||
## Naming Convention
|
||||
|
||||
Suggested run name pattern:
|
||||
|
||||
- pretraining comparison:
|
||||
- `sim-transfer-baseline-pretrain-on-emb128-layer4-bs{B}-lr{LR}-e200`
|
||||
- `sim-transfer-pretrain-off-bs{B}-lr{LR}-e200`
|
||||
- architecture sweep:
|
||||
- `(128,4)` reuses `sim-transfer-baseline-pretrain-on-emb128-layer4-bs{B}-lr{LR}-e200`
|
||||
- `sim-transfer-emb{E}-layer{L}-bs{B}-lr{LR}-e200`
|
||||
|
||||
## Success Criteria
|
||||
|
||||
The sweep is successful when:
|
||||
|
||||
1. A shared safe batch size has been found.
|
||||
2. All 10 unique runs have launchable commands and isolated run directories.
|
||||
3. Each run trains for about 200 epochs worth of steps.
|
||||
4. Rollout validation occurs every 20 epochs.
|
||||
5. SwanLab records train metrics and rollout rewards.
|
||||
6. Final comparison can be made using best rollout average reward and training traces across:
|
||||
- 2-point pretraining comparison
|
||||
- 9-point emb/layer sweep
|
||||
|
||||
## Non-Goals
|
||||
|
||||
- No change to dataset format.
|
||||
- No change to model architecture beyond the requested pretraining toggle and Transformer width/depth sweep.
|
||||
- No attempt to equalize wall-clock time across models.
|
||||
- No multi-GPU scheduling.
|
||||
|
||||
## Risks and Mitigations
|
||||
|
||||
### Risk: larger models may OOM
|
||||
|
||||
Mitigation: probe the largest model first and choose one shared safe batch size.
|
||||
|
||||
### Risk: long runs may be interrupted
|
||||
|
||||
Mitigation: use per-run directories, periodic checkpoints, and resumable launch commands.
|
||||
|
||||
### Risk: rollout validation increases total training time
|
||||
|
||||
Mitigation: keep rollout headless, CPU-based, and limited to 3 episodes every 20 epochs.
|
||||
Reference in New Issue
Block a user