opentau.datasets.speed_percentiles

Per-task percentile bucketing for the speed optional key.

For each (dataset, task) pair we rank the episodes by length-in-frames and bucket each episode’s length into one of {0, 10, 20, ..., 100} based on where it falls in its task’s distribution. Shorter episodes get lower bucket labels (= “faster”); longer episodes get higher labels.

The 10 percentile boundaries (p5, p15, …, p95) per task are persisted to meta/speed_percentiles.jsonl so the compute happens at most once per dataset on disk. Existence of the file is the sole gate; staleness (after meta/episodes.jsonl is appended to) is detected by comparing the on-disk per-task n_episodes total against the current load — a mismatch logs a WARNING but the file is still trusted (delete it to force recompute).

Tasks with fewer than MIN_EPISODES_FOR_PERCENTILES distinct episode lengths are written with "percentiles": null and bucket every episode to SPARSE_TASK_BUCKET (= 50). The distinct-count threshold catches both small-N and degenerate (all-equal-length) distributions in one rule.

Functions

`bucket_episode_length`(length, percentiles)	Map an episode length to its bucket label.
`compute_task_percentiles`(...)	Compute p5..p95 of episode lengths for each task.
`episode_to_task_index_from_hf_dataset`(...)	Build `{ep_idx: task_idx}` by reading the per-frame parquet.
`load_or_compute_speed_percentiles`(root, ...)	Return per-task percentile lookup for the dataset rooted at `root`.

opentau.datasets.speed_percentiles.bucket_episode_length(length: int, percentiles: list[float] | None) → int[source]

Map an episode length to its bucket label.

Parameters:

length – Episode length in frames.
percentiles – The 10 percentile boundaries for this episode’s task, or None when the task is too sparse to rank.

Returns:

One of SPEED_BUCKET_LABELS (i.e. {0, 10, 20, ..., 100}). Sparse tasks (percentiles is None) return SPARSE_TASK_BUCKET.

Tie semantics: length == p_X lands in the upper bucket, e.g. length == p25 → 30. Uses np.searchsorted(side='right') so the bucket boundary intervals read as [p_{i-1}, p_i) with the lowest bucket being (-inf, p5) and the highest [p95, +inf).

opentau.datasets.speed_percentiles.compute_task_percentiles(episode_lengths_per_task: dict[int, list[int]]) → dict[int, list[float] | None][source]

Compute p5..p95 of episode lengths for each task.

Parameters:: episode_lengths_per_task – Maps task_index to the list of episode lengths (in frames) for episodes belonging to that task.
Returns:: Dict mapping task_index to a length-10 list of floats [p5, p15, ..., p95] (ascending), or None for tasks with fewer than MIN_EPISODES_FOR_PERCENTILES distinct lengths. The None sentinel covers both the small-N case (one or two episodes) and the degenerate all-equal-length case (a fixed-length synthetic dataset). Both produce zero-information rankings, so the bucket lookup falls back to SPARSE_TASK_BUCKET.

opentau.datasets.speed_percentiles.episode_to_task_index_from_hf_dataset(hf_dataset: datasets.Dataset, episodes: list[int], episode_data_index: dict[str, torch.Tensor], epi2idx: dict[int, int], valid_task_indices: set[int]) → dict[int, int][source]

Build {ep_idx: task_idx} by reading the per-frame parquet.

Each row of the per-episode parquet carries an authoritative integer task_index column (written by LeRobotDataset._save_episode_table() and consumed by LeRobotDataset.__getitem__()). Resolving the per-episode task this way bypasses any drift between episodes.jsonl::tasks[0] (a string) and tasks.jsonl::task (the string keyed by task_index) — a paraphrased tasks.jsonl no longer breaks the lookup.

Only the first row of each episode is read; the parquet’s task_index is constant within an episode by construction.

Episodes whose resolved task_index is not present in valid_task_indices (i.e. the index points to a task that tasks.jsonl doesn’t define — genuine metadata corruption, not paraphrasing) are skipped with a deduped WARNING. Skipped episodes are still trained on; they just fall back to SPARSE_TASK_BUCKET in the downstream speed-bucket lookup, which already tolerates a missing episode_to_task_index entry.

Parameters:

hf_dataset – The dataset’s per-frame HF Dataset (parquet-backed).
episodes – Selected episode indices (the dataset’s self.episodes).
episode_data_index – {"from": Tensor, "to": Tensor} giving the start/end row in hf_dataset for each selected episode.
epi2idx – Maps episode_index to its position in episode_data_index["from"].
valid_task_indices – All task_index values defined in tasks.jsonl (i.e. set(meta.task_to_task_index.values())). Used solely to detect parquet rows that point at a task the metadata doesn’t know about.

opentau.datasets.speed_percentiles.load_or_compute_speed_percentiles(root: Path, episode_lengths: dict[int, int], episode_to_task_index: dict[int, int], task_to_task_index: dict[str, int]) → dict[int, list[float] | None][source]

Return per-task percentile lookup for the dataset rooted at root.

If root / meta/speed_percentiles.jsonl already exists, load it verbatim — staleness is accepted by design (a WARNING is logged when the on-disk file was computed from fewer episodes than the current load is using, i.e. episodes were appended since; subset loads, where the current load uses fewer episodes than on-disk, are silent because the on-disk percentiles are a more robust sample). The file is trusted in both cases — delete it to force a recompute. Otherwise compute the percentiles from episode_lengths and persist the file.

Append-only migration: when the on-disk file is missing rows for tasks that do have resolved episodes in this load, percentiles are computed for the missing tasks only and appended to the file (logged at WARNING so the rewrite is visible). Existing rows are preserved verbatim. This recovers datasets whose existing percentile file was written before episode_to_task_index_from_hf_dataset() — those files were missing rows for tasks whose episodes.jsonl::tasks[0] string didn’t match tasks.jsonl, so every episode of those tasks bucketed to SPARSE_TASK_BUCKET. The append-only design avoids clobbering existing percentiles with subset-derived samples when the current load is a mixture / subset; force a full recompute by deleting the file.

Distributed-safe: every rank computes the percentiles in-memory (the result is deterministic from the inputs), but only the main process writes the file; the Accelerator.wait_for_everyone() barrier afterwards ensures non-main ranks don’t sail past the function while the file is still mid-rename. Each rank returns its own in-memory copy rather than re-reading the just-written file. The write itself is atomic (per-writer UUID-suffixed tmp file + os.replace) so two independent processes (e.g. concurrent training jobs sharing a dataset root) can’t corrupt each other’s output even outside the rank-gated case.

Persistence schema (one line per task):

{"task_index": 0, "task": "pick up the red block",
 "n_episodes": 47, "percentiles": [120.0, 145.0, ..., 340.0]}
{"task_index": 3, "task": "open the drawer",
 "n_episodes": 2, "percentiles": null}

On read-only roots (write raises OSError / PermissionError) the in-memory dict is still returned and a warning is logged once per root.

Parameters:

root – Dataset root directory (the parent of meta/).
episode_lengths – {ep_idx: length_in_frames} — typically LeRobotDataset.episode_lengths so the percentile compute and the per-episode bucket pre-fill share one source.
episode_to_task_index – {ep_idx: task_idx} — typically built via episode_to_task_index_from_hf_dataset().
task_to_task_index – {task_string: task_index} — used only to denormalize the task string into each row for human inspection.

Returns:

Dict mapping task_index to its 10 percentile boundaries (or None for sparse tasks). Indexed lookup uses bucket_episode_length().