[RyanPCo]

ZarrDataset.__getitem__ now makes one zarr call per key, not one per keymap entry. Several entries name the same zarr key with different windows. In the mecka keypoint keymap at proprio_history=4, left/right.obs_keypoints and the two wrist poses are each read once as a forward action chunk (i, i+H) and once as a backward proprio window (i-K+1, i+1). That was ten reads over six distinct keys, each paying zarr's per-call metadata hop.

Behaviour changes

Output is unchanged, with one exception: an array shorter than a window now raises a ValueError naming the key. The old per-entry slice returned the short rows and padded them. A lookup miss here would otherwise escape as StopIteration, which the DataLoader reads as the end of the epoch.

Speed

Measured on a real mecka episode with that keymap, back to back in one process:

	before	after
read path, per sample	27.5 ms	15.6 ms (1.76x)
read calls	10	1
metadata hops	10	6

Read is ~70% of __getitem__. It also sets the epoch-boundary stall, since an epoch's first batch is one worker building the whole batch.

How reads are planned

• ZarrEpisode.plan_reads merges the windows on one key into a single contiguous read only when every row between them lies in a chunk one of them already touches. Otherwise the groups are fetched together with one orthogonal index (ZarrEpisode.read_intervals).
• Numeric keys are chunked 100 rows, so their windows always merge. The numbers above are this path.
• JPEG keys are one frame per chunk, where a plain union would read and ship every frame between two distant windows. The image-history window loader: image history (a past front frame, gap in seconds) and the episode fps in every sample #678 adds on the front camera is that case.
• The union is sliced first and padded second (_pad_to_horizon), so a chunk running off the end of the episode pads exactly as before.
• Each slice is copied. A slice of a shared span is a view where a per-entry read was not, and transforms may write into these arrays.

Tests

tests/unit/test_loader_read_coalescing.py:

• Compares against a per-entry oracle over a ramped episode ([t, d] = t*1000 + d). The shared synthetic fixture holds every numeric key constant in time, so it can't see a wrong frame offset. The oracle covers 5 keymaps x 9 indices x {ZarrDataset, ZarrAnnotationCutoffDataset}.
• Covers two paths nothing else in the repo does: a key read as both a window and a single frame, and the JPEG-decode retry, which changes idx and must recompute the windows and the read.
• The oracle compares values, so it would pass a loader that read the whole episode. Separate tests count the rows read per key (test_rows_read_are_the_union_of_the_windows, test_far_apart_jpeg_frames_are_two_rows_not_the_gap) and pin plan_reads on both chunk sizes.
• Also: one read per distinct key, independence from keymap order, no aliasing between slices, and test_an_array_shorter_than_its_window_fails_loudly.

🤖 Generated with Claude Code

https://claude.ai/code/session_01HjCQBqQCr2igFEMR7yUovw

[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

• pi: initialize from PaliGemma with a scratch action expert #677 : 2 dependent PRs (#678 , #680 )
• perf: vectorized image augs and a torch.compile knob for HPT and Pi; Pi clips gradients like openpi #676
• loader: one zarr read per key, not one per keymap entry #674 👈 (View in Graphite)
• eval: build the pinned-episode video batches only where they are used #673
• checkpoints: last.ckpt as a symlink, not a second full copy #672
• trainHydra: metric_frames_per_episode: auto, and warn on empty prompts #671
• transforms: per-step wrist framing for proprio history #670
• viz: one keypoint renderer, the horizon trace #669
• configs: history and backbone-LR knobs on the keypoint wrist HPT model #668
• trainer: ddp_bf16_compress, DDP with a bf16 gradient all-reduce #667
• eval: overlay videos play at the source rate under DDP #666
• hpt: build the policy on CPU and let Lightning place it #664 : 1 other dependent PR (#665 )
• hpt: never re-initialise pretrained stems; hash-check them; mask text padding #645 : 1 other dependent PR (#646 )
• obs: proprio history (K proprio steps, one token each, learned time embedding) #650
• loss: mask the repeat-last action padding out of the HPT loss #649
• resnet: keep a frozen backbone in eval mode; give backbones their own LR group #652
• augs: per-sample image augmentation for HPT train-mode augs #651
• eval: deterministic validation (seeded prompts + flow sampling); log-only grad-norm MAD #644
• data: mecka_fold_flagship_top3_hpt_6d, the 3-operator flagship train split #643
• val: subsampled metric loaders + pinned contiguous video loaders #642
• infra: mecka 6D / keypoint recipes, 300M models, defaults, and the rollout revert #640
• eval: layout-aware frame-invariant metrics, and a train-split viz head #659
• data: stop letting one bad episode or a rotation channel take a run down #658
• norm: guard degenerate channels, horizon pooling, and no rot6d normalization #657
• keypoints: the 144-D wrist-first hand-keypoint action becomes the human default #656
• mecka: fix the double-mirrored LEFT wrist frame, and correct existing zarrs #655
• 6d: PI goes 6D-only; Eva and Human gain the continuous-rotation modes #654
• 6d: continuous rotation representation, bimanual layouts, shared norm helpers #653
• train: give each Slurm slurm job gets its own cache dir to avoid conflict #638 : 1 other dependent PR (#639 )
• write git version into the run dir for reproducibility #637
• train: don't load PI base weights on resume and verify that checkpoint path exists #636
• configs: declare proprio and actions dims once per HPT model config #634 : 1 other dependent PR (#635 )
• cache norm stats for data configs to minimize recomputation #633
• import cloud packages lazily to support local training #631 : 1 other dependent PR (#632 )
• cache dataset resolver for faster training initialization #630
• validate episode hashes in data config before training #628
• modify pi camera map to match data #629
• add tests to ensure training works between every model and data vendor #626
• run CI tests on every PR #625
• add env variables for user-specific paths #624
• build: set pyproject and uv as the dependency managers #623
• fix(configs): fix paths for data hydra yaml configs #622
• refactor(algo): delete ACT #627 : 1 other dependent PR (#621 )
• tests: create dedicated test folder with unit and integration tests #620
• add pyright type checker #619
• main

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

[RyanPCo]

nit (dead-code): ZarrDataset no longer uses _pad_sequences. Its only caller is zarr_dataset_action_expert.py:121, which pads a single key. Call self._pad_to_horizon(raw[zarr_key], horizon) there and delete _pad_sequences.

[RyanPCo]

Automated adversarial review (read-only; findings were checked by a second skeptical pass before posting): 2 nit.

Comments that could not be anchored to a diff line:

• egomimic/rldb/zarr/zarr_dataset_multi.py:2321 nit (duplication): _read_windows' (start, end-or-None) encoding is decoded twice (start + 1 if end is None else end here and in the intervals loop). The stride is re-parsed with int(spec.get("history_stride", 1) or 1) in _read_windows and again below, though _proprio_history already returns (K, stride). Return half-open (start, stop) plus a single-frame flag from _read_windows, and reuse one (K, stride) helper, so getitem stops re-deriving them. Before this PR the stride was a local.

[RyanPCo]

Automated adversarial review (read-only; findings were checked by a second skeptical pass before posting): 1 nit.

[RyanPCo]

nit (duplication): _read_windows' (start, end-or-None) encoding is decoded twice (start + 1 if end is None else end here and in the intervals loop). The stride is re-parsed with int(spec.get("history_stride", 1) or 1) in _read_windows and again below, though _proprio_history already returns (K, stride). Return half-open (start, stop) plus a single-frame flag from _read_windows, and reuse one (K, stride) helper, so getitem stops re-deriving them. Before this PR the stride was a local.

[RyanPCo]

Review findings: status after the fixup pass. The fixes are squashed into this PR's commits, and the whole stack was force-pushed.

Fixed

• _pad_sequences kept only for zarr_dataset_action_expert: Deleted _pad_sequences. zarr_dataset_action_expert, outside the PR, now calls _pad_to_horizon(raw[zarr_key], horizon).

Partial

• Window semantics (end=None, history_stride parsing) decoded in three places: Fixed in loader: one zarr read per key, not one per keymap entry #674: _read_windows returns (zarr_key, start, stop, single_frame) with stop exclusive, and one _history_spec helper is used by _read_windows, getitem and _proprio_history. rdt: long-range visual memory; image memory window and history stride in seconds #689 then replaces history windows with point-read index arrays (its own (zarr key, window) form); that rewrite was kept as is, so from rdt: long-range visual memory; image memory window and history stride in seconds #689 on the tuple form and _history_spec are gone.