feat(scripts): add annotate_mistakes.py for per-frame mistake labels

[WilliamYue37]

What this does

Adds src/opentau/scripts/annotate_mistakes.py, a sibling to annotate_subtasks.py. For every episode in a dataset mixture config, the script:

1. Reads the per-frame response column already written by annotate_subtasks.py from each episode parquet, and treats every contiguous run of identical response values as a single subtask segment.
2. Decodes the dataset's camera0 video (resolved with the same lookup chain as annotate_subtasks.py) once per episode and pulls the last frame of each contiguous run — no temporal subsampling, just one frame per segment.
3. Sends that single frame plus the segment's subtask string to a VLM (default: gemini-robotics-er-1.6-preview; Anthropic Claude also supported via --model) and asks for a {"success": bool, "reason": str} JSON verdict.
4. Sets every parquet row in the segment to mistake=1 if the VLM reports failure, 0 otherwise. Any parse / API failure defaults to 0 (no mistake).
5. Atomically rewrites the episode parquet with the new int64 mistake column and registers mistake in meta/info.json features ({"dtype": "int64", "shape": (1,), "names": None}) the first time it is added to a dataset.

Resumability semantics:

• Episodes whose parquet already has a mistake column in the schema are skipped (cheap O(1) check via pq.read_metadata).
• Episodes whose parquet has no response column are skipped with a warning — run annotate_subtasks.py first.
• Frames are still spatially downsampled / center-cropped to --target-size (default 448) to keep image-token cost bounded; we never upsample.

Helpers shared with annotate_subtasks.py (_resize_and_center_crop, _to_jpeg_bytes, _to_b64_jpeg, _is_gemini_model, _resolve_camera0_video_key, _resolve_root, _load_datasets_from_config) are imported rather than duplicated.

How it was tested

• python -c "from opentau.scripts import annotate_mistakes" — module imports cleanly.
• pre-commit run --files src/opentau/scripts/annotate_mistakes.py — all hooks pass (ruff, ruff-format, pyupgrade, bandit, gitleaks, license header).
• No automated tests added — the script is a thin orchestrator over the Anthropic / Gemini SDKs and PyAV decoding, identical in shape to annotate_subtasks.py which also has no tests. End-to-end exercise requires API keys + a real LeRobot dataset that has already been processed by annotate_subtasks.py.

How to checkout & try? (for the reviewer)

# Default Gemini ER path (requires GEMINI_API_KEY or GOOGLE_API_KEY)
python src/opentau/scripts/annotate_mistakes.py \
    --config-path configs/examples/train_mixture_config.json \
    --max-episodes-per-dataset 1

# Claude path
python src/opentau/scripts/annotate_mistakes.py \
    --config-path configs/examples/train_mixture_config.json \
    --model claude-opus-4-7 \
    --max-episodes-per-dataset 1

Checklist

• I have added Google-style docstrings to important functions and ensured function parameters are typed.
• My PR includes policy-related changes.
  • If the above is checked: I have run the GPU pytests (pytest -m "gpu") and regression tests.

Note: Before submitting this PR, please read the contributor guideline.

[claude]

superseded by next review

[claude]

Inline findings below. None are blocking — primary concerns are (1) bool(parsed["success"]) returning True for stringy "false" and (2) info.json being updated before any parquet is annotated. See summary comment for the full list.

[claude]

[claude-review] summary for commit bcfaebb

Latest commit bcfaebb ("address review feedback on #280") cleanly addresses every prior finding from the review history. Verified by re-reading the current state of src/opentau/scripts/annotate_mistakes.py, tests/scripts/test_annotate_mistakes.py, docs/source/tutorials/datasets.rst, and configs/examples/annotate_mistakes_example.json:

• _parse_success_response now rejects non-bool success values, with test_rejects_int correctly anticipating that bool subclasses int (so a literal 1 is not silently accepted).
• Claude max_tokens=32000 → model-agnostic max_tokens=1024, with a comment explaining the choice fits the smallest 3.5-family ceiling without underutilizing larger models.
• _call_gemini_single carries a comment documenting why the prior thinking_budget=0 workaround was retired (default budget verified empirically to fit reasoning + one-line JSON).
• _query_subtask_success returns (success, ok); _annotate_episode returns (processed, n_api_failures, n_missing_frames); the per-dataset summary log surfaces both totals so quiet API outages and short-video frame misses no longer masquerade as "no mistakes found".
• mistake feature registration in info.json is deferred inside the per-episode loop until after the first parquet has actually been rewritten — matches the resumability claim in the docs.
• New tests/scripts/test_annotate_mistakes.py covers _find_response_runs (empty / single / multiple runs / None / empty-string / alternating / singleton) and _parse_success_response (plain true/false, JSON / bare fences, surrounding whitespace, stringy "false" rejection, stringy "true" rejection, int 1 rejection, missing key, non-object, invalid JSON).
• --hub-cache-dir help text explains the deliberate match with annotate_subtasks.py for cache reuse; the docs flag table mirrors the same note.
• configs/examples/annotate_mistakes_example.json pins "revision": "v2.1" on the Hub repo_id, mirroring the docs' v2.1-only support note. Tutorial section and module docstring both link the new example file.

Stale inline findings from prior reviews have been deleted (the GitHub dismissals API rejects COMMENTED reviews, so deletion is the only way to clear line-anchored findings carried forward across runs).

No blocking issues found.

[claude]

Inline findings below. None are blocking — primary concerns are (1) bool(parsed["success"]) returning True for stringy "false" and (2) info.json being updated before any parquet is annotated. The new docs commit only adds tutorial prose; the code findings carry over unchanged. See summary for the full list.

[claude]

Latest commit e12aa55 only adds configs/examples/annotate_mistakes_example.json — annotate_mistakes.py is byte-identical to db60953, so the prior six inline findings carry over unchanged on the same lines. One new nit on the example config itself: the docs (docs/source/tutorials/datasets.rst) and the script's own docstring (src/opentau/scripts/annotate_mistakes.py:53) still cite configs/examples/train_mixture_config.json rather than the newly-added example. See summary for the full list.

[claude]

Inline findings below — none blocking. Latest commit's thinking_budget=0 workaround is correct; flagged a follow-up since annotate_subtasks.py calls the same Gemini ER family without it. See summary comment for the full list.

[claude]

Inline findings below — none blocking. Latest commit (c7eeaf9 removing the output token caps) is correct in both branches; raised one minor follow-up on the hardcoded Claude max_tokens=32000. Prior code-level findings carry over unchanged on the new line numbers (re-anchored after the +2-line shift in _call_claude_single). See summary for the full list.

[claude]

Inline findings below — none blocking. Latest commit re-enables Gemini ER thinking; this resolves the prior cross-script inconsistency vs annotate_subtasks.py. Remaining findings carry over with re-anchored line numbers (-6 lines after the gemini_config simplification).

[WilliamYue37]

@claude fix

[claude]

Latest commit bcfaebb cleanly addresses every prior finding: _parse_success_response now strict-validates isinstance(success, bool) (with a test that rejects int 1 since bool subclasses int); Claude max_tokens dropped from 32000 to a model-agnostic 1024; Gemini-thinking-retired comment added; (success, ok) tuple plumbed through so API failures surface in the dataset summary; n_missing_frames likewise; write_info deferred until after the first parquet rewrite; 20 unit tests added covering both helpers; --hub-cache-dir help text now explains the shared default; example config pins revision: v2.1 and is referenced from the tutorial. No new blocking issues. Cleared all prior inline comments; see summary for the updated state.