add custom configs to specify tasks for viz

[RyanPCo]

No description provided.

[RyanPCo]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• add custom configs to specify tasks for viz #481 👈 (View in Graphite)
• remove rotation bounds checks for norm stats #480 : 1 other dependent PR (#490 )
• optional viz videos on train dataset #476
• optional flag to generate validation videos for each task #475
• pi train human new #473
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[github-actions]

Claude Code Review

PR #481 Review

Summary

Replaces the one_video_per_task boolean with an explicit tasks allowlist + videos_per_task knob on EvalVideo, threaded through data.viz_tasks in data configs. Also adds an OpenPI tokenizer fallback for PI when HF access is unavailable, and bumps the CPU PACE partition.

Key concerns

1. Unrelated changes mixed into a "viz config" PR. The OpenPI tokenizer fallback in egomimic/algo/pi.py and the submitit_cpu_pace.yaml partition swap (cpu-small → cpu-medium) have nothing to do with the stated PR purpose. These should be separate PRs — the partition change in particular can silently break queued preprocessing jobs for collaborators and deserves its own review/CI confirmation.

2. HPT eval inherits tasks/videos_per_task but PI was the only path that forwards batch["task"]. You added task forwarding in hpt.py (good), and eval_hpt.yaml now selects ${oc.select:data.viz_tasks,[]}. Confirm HPTEvalVideo's validation path actually reads batch[emb]["task"] the same way PIEvalVideo does — otherwise enabling viz_tasks on an HPT run will hit the KeyError("Per-task video output requires 'task'...") only at first validation, well into training. Worth a quick smoke test.

3. _task_video_chunks splits a single contiguous buffer into N near-equal slices. With videos_per_task=3 and max_frames_per_task=1000, you buffer 3000 contiguous frames then cut into 3 sequential mp4s. Those three videos will be three temporally-adjacent windows from the same validation pass — not three diverse samples across the task, which is presumably what a user asking for "3 videos per task" wants. If the intent is variety, you'd want to sample episodes/batches into separate buckets at validation_step time. At minimum, document the current behavior in the config comment so users don't get surprised.

4. _OpenPIPaligemmaPromptTokenizer is missing methods the rest of the pipeline likely calls. AutoTokenizer instances expose pad_token_id, eos_token_id, vocab_size, decode, batch_decode, etc. The fallback only implements __call__. If anything downstream (logging, decoding generated tokens for debugging, the PI model's own internals) touches those attributes, it will AttributeError only on machines without HF access — i.e., exactly the machines this fallback was added to support. Either add the missing surface or constrain the fallback to a known-narrow code path.

5. oc.select:data.viz_tasks,[] default semantics. oc.select returning [] when data is set but doesn't define viz_tasks is the desired behavior; just confirm data.viz_tasks is not required at the schema level on data configs that don't define it (e.g., mecka_pi_eval.yaml's comment now says "set data.viz_tasks or evaluator.tasks" but mecka_pi_eval.yaml itself doesn't set it — verify that resolves cleanly).

Suggestions

• Split this into three PRs: (a) viz task config (the actual stated change), (b) PI tokenizer fallback, (c) PACE partition bump.
• Add a one-line test or at least a manual-run note: PI eval with data.viz_tasks=[task_a, task_b] writes exactly two mp4s named task_a.mp4, task_b.mp4.
• Update the comment in eval_pi.yaml / eval_hpt.yaml to clarify that videos_per_task>1 produces sequential chunks of the buffered frames, not independent samples.
• For the tokenizer fallback, gate behind an explicit env flag (e.g., EGOMIMIC_OFFLINE_TOKENIZER=1) rather than catching OSError — silent fallback to a different tokenizer implementation during training is a footgun.
• Confirm HPTEvalVideo's validation_step reads batch[emb]["task"]; if not, this PR enables a config path that crashes at first validation.
• Justify the cpu-small → cpu-medium change in a commit message or PR description; the original comment said "confirm before first submit" and that confirmation isn't visible here.

Verdict

Request Changes — primarily to split out the unrelated tokenizer and SLURM changes, and to verify HPT actually forwards task end-to-end so the new config default doesn't break HPT eval runs.

---

Reviewed by Claude · Review workflow