README_ablation.md

Variable-Speed Ablation Sweeps

Workflow for running ablation studies that compare different target-speed sets, different ways of integrating the speed signal into the model, and different soft-prompt capacities, against a shared LIBERO evaluation harness.

For the underlying single-pipeline workflow (build one dataset, train one model), see VARIOUS_SPEED_README.md.

1. Research questions

The default ABLATIONS table in scripts/run_ablations.py answers three research questions:

1. Step size — holding the speed range fixed at [0.5, 2.0], does denser speed coverage improve generalization?
2. Range — holding step size roughly fixed, how much does the speed range matter? Where does the policy break?
3. Integration strategy — given a fixed training speed set, does it matter how the speed signal reaches the model: as text in the instruction prompt, as a continuous feature that modulates the action expert (adaRMSNorm), or as a learned soft-prompt inserted between the image and the language tokens?

Plus a hyperparameter sweep on the soft-prompt capacity:

4. Soft-prompt P — how many learnable tokens per speed anchor are actually needed?

Speed-set sweep (text conditioning across all entries)

name	speeds	role
g1_baseline	[1.0]	no augmentation, baseline
g2_coarse	[0.5, 1.0, 1.5, 2.0]	wide range, coarse step (0.5)
g3a_step025	[0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0]	wide range, fine step (0.25)
g4_narrow	[0.75, 1.0, 1.25, 1.5]	narrow range probe
g5_extreme	[0.25, 0.5, 1.0, 2.0, 4.0]	extreme range probe

Speed-integration sweep (fixed speeds [0.75, 1.0, 1.25, 1.5])

name	speed_integration	model surgery
speedint_text	text	none (instruction prompt only)
speedint_modulation	modulation	requires Pi0Config.speed_modulation=True -> MLP head + adaRMSNorm in action expert
softprompt_p8	soft_prompt	K=4 × P=8 learnable tokens (also serves as the P=8 arm of the P-sweep below)

Soft-prompt P-length sweep (same speeds, varies tokens per anchor)

name	P
softprompt_p1	1
softprompt_p4	4
softprompt_p8	8
softprompt_p16	16
softprompt_p32	32

g4_narrow, speedint_text, speedint_modulation, and all five softprompt_p* entries share the same speed set (0.75, 1.0, 1.25, 1.5) -- the runner builds that dataset only once.

Run-time dedup

• Build dedup: ablations with identical speeds reuse one built dataset directory.
• Norm-stats dedup: ablations whose effective asset_id collides reuse one norm_stats.json. The five softprompt_p* entries declare shared_norm_key="softprompt_shared" because norm stats depend only on the dataset, not on the soft-prompt parameter P.
• For the default 12-ablation table, this means 5 dataset builds, 8 norm stats computes, and 12 training runs.

Note on data-volume confound

Each ablation produces len(speeds) × source_episodes training samples, so larger speed sets train on more total data. When comparing groups, prefer plotting metrics against samples seen (or epochs), not raw step count, to disentangle "more granular speeds" from "more data".

2. Calibrate eps once (data-driven cleaning thresholds)

Before the first build, profile the source dataset to pick --clean-*-eps:

uv run python scripts/profile_action_norms.py \
  --src "$SRC" \
  --output "$SRC/action_norm_profile.json"

If you want near-zero cleaning, use the printed P1 / P5 percentiles as --clean-eps overrides. The default is 0.0 (no cleaning), since LIBERO demos never fall below 1e-4 in practice. Set --clean-eps 1e-4 (or the percentile value) if you bring in a noisier dataset.

3. Two ways to drive a sweep

(a) Data-prep only (separate machine)

export SRC=/path/to/libero_data
export ROOT=/path/to/ablation_outputs

uv run python scripts/build_ablation_datasets.py \
  --src "$SRC" --out-root "$ROOT" \
  --num-workers 16 \
  --clean-eps 0.0  # 0 = no cleaning (LIBERO default); set 1e-4 for noisier data

This shells out to scripts/build_libero_speed_dataset_mp.py once per unique speed set. Each dataset lands at $ROOT/libero_speed_<speed_token>_<run_tag>/.

(b) End-to-end (build + norm-stats + train)

uv run python scripts/run_ablations.py \
  --src "$SRC" --out-root "$ROOT" \
  --train-config pi05_libero_various_speed_all \
  --base-asset-id libero_various_speed_all_pi05 \
  --exp-prefix pi05_ablation \
  --num-train-steps 30000 \
  --build-num-workers 16 --train-num-workers 8

Useful flags: --only NAME[,NAME] to run a subset, --skip-build / --skip-norm-stats / --skip-train to bypass stages, --dry-run to print commands without running them. Always do a --dry-run once before the first real run to sanity-check the generated CLI strings (especially --model.soft-prompt-speeds 0.75 1 1.25 1.5 which must be space-separated).

4. How configs are scoped per ablation

A single TrainConfig name (e.g., pi05_libero_various_speed_all) is reused across all ablations. Per-ablation differences are applied at the CLI via tyro's overridable_config_cli:

Stage	Override
build	--speeds <ab.speeds> and per-(speeds) --dst
compute_norm_stats	--repo-id <dataset_dir> --asset-id <effective_asset_id>
train_pytorch	--data.repo-id, --data.assets.asset-id, --data.speed-integration, --exp-name, --eval-speed-set, plus per-ablation extra_train_args

effective_asset_id = "<base>_<ablation.name>" unless Ablation.shared_norm_key is set, in which case it is "<base>_<shared_norm_key>". Norm stats are written to assets/<train_config_name>/<effective_asset_id>/norm_stats.json.

--eval-speed-set makes the wandb per-speed loss breakdown follow the ablation's own speeds. It only fires when use_flow_control=True (i.e. the modulation path); for text and soft_prompt paths it's harmless dead-code that nonetheless stays in sync.

5. Diagnostic outputs after each build

Three artifacts under <dataset_dir>/meta/:

• cleaning_summary.json — frame counts and ratios that were near-zero and got zeroed by clean_near_zero_actions. Cross-check against the percentiles in action_norm_profile.json from §2.
• replay_summary.json — per-target-speed mean/median/max of integrated translation/rotation L2 error vs. source, path-length ratios, padded ratio, and gripper-switch delta. Also printed to stdout at end of build.
• speed_metrics.jsonl — one row per (source_episode, target_speed) with the same raw fields as above; useful for plotting.

If transl_L2 / rot_L2 medians are not at floating-point noise level (~1e-7), or if gripper_delta sum != 0, something is wrong with that build.

6. LIBERO evaluation across 8 GPUs

For each trained model, evaluate at multiple speeds — including out-of-distribution values — to measure interpolation vs. extrapolation:

training:   g3a_step025 ([0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0])
evaluation: [0.4, 0.5, 0.6, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.5, 3.0]

For each speed, run all four LIBERO suites on the 8-GPU server. Use scripts/eval_libero_8gpu.sh, which partitions the work to roughly equal wall-clock per GPU:

GPU	suite	task_ids	episodes
0	libero_spatial	all	500
1	libero_goal	all	500
2	libero_object	all	500
3	libero_10	0,1	100
4	libero_10	2,3	100
5	libero_10	4,5	100
6	libero_10	6,7	100
7	libero_10	8,9	100

libero_10 (libero_long) has the longest rollouts (max_steps=520), so splitting it 5 ways balances total wall-clock against the three short suites at max_steps ≤ 300.

Step 1: launch one policy server per GPU

for g in 0 1 2 3 4 5 6 7; do
  CUDA_VISIBLE_DEVICES=$g uv run python scripts/serve_policy.py \
    policy:checkpoint --policy.config=pi05_libero_various_speed_all \
    --policy.dir=checkpoints/<run>/<step> --port=$((8000 + g)) &
done

Step 2: launch the 8-way eval driver

SPEED=1.25 ./scripts/eval_libero_8gpu.sh

This dispatches scripts/eval_libero_speed.py 8 times in parallel (one per GPU/server), each writing per-episode results to results/libero_eval_<speed>x_<ts>/<label>_<speed>x.json and a video directory under the same root.

What gets recorded per episode

scripts/eval_libero_speed.py records:

• success (bool, from env.step returning done=True before max_steps)
• steps (int, policy steps actually executed, excluding the num_steps_wait warmup frames)
• task_id, episode_idx, task_description, suite, speed

Successful rollouts terminate as soon as the env returns done; failures run to max_steps for the suite.

What gets printed per rank and globally

Each rank prints a summary line like:

[rank=0] libero_spatial speed=1.25x  success=412/500 (82.4%)  mean_steps_success=87.3 median=82.0  mean_steps_failure=220.0  mean_steps_all=110.6

After all 8 ranks finish, the driver auto-aggregates per suite (merging the 5 libero_10 shards back into one row) and globally:

--- per-suite rollup ---
  libero_spatial    success=412/500 (82.4%)  mean_steps_success=87.3   mean_steps_all=110.6
  libero_goal       success=...
  libero_object     success=...
  libero_10         success=...   (merged across ranks 3-7)

GLOBAL (speed=1.25): success=1342/1500 (89.5%)  mean_steps_success=128.5  mean_steps_all=156.8

Use the mean_steps_success vs. mean_steps_all gap as a fast read-out: when a policy fails, it tends to walk to the time limit, so mean_steps_all rises sharply while mean_steps_success stays roughly constant.

How speed reaches the policy

The eval client adds both "speed": float(speed) and "speed_label": "1p25x" (etc.) to the element dict it sends to the websocket server. Whichever integration strategy was trained, the corresponding pipeline consumes the appropriate key:

• text → data-side SpeedConditionedPrompt reads speed_label and rewrites the instruction prompt
• modulation → model-side reads observation.speed (raw scalar) and feeds it through an MLP that fuses with the timestep embedding to drive adaRMSNorm in the action expert
• soft_prompt → model-side reads observation.speed for nearest-anchor lookup over the K learnable token groups

modulation and soft_prompt share the same observation.speed field end-to-end; there is no separate flow_control channel anymore.

For OOD speeds at eval time, text benefits from token-level extrapolation, modulation is naturally continuous, but soft_prompt falls back to the nearest training anchor (e.g., speed=2.0 with anchors (0.75, 1.0, 1.25, 1.5) clamps to 1.5). This is by design and is the main reason to compare the three strategies on the same OOD test set.

7. Implementation notes

The speed-integration paths (modulation + soft-prompt) were implemented / refactored across these files (branch 0502_mp_process):

• src/openpi/models/model.py — Observation.speed is the single end-to-end channel for raw target speed. The legacy flow_control field (and the LogSpeedControl data transform that produced it) was removed.
• src/openpi/models/pi0_config.py — added speed_modulation: bool (replaces the old flow_control_dim: int), soft_prompt_speeds (tuple of K anchor speeds), and soft_prompt_p (tokens per anchor).
• src/openpi/models_pytorch/pi0_pytorch.py:
  • __init__ registers two optional sub-modules:
    • speed_mod_mlp_in/out + speed_condition_mlp_in/out when speed_modulation=True. Reads raw observation.speed (shape (B, 1)), pushes through MLP, fuses with timestep embedding, sends to action expert as adarms_cond.
    • soft_prompt_tokens: nn.Parameter of shape (K, P, paligemma_width) with N(0, 0.02) init, plus a non-persistent buffer soft_prompt_anchors: tensor(K,), when soft_prompt_p > 0.
  • embed_prefix does an argmin nearest-anchor lookup, then inserts (B, P, hidden) soft-prompt tokens between vision and language with full attention.
  • embed_suffix reads speed directly (raw scalar) for the modulation path; the MLP head is wide enough to learn any monotonic transform internally so log-scaling is unnecessary.
  • _preprocess_observation, forward, sample_actions, and denoise_step plumb speed through.
• src/openpi/models_pytorch/preprocessing_pytorch.py — passes speed through SimpleProcessedObservation. (Previously it silently dropped flow_control, which is why the old modulation path appeared half-baked.)
• src/openpi/policies/libero_policy.py — LiberoInputs passes speed through to the observation dict.
• src/openpi/models/pi0.py — JAX Pi0 mirrors the PyTorch rename (speed_modulation, speed_mod_mlp_*, speed_condition_mlp_*, reads obs.speed).
• src/openpi/training/config.py — LeRobotVariousSpeedLiberoDataConfig has a single high-level speed_integration: Literal["text", "modulation", "soft_prompt"] switch (no more auto, use_flow_control, or use_speed_prompt plumbing). The pi05_libero_various_speed_all_flow_prompt/noprompt configs were replaced by pi05_libero_various_speed_all_modulation.
• tests/test_soft_prompt_smoke.py — light tests for config validation and the argmin nearest-neighbor logic. The full forward-pass test is gated to manual GPU runs because PaliGemma weights are heavy.

Token budget

For PI0.5, max_token_len=200 counts only the language tokens (image and soft-prompt tokens come from a separate budget). With P=8, the soft prompt adds 8 prefix tokens; P=32 is still well under typical attention budgets. For sanity, pi0_config.py:max_token_len only constrains the tokenized prompt, not the visual or soft-prompt embeddings.

Inference at OOD speeds (soft-prompt only)

embed_prefix uses an argmin lookup, so a speed not in soft_prompt_speeds falls back to the nearest training anchor. If you want linear interpolation between two adjacent anchors instead, modify the lookup in pi0_pytorch.py:embed_prefix. (text and modulation paths do not have this limitation.)

8. Adding a new ablation

1. Append Ablation(name, speeds, ...) to the ABLATIONS tuple in scripts/run_ablations.py. Use a short unique name -- it is embedded in dataset directory, asset_id, and exp_name.
2. (Optional) If the new group uses very different speeds than what profile_action_norms.py was last run for, re-run it.
3. build_ablation_datasets.py --only <new_name> to produce the dataset (or skip if it shares speeds with an existing one).
4. run_ablations.py --only <new_name> --skip-build to compute norm-stats and train.
5. After training, eval_libero_8gpu.sh to evaluate at the speeds you care about.

9. Files added by this workflow

# New
scripts/profile_action_norms.py            # action-norm profiler (eps calibration)
scripts/build_libero_speed_dataset_mp.py   # multi-process speed dataset builder
scripts/build_ablation_datasets.py         # data-prep stage of ablation sweep
scripts/run_ablations.py                   # end-to-end ablation runner
scripts/eval_libero_speed.py               # LIBERO eval client w/ speed + step tracking
scripts/eval_libero_8gpu.sh                # 8-GPU eval driver
tests/test_soft_prompt_smoke.py            # smoke tests for soft_prompt logic

# Updated
scripts/build_libero_speed_dataset.py      # writes cleaning_summary.json + replay_summary.json
scripts/compute_norm_stats.py              # accepts --repo-id / --asset-id overrides
scripts/train_pytorch.py                   # eval_speed_set config-driven; wandb image logging removed
src/openpi/models/model.py                 # Observation.speed field
src/openpi/models/pi0_config.py            # soft_prompt_speeds / soft_prompt_p fields
src/openpi/models_pytorch/pi0_pytorch.py   # soft_prompt model surgery
src/openpi/models_pytorch/preprocessing_pytorch.py  # bug fix: flow_control + speed pass-through
src/openpi/policies/libero_policy.py       # speed pass-through
src/openpi/training/config.py              # eval_speed_set, speed_integration fields
src/various_speed/core.py                  # cleaned_any/both metrics + ratios