in-context-learning: human↔robot pairing + RICL on EgoVerse pi0.5

[RyanPCo]

This branch adds the in-context learning line of work on top of main.

1. Human↔robot episode pairing (ed870df5, f2adfcd3)

Dense-language-annotation pairing between human (aria) and robot (eva) episodes for side-by-side eval: egomimic/scripts/human_robot_pairs.json (Tier-1 co-located alignment sets + Tier-2 similar-task object-overlap pairs), the matcher, and a findings doc for future agents.

2. RICL — retrieval-based in-context learning on pi0.5 (5b6a442a)

kNN-retrieved in-context demonstrations injected into the pi0.5 prefix, reusing the existing zarr / Cartesian / trainHydra / PI-algo / DINOv3 stack. For each query observation we retrieve the k nearest demos and add their (image, state, action) to the prefix so the policy can imitate a task it never trained on. Cross-embodiment: retrieval bank = aria (human), query = eva (robot).

Key point — no PI0Pytorch surgery. The flow pi0.5 embed_prefix iterates over all images in the observation and embeds the full prompt, with the entire prefix attended bidirectionally (and EgoVerse already feeds it a variable image set). So PIRicl only: (1) appends the k retrieved base_0_rgb frames as extra entries in the obs images dict, and (2) splices each retrieved demo's discretized (state, action) into the prompt text (same binning as the pi0.5 State block). eva (14-D) and aria (12-D, no gripper) already share one 32-D action space via the converters.

Added

• egomimic/ricl/ — retrieval (DINOv3 → cKDTree → per-query top-k cache), conditioning (the prefix surgery), data (collate + frame-idx wrapper + bank provider), metrics (+ CPU unit tests for each)
• egomimic/algo/pi_ricl.py (PIRicl(PI), three small overrides) and egomimic/eval/pi_ricl_eval.py (PIRiclEval: retrieval vs zero-context floor)
• RiclDataModuleWrapper in pl_utils/pl_data_utils.py; configs model/pi0.5_ricl, data/cotrain_pi_ricl, evaluator/eval_pi_ricl
• egomimic/ricl/README.md — design notes + cluster runbook

Status. All logic is CPU-verified locally (23 checks: retrieval smoke, conditioning, data collate, metrics; both Hydra configs validated; ruff-clean; imports resolve in the emimic env). The finetune (from the openpi pi0.5 base), DINOv3 embedding at scale, and the D0 (eva→eva sanity) → D1 (aria→eva) eval are GPU-cluster steps documented in the README — they can't run on the dev box (no openpi/CUDA).

🤖 Generated with Claude Code

[RyanPCo]

• in-context-learning: human↔robot pairing + RICL on EgoVerse pi0.5 #491 👈 (View in Graphite)
• main

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

[github-actions]

Claude Code Review

Review of PR #491: RICL on EgoVerse pi0.5

Summary

Adds retrieval-based in-context learning to pi0.5 via a thin PIRicl(PI) subclass that splices k-NN-retrieved (image, state, action) demos into the prefix without touching PI0Pytorch. Clean separation, good test coverage at the unit level, but several cluster-side correctness questions remain.

Key concerns

1. frame_idx semantics may misalign with retrieval cache

RiclQueryDataset derives frame_idx from MultiDataset.index_map[idx] -> (dataset_name, local_idx). The docstring notes this is "best-effort" with a clamp fallback. However:

• The retrieval cache is built per-frame from DINOv3 embeddings indexed 0..T-1 over the raw episode. If local_idx is the chunk start frame but the embedding cache indexes raw frames, this works. If MultiDataset uses any subsampling/filtering (it often does), local_idx and the embedding row index will silently drift, and retrieval will return neighbors for the wrong observation.
• This needs explicit verification against ZarrDataset.__getitem__ and the embedding pipeline before training. Recommend a runtime assertion that the cache's per-episode T matches the dataset's episode length.

2. Bank-side normalization is unverified

ZarrBankFrameProvider reads raw observations.state.ee_pose and actions_cartesian from the bank zarr and applies only converter.to32(action). The retrieved state is never passed through any normalization, but the comment in data.py claims it's "normalized to the query's 32-D convention." This is a bug:

• The query's state in the prompt is normalized by norm_stats before discretization (in PI._discretize_state_for_sample). Retrieved state must go through the same normalization to live in [-1, 1] and yield comparable bins.
• Without this, retrieved state bins will saturate at 0 or 255, making the in-context "state" tokens meaningless or actively misleading.

Same concern for actions: the converter's to32 does layout mapping, but does it also produce values in the [-1, 1] normalized range that pi0.5's action discretizer expects? If not, the same saturation issue applies.

3. Embedding key inconsistency

• DINOV3_ZARR_KEY = "observations.embeddings.dinov3.front_1" in retrieval.py
• README runbook uses --output-keys observations.embeddings.dinov3.front_1 ✓
• But ZarrBankFrameProvider reads images.front_1 (not observations.images.front_1). This images.front_1 key path is unusual for EgoVerse zarr — please confirm against the actual zarr layout (typically observations.images.{cam}).

4. Coordinate frame at retrieval time

The codebase convention is poses stored in SLAM world frame, re-expressed to head frame at training time. Retrieved state comes raw from zarr (observations.state.ee_pose) without any frame transform. If the query state has been transformed to the head frame by the training pipeline but the retrieved state is in world frame, bin comparison is meaningless. This needs the same transform_lists applied as the query.

5. Image normalization mismatch

build_ricl_collate scales uint8 → [0, 1], then augment_images_with_retrieved checks if img.dtype not int → maps [0,1] → [-1,1]. This works but is fragile across two modules. One float path through _to_minus1_1 already in augment_images_with_retrieved handles uint8 directly — consider passing uint8 through and centralizing the normalization in one place.

6. Token budget estimation under-counts

estimate_prompt_tokens assumes ~1.3 tokens per number, but PaliGemma's tokenizer often splits 3-digit numbers into 2-3 tokens. At k=4, action_steps=1, state_dim=32, action_dim=32: actual usage is likely closer to 600+ tokens. The bump to max_token_len=512 may silently truncate. Run a tokenizer measurement on a real prompt before training.

7. Missing test: retrieval correctness end-to-end

All tests use mocks. There's no test that:

• A built cache + provider round-trips a known (state, action) into the prompt correctly.
• The actual PIRicl._build_prompts output is parseable / token-count is within budget.

The smoke test in retrieval.py --smoke exercises only the index, not the integration.

Suggestions

1. Add a --validate mode to the retrieval CLI that loads the cache + opens one bank zarr and prints (state, action, image) for a sampled query, so cluster users can sanity-check before training.
2. Normalize retrieved state/action through norm_stats in ZarrBankFrameProvider (or in the collate). Pass the algo's norm_stats into the provider for the appropriate embodiment.
3. Apply the same coordinate-frame transforms to retrieved state as the query loader. Reuse the existing transform_lists from the embodiment config.
4. Assert cache/episode length consistency at RiclDataModuleWrapper.__init__ or first batch: for each cached query hash, confirm cache.T == dataset_episode_length.
5. Image key: confirm images.front_1 vs observations.images.front_1 against actual zarr; align with image_zarr_key default.
6. Measure tokenizer length on a sample prompt and either raise max_token_len or document the truncation behavior. Convert the warning into a hard error if estimated > max.
7. Add frame_idx to the base ZarrDataset return dict rather than wrapping; the current best-effort approach is brittle and hides the index-map dependency.
8. AGENTS.md: miss

---

Reviewed by Claude · Review workflow