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`:

```bash
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)

```bash
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)

```bash
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:

```text
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

```bash
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

```bash
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

```text
# 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
```