SPEED_TRANSFORM_ALGORITHM.md

# Variable-Speed Trajectory Transformation — Algorithm Reference

This document describes the **action-space speed transformation** that turns
one LIBERO source episode at speed 1.0x into one or more re-timed episodes
at target speeds (e.g., 0.75x, 1.25x, 1.75x). The implementation lives in
`src/various_speed/core.py`; build orchestration is in
`scripts/build_libero_speed_dataset{,_mp}.py`.

For higher-level usage, see `VARIOUS_SPEED_README.md` and
`README_ablation.md`. This file documents the **algorithm itself**.

---

## 1. Motivation

A LIBERO demonstration is a sequence of 7-D delta-EE-pose commands at fixed
control frequency. To train a VLA that can execute the same task at multiple
speeds, we need supervision data at those speeds. We obtain it by re-timing
existing 1.0x demonstrations: integrate the cumulative motion over different
time windows so the policy issues the same total motion in a different number
of controller steps.

The naive approach — uniformly stretch / compress every action — fails for
two reasons:
1. Direction reversals (e.g., approach → grasp → retreat) cause the
   integrated mid-trajectory motion to nearly cancel out, distorting both the
   action shape and the integrated displacement. **Solution**: segment the
   episode at motion-class transitions and direction reversals; treat each
   segment as locally-linear and re-time it independently.
2. Gripper open/close are discrete events with a roughly constant physical
   duration; resampling them with the same speed factor as motion changes
   their wall-clock timing in a non-physical way. **Solution**: anchor every
   gripper switch as a hard boundary; never compress / stretch a switch event.

## 2. Input / Output

**Source format** (per episode):
- `actions[T, 7]` — 7-D delta action, columns
  `[dx, dy, dz, droll, dpitch, dyaw, gripper]`. The first 6 are continuous
  delta poses; column 6 is a discrete gripper command (≈ ±1).
- `states[T, S]` — proprioceptive state at each control step.
- `source_frame_indices[T]` — original LeRobot frame indices.

**Output format** (per (source_episode, target_speed) pair, length T'):
The output preserves the input's 7-D action / state schema and adds
metadata columns: `speed`, `speed_index`, `speed_label`, `valid_mask`,
`observation_mask`, `action_mask`, `is_padded`, `segment_id`, `motion_class`,
`source_episode_index`, `source_frame_index`, `source_step_index`,
`source_index`, `cleaned_translation`, `cleaned_rotation`. See §11 for
semantics.

## 3. Pipeline overview

```
   actions[T, 7], states[T, S], speed s
     │
     ▼  ① clean near-zero noise   (clean_near_zero_actions)
     │
     ▼  ② segment by motion class (segment_actions)
     │
     ▼  ③ for each segment:        (_resample_segment)
     │     ├─ build resampling boundaries  (_segment_boundaries)
     │     ├─ integrate actions over each bin (cumulative + lin-interp)
     │     ├─ assign per-output state, gripper, masks
     │     └─ tag chunk-start (chunk_aligned mode only)
     │
     ▼  ④ concatenate segments → output episode
     │
     ▼  ⑤ compute per-episode metrics (replay fidelity, segment stats,
            cleaning stats)
```

The five stages are explained below.

---

## 4. Stage ① — Near-zero action cleaning

Function: `clean_near_zero_actions(actions, transl_eps, rot_eps)`.

For each frame, if `‖action[:3]‖ < transl_eps` set translation to zero;
similarly for rotation. Returns the cleaned `actions` and a `(T, 2)` boolean
mask recording which frames had translation / rotation zeroed.

**Defaults**: `clean_transl_eps = clean_rot_eps = 0.0` (cleaning OFF). LIBERO
demos are clean enough that no frame falls below 1e-4. To enable on noisier
data, set positive eps via CLI.

The cleaning mask is propagated to the output as the
`cleaned_translation` and `cleaned_rotation` columns (per-output-frame),
indexed by which source frame each output represents.

## 5. Stage ② — Segmentation

Function: `segment_actions(actions, config)` returns
`list[(start, end, motion_class)]` with `end` exclusive.

### 5.1 Per-frame motion class (`_motion_class`)

For each frame:
- `has_translation = ‖action[:3]‖ ≥ transl_eps`
- `has_rotation    = ‖action[3:6]‖ ≥ rot_eps`

Class:
| code | name | condition |
|---|---|---|
| 0 | `still` | neither |
| 1 | `translate` | only translation |
| 2 | `rotate` | only rotation |
| 3 | `translate_rotate` | both |

The gripper channel (column 6) is **never** used in this classification.

### 5.2 Segment boundaries

Walking through frames, a new segment starts at frame `i` whenever:
- **Class change**: `class[i] != class[i-1]`, OR
- **Direction reversal within same class**:
  - For class 1 or 3: `cos(action[i-1, :3], action[i, :3]) < direction_cos_threshold`
  - For class 2 or 3: `cos(action[i-1, 3:6], action[i, 3:6]) < direction_cos_threshold`

Default `direction_cos_threshold = -0.25` (≈ 104° reversal). A class-3 segment
splits if **either** the translation OR the rotation reverses sharply (OR
semantic, not AND), to keep within-segment motion homogeneous so the
cumulative-integration assumption holds.

### 5.3 Min-segment merge

When `min_segment_len > 1`, runs of segments shorter than the threshold get
folded into the previous segment with the previous segment's class. Default
`min_segment_len = 1` (no merging).

---

## 6. Stage ③ — Boundary determination

Function: `_segment_boundaries(src_actions, speed, *, chunk_aligned=False)`
returns a sorted, deduplicated array of float boundary positions in
source-time. The number of output frames within a segment is
`n_out = len(boundaries) - 1`; output bin `j` covers source-time interval
`[boundaries[j], boundaries[j+1])`.

Two boundary modes, plus a universal gripper-anchor step:

### 6.1 Default mode (uniform linspace)

```text
boundaries = np.linspace(0, n_src, n_out + 1)
where n_out = max(1, round(n_src / speed))
```

Outputs are evenly distributed in source-time. Per-output time-span =
`n_src / n_out ≈ speed`.

### 6.2 Chunk-aligned mode (`chunk_aligned_observation=True`)

Decompose `speed = q / p` as a small-denominator rational
(`Fraction(speed).limit_denominator(32)`), so `q` source frames map exactly
to `p` output frames per chunk. For all ablation speeds (0.25, 0.5, 0.75, 1.0,
1.25, 1.5, 1.75, 2.0, 4.0), `p ≤ 4`.

Algorithm:

```text
n_full = n_src // q                # number of complete chunks
for k = 0, 1, ..., n_full - 1:
    boundaries += linspace(k*q, (k+1)*q, p + 1)   # p outputs per chunk

leftover = n_src - n_full * q      # remaining source frames (< q)
if leftover > 0:
    # Passthrough: each leftover source becomes one output verbatim
    boundaries += [n_full*q, n_full*q + 1, ..., n_src]
```

**Effect on full chunks**: each output integrates exactly `q/p` source frames
(magnitude = `q/p` × typical source action). Chunk-start outputs land at
integer source positions `k*q`, giving exact state-action temporal alignment.

**Effect on leftover (passthrough)**: each leftover source frame becomes
**one** output with action / state / gripper copied verbatim. These outputs
behave at 1.0x for that small trailing region (typically < q < 8 frames at
1.75x).

### 6.3 Gripper-switch anchoring (both modes)

After the mode-specific boundaries are built, every source-time index where
`|diff(action[:, 6])| > 0.5` (gripper close ↔ open transition) is **added**
to the boundary set. The boundaries are then deduplicated and sorted.

This guarantees that no output bin straddles a gripper switch — the switch
event always lands on a boundary, so the discrete close/open command is
preserved in the output as its own distinct controller step. This is what
makes `gripper_switch_delta == 0` true by construction.

---

## 7. Stage ④ — Per-segment resampling

Function: `_resample_segment(actions, states, source_frame_indices, start, end,
motion_class, speed, chunk_aligned_observation=False)`.

For each segment `[start, end)` with `n_src = end - start`:

### 7.1 Action computation (cumulative + linear interpolation)

Build the cumulative sum of the 6-D motion delta:

```text
cumulative[k] = sum of src_actions[0:k, :6]   (shape (n_src+1, 6))
```

For each output bin `j`:

```text
left_t  = boundaries[j]
right_t = boundaries[j+1]

out_action[j, :6] = interp(cumulative, right_t) - interp(cumulative, left_t)
```

where `interp` is piecewise-linear interpolation between integer cumulative
positions. **Mathematical property**: for any partition of `[0, n_src]`,
`sum(out_actions[:, :6]) == sum(src_actions[:, :6])` to floating-point
precision — total integrated motion is **exactly** preserved.

### 7.2 State assignment

```text
source_local       = floor(left_t)             # clipped to [0, n_src-1]
out_state[j]       = src_states[source_local]
out_source_frames[j] = src_frame_indices[source_local]
```

In chunk-aligned mode, at every chunk-start (where `left_t = k*q`) this
resolves to an exact source frame index — no offset. In default mode there
can be sub-step offsets but state-action is still consistent (state at
left edge, action over the bin).

### 7.3 Gripper handling

```text
grip_local         = ceil(right_t) - 1         # clipped to [0, n_src-1]
out_action[j, 6]   = src_actions[grip_local, 6]
```

Gripper is taken from "just before the bin's right edge". Combined with
gripper-switch anchoring (§6.3), this guarantees that the close/open event
appears in the output as its own distinct step at the same wall-clock-relative
time as in the source.

### 7.4 Observation mask

The mask marks which output frames are valid training samples (mask=1) vs
padding to be skipped (mask=0).

**Default mode**: at slow speeds (`speed < 1.0`) the same source frame is
the source for multiple outputs; the first occurrence is mask=1 and
duplicates are mask=0. At equal/fast speeds, mask=1 everywhere.

**Chunk-aligned mode** (`chunk_aligned_observation=True`):

```text
n_full                 = n_src // q
passthrough_threshold  = n_full * q
for j in range(n_out):
    left_t = boundaries[j]
    if left_t >= passthrough_threshold and n_full > 0:
        # Passthrough region: every output is a 1:1 verbatim source frame
        mask[j] = 1
    elif left_t == k*q for some integer k with k*q < n_src:
        # Chunk-start in the full-chunk region
        mask[j] = 1
    else:
        # Padding inside a chunk
        mask[j] = 0
```

The chunk-start check uses `abs(left_t / q - round(left_t / q)) < 1e-6` for
robust float comparison, and is unaffected by gripper-anchor-induced
extra boundaries (those land on non-chunk-aligned positions and stay
mask=0).

### 7.5 Other per-output fields

- `source_step_index = start + source_local` (global source-frame pointer)
- `segment_id` is overwritten by the caller (`transform_episode`)
- `motion_class` is broadcast across all outputs of the segment

---

## 8. Stage ⑤ — Episode-level aggregation

`transform_episode` concatenates per-segment outputs and adds:

- `speed`: broadcast scalar = target speed
- `action_mask`: all 1 (reserved for future loss masking)
- `is_padded`: `1 - observation_mask`
- `cleaned_translation` / `cleaned_rotation`: indexed from
  `clean_mask[source_step_index]`

Replay-fidelity metrics (`compute_replay_metrics`) are then computed on the
full output, plus segmentation stats (counts, length distribution, motion
class distribution). These get aggregated across the dataset into
`meta/cleaning_summary.json`, `meta/segment_summary.json`, and
`meta/replay_summary.json`.

---

## 9. Mathematical guarantees

These hold regardless of speed value, segment length, or mode:

| Invariant | Always | Reason |
|---|---|---|
| `sum(out_action[:, :6]) == sum(src_action[:, :6])` (per segment, machine precision) | ✓ | Cumulative + linear interp is exact integration |
| `gripper_switch_count(out) == gripper_switch_count(src)` | ✓ | Gripper-switch anchoring forces every transition onto a boundary |
| `out_state[chunk_start_j] == src_state[k*q]` exactly (chunk_aligned mode) | ✓ | Chunk boundaries = integer multiples of q; floor at integer = identity |
| Output never references a source frame index outside `[0, n_src-1]` | ✓ | `min(..., n_src-1)` clipping in `_resample_segment` |

---

## 10. Configuration reference (`SpeedTransformConfig`)

| Field | Default | Effect |
|---|---|---|
| `transl_eps` | `1e-4` | Translation threshold for `_motion_class`. Below = "still" axis. |
| `rot_eps` | `1e-4` | Rotation threshold for `_motion_class`. |
| `clean_transl_eps` | `0.0` | If > 0, `‖transl‖ < eps` frames have translation zeroed. OFF by default for LIBERO. |
| `clean_rot_eps` | `0.0` | Same for rotation. |
| `direction_cos_threshold` | `-0.25` | Cosine threshold for in-class direction-reversal split (≈ 104°). |
| `min_segment_len` | `1` | Merge sub-threshold segments into previous (1 = no merge). |
| `keep_still_segments` | `True` | If False, segments with `motion_class == 0` are dropped before resampling. |
| `fps` | `20` | Used only to populate the output `timestamp` column. |
| `chunk_aligned_observation` | `False` | See §6.2 / §7.4. Enables chunk-by-chunk boundaries with passthrough leftover. |

---

## 11. Worked examples

All examples assume default thresholds, no cleaning, default segmentation,
and a single segment covering the whole episode (no gripper switches).

### Example A — 1.25x, n_src = 11, default mode

```text
speed = 1.25, n_out = round(11 / 1.25) = 9
boundaries = linspace(0, 11, 10) = [0, 1.222, 2.444, ..., 11]
```

Each output integrates `~1.222` source frames. Output magnitude ≈ 1.25 ×
typical source magnitude. observation_mask = all 1 (fast speed, no
duplicates).

### Example B — 1.25x, n_src = 11, chunk-aligned mode

```text
q, p = 5, 4    (since 1.25 = 5/4)
n_full = 11 // 5 = 2 chunks
leftover = 1

Full-chunk boundaries:
    chunk 0: [0, 1.25, 2.5, 3.75, 5]
    chunk 1: [5, 6.25, 7.5, 8.75, 10]
Passthrough boundaries (1 leftover source):
    [10, 11]

Final boundaries (deduplicated, sorted):
    [0, 1.25, 2.5, 3.75, 5, 6.25, 7.5, 8.75, 10, 11]
n_out = 9
```

Output magnitudes:
- Outputs 0-3: each integrates `1.25` source frames → magnitude `1.25 ×`
- Outputs 4-7: each integrates `1.25` source frames → magnitude `1.25 ×`
- Output 8: integrates exactly `src_actions[10]` (1 source frame) → magnitude `1.0 ×`

observation_mask:
- `j=0`, `j=4`: chunk-starts (left_t = 0, 5 = `k*q`) → mask=1
- `j=8`: passthrough output (left_t = 10 ≥ n_full*q = 10) → mask=1
- Others: mask=0

States:
- `out_state[0] = src_state[0]` (chunk-aligned)
- `out_state[4] = src_state[5]` (chunk-aligned)
- `out_state[8] = src_state[10]` (verbatim passthrough)
- Other outputs have states from the floor formula but are masked out

### Example C — 0.75x, n_src = 11, chunk-aligned mode

```text
q, p = 3, 4    (since 0.75 = 3/4)
n_full = 3 chunks (covers source 0-8)
leftover = 2

Full-chunk boundaries: 3 × [linspace of 5 entries each]
Passthrough boundaries: [9, 10, 11]

n_out = 14, n_valid = 5
```

Trailing 2 source frames (9, 10) become 2 verbatim outputs.

### Example D — 4.0x, n_src = 13, chunk-aligned mode

```text
q, p = 4, 1
n_full = 3 chunks (covers source 0-11)
leftover = 1

Full-chunk boundaries: [0, 4, 8, 12]
Passthrough boundaries: [12, 13]

n_out = 4, n_valid = 4
```

3 outputs at 4.0x (each merges 4 source frames) + 1 passthrough output at
1.0x. `mask = [1, 1, 1, 1]` (full chunks have p=1 so every output is a
chunk-start; passthrough is also valid).

---

## 12. Output schema

All columns in the output parquet, indexed by output frame `j`:

| Column | Dtype | Semantics |
|---|---|---|
| `state` (or `observation.state`) | `float32 [S]` | proprioceptive state at output `j` |
| `actions` (or `action`) | `float32 [7]` | re-timed delta action |
| `timestamp` | `float32` | `j / fps` (output-time) |
| `frame_index` | `int64` | local frame index `j` |
| `episode_index` | `int64` | output episode index |
| `index` | `int64` | global frame index across the dataset |
| `task_index` | `int64` | inherited from source |
| `speed` | `float32` | target speed (broadcast) |
| `speed_index` | `int64` | which speed slot in the build's `--speeds` list |
| `speed_label` | `string` | e.g. `"1p25x"` |
| `valid_mask` | `int8` | (currently == observation_mask) |
| `observation_mask` | `int8` | 1 = valid for training, 0 = padded |
| `action_mask` | `int8` | reserved (always 1 today) |
| `is_padded` | `int8` | `1 - observation_mask` |
| `segment_id` | `int64` | which segment this output came from |
| `motion_class` | `int64` | `_motion_class` of the segment |
| `source_episode_index` | `int64` | source episode |
| `source_frame_index` | `int64` | source frame index used for state |
| `source_step_index` | `int64` | global source step index used for state |
| `source_index` | `int64` | global source dataset row index |
| `cleaned_translation` | `int8` | 1 iff this output's source frame had transl zeroed |
| `cleaned_rotation` | `int8` | 1 iff this output's source frame had rot zeroed |

Plus the original image / wrist_image columns, with mask=0 frames replaced
by black images (so the dataloader can detect and skip).

---

## 13. Diagnostic outputs

After a full build, the following appear under `<dataset_dir>/meta/`:

- `cleaning_summary.json` — counts and ratios of source frames whose
  translation / rotation was zeroed by `clean_near_zero_actions` (per source
  episode, deduplicated across speeds).
- `segment_summary.json` — per-source-episode segment-length distribution
  (mean / median / P10 / P90), motion-class distribution.
- `replay_summary.json` — per-target-speed mean / median / max of integrated
  translation / rotation L2 error vs. source, path-length ratios, padded
  ratio, gripper-switch delta.

The build also prints a single-line summary of each at the end of the run.
Use these to verify that cleaning is doing nothing harmful (LIBERO default),
that segments are reasonable for the chosen thresholds, and that integrated
motion is preserved at machine precision (rot/transl L2 errors should be at
1e-7 or below).