Attaching segment metadata to a LeRobot v2.1 dataset

opentau.scripts.attach_metadata produces a copy of a LeRobot v2.1 dataset with three new per-frame parquet columns and two new per-episode episodes.jsonl fields. Downstream LeRobotDataset loaders pick these up automatically and expose them as optional keys in the Standard Data Format.

What it writes

Given a source dataset at --root and a JSON description of per-episode segments at --annotations, the script writes the modified dataset to --copy-to:

• data/chunk-XXX/episode_XXXXXX.parquet gains three columns:

  • response (string) — the current segment’s subtask text.

  • memory (string) — a cumulative memory summary produced by opentau.scripts.pi_mem_data_generator (or a placeholder when --skip-memory is passed).

  • mistake (int64, 0/1) — not segment.success for the current segment.

• meta/info.json gains feature entries for the three new columns.

• meta/episodes.jsonl gains quality (int, 1–5) and segments (list of segment start frame indices) per episode.

• Stats files (meta/episodes_stats.jsonl, meta/stats.json) are left untouched — the new columns do not participate in normalization.

The script never modifies the source dataset in place. --copy-to is required.

Annotations JSON schema

The top-level JSON value is a list of episode annotations. Episode IDs must be a subset-equal match to the dataset’s episode_index values (extras or missing episodes are rejected).

A runnable example lives at configs/examples/attach_metadata_annotations.json. Its shape, in brief:

[
  {
    "episode_id": 0,
    "quality": 3,
    "segments": [
      {"start": 0,  "subtask": "approach the cup", "success": false},
      {"start": 50, "subtask": "pick up the cup",  "success": true}
    ]
  },
  {
    "episode_id": 7,
    "quality": 5,
    "segments": [
      {"start": 0, "subtask": "place the block", "success": true}
    ]
  }
]

Validation rules (all enforced before any file is written):

• episode_id must be an int and exactly match one dataset episode.

• quality must be an int in 1..5.

• segments must be non-empty.

• segments[0]["start"] must equal 0.

• Each start must be an integer frame index (floats are rejected to catch “time in seconds” mistakes).

• Segment starts must be strictly increasing.

• The last start must be less than the episode length.

• success must be a Python bool (truthy strings/ints are rejected).

mistake is derived as not segment.success; a successful segment produces mistake=0, a failure produces mistake=1.

Usage

Offline smoke / unit tests (no OpenAI calls, memory column is filled with the placeholder f"memory{frame_index}"):

python -m opentau.scripts.attach_metadata \
    --root /path/to/source/dataset \
    --annotations /path/to/annotations.json \
    --copy-to   /path/to/output/dataset \
    --skip-memory

Production run (memory summaries are generated by OpenAI). Place your key in a .env at the repo root (or export OPENAI_API_KEY):

# .env
OPENAI_API_KEY=sk-...

python -m opentau.scripts.attach_metadata \
    --root /path/to/source/dataset \
    --annotations /path/to/annotations.json \
    --copy-to   /path/to/output/dataset \
    --model gpt-4o-mini \
    --delay 0.1

Memories are generated before any parquet is written, so an OpenAI failure leaves the destination partially copied but the parquet tree untouched — rerun with --overwrite to retry cleanly.

CLI arguments

Flag	Description
--root	Required. Path to the source LeRobot v2.1 dataset root.
--annotations	Required. Path to the annotations JSON (schema above).
--copy-to	Required. Destination dataset root.
--overwrite	Optional. Allow clobbering --copy-to if it already exists.
--skip-memory	Optional. Skip OpenAI calls; memory column becomes f"memory{frame_index}". Useful for tests and offline runs.
--model	OpenAI chat completion model name. Defaults to gpt-4o-mini (also overridable via OPENAI_MODEL).
--delay	Seconds to sleep between OpenAI calls (rate limiting).

Consuming the metadata at training time

Once a dataset is annotated, a plain LeRobotDataset loader reading from --copy-to will automatically surface the new columns as optional keys in the standard data format — no changes to training code are needed beyond configuring the optional-key dropout probabilities on DatasetMixtureConfig (see Optional Standard-Format Keys).