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.