roboimi/docs/superpowers/specs/2026-03-30-vla-experiment-sweep-design.md

# 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.