VLAwithVariousSpeed / SPEED_TRANSFORM_ALGORITHM.md
Alan0928's picture
Upload folder using huggingface_hub
08ff31f verified
|
Raw
History Blame Contribute Delete
18.2 kB

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)

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:

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:

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

For each output bin j:

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

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

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

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

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

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

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

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