loader: image history (a past front frame, gap in seconds) and the episode fps in every sample by RyanPCo · Pull Request #678 · GaTech-RL2/EgoVerse

RyanPCo · 2026-09-20T16:14:59Z

Lets the loader return a past front-camera frame next to the current one, and puts the episode's fps in every sample. RDT (#680) uses both: the past frame for its 2-frame image history, the fps for its control-frequency token.

Behaviour changes

None by default. Without image_history_gap_s the keymap is unchanged. Every sample gains an fps key (float32) that nothing in this PR reads.

Usage

Human.get_keymap(image_history_gap_s=...) adds a <front key>_hist camera entry, a copy of the front camera's entry with lag_s set.
ZarrDataset reads that entry round(gap x fps) frames back, at least one frame, clamped at the episode start (so the first samples read frame 0).
fps comes from the episode's zarr fps attr, 30 when it is absent.
Nothing consumes <front>_hist here. Under plain HPT it has no stem and is decoded and dropped; the keymap docstring says so.
A norm_mode keymap drops the history camera along with the other images.

Read cost

The past frame names the front camera's zarr key, so the loader's one-read-per-key covers both windows. JPEG keys hold one frame per chunk, so a plain union span would read every frame in between. The read planner from #674 keeps them two rows.

Per sample on a synthetic 30 fps episode of 360x640 JPEGs (median of 7x40 samples, login node, noisy):

gap	union span	read planner
no history	11-12 ms	11-12 ms
0.1 s	21-24 ms	19-23 ms
0.5 s	30-34 ms	20-21 ms
2 s	53-61 ms	19 ms

With the planner the cost is flat in the gap; what remains is the second decode, which history needs.

Tests

tests/unit/test_image_history.py:

the past frame is the gap in the episode's own fps (30, 10 and 60 fps), clamps at the episode start, and is at least one frame back;
the default keymap has no history entry, and samples still carry fps;
norm_mode drops the history camera;
the JPEG rows read at 0.1, 0.5 and 2 s are 2.

🤖 Generated with Claude Code

https://claude.ai/code/session_01UkAZwERXcFKjueZQKfRF6J

RyanPCo · 2026-09-20T16:15:29Z

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

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

ModelCheckpoint with `save_last: true` calls _save_checkpoint a second time for last.ckpt, so every checkpoint event writes the same tensors twice -- ~21 GB a time on Pi 0.5 (the size tests/integration/test_train_step_gpu.py records). It was measured at 9 GB and ~1.3 min on the now-retired ABC-DiT-L, inside an event that blocked the run for 2.7-3.4 min of a 42.3-min 1000-step cycle. "link" writes a relative symlink to the epoch checkpoint just saved instead. Safe here because save_top_k is -1, so the target is never collected, and _should_remove_checkpoint returns False when previous == current, which is always the case for last.ckpt. Resume is unaffected: _requeue_resume_path gates on os.path.isfile, which follows symlinks, and the opscale launcher already resumes through a hand-made symlink. A link can dangle where a file could not: the target deleted by hand, last.ckpt dropped by an rsync without -l, or a preemption between Lightning's unlink and relink. isfile is False in all three, and the requeue path used to warn and silently restart the run from scratch. It now raises when last.ckpt dangles or is missing next to other checkpoints; only an empty checkpoint dir falls back. One behaviour difference worth knowing: with "link" the callback relinks on EVERY train epoch once any checkpoint exists, not only on checkpoint epochs, so there is an extra os.remove + os.symlink + strategy.barrier per epoch. The test drives a real two-epoch fit and asserts both halves: "link" leaves a relative symlink to the newest epoch checkpoint, and `true` leaves a regular file the same size as it -- the copy this removes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01M5dTGNydAZbRP1EsXgbp4W

Two problems on the <head>_video loaders, one of them a correctness bug. EvalVideo._write_video had no rank guard. The val heads come back from val_dataloader() inside CombinedLoaders, which Lightning does not wrap in a DistributedSampler, so every rank iterates every val batch and renders the same frames -- and every rank then wrote the same mp4 path. The file was whichever rank finished last, and could be torn. Separately, the video loaders were iterated in full on every validation pass. ModelWrapper.validation_step already returned before the forward on a non-viz pass, but Lightning had built the batch by then, so limit_val_batches x 64 pinned frames were decoded and discarded on every rank on the ~90% of passes that render nothing: ~1.7 min a pass, ~4% of a 1000-step cycle. _VizGatedLoader yields nothing when its gate is closed, which is on a non-viz pass and on every rank but 0. The val loaders are built once (_EvaluationLoop.setup_data returns early once _combined_loader is set), so the gate cannot be a length; it is an empty iterator, which _Sequential already reads as "move to the next loader". __len__ answers WITHOUT creating an iterator, because creating one forks that loader's workers and starts them prefetching -- doing the skipped work at setup instead, and never freeing it. It asks a fresh _MaxSizeCycle for the length, which stores its loaders and limits and builds no iterator until __iter__ (lightning 2.6.1). Ranks therefore iterate different numbers of val batches. That is safe because what they LOG is not: the metric loaders run everywhere and log the same keys, the video loader logs nothing, so the sync_dist reduction stays symmetric. It does mean the other ranks sit in that reduction while rank 0 renders -- 8.4 min for 76 batches back when all four ranks rendered at once, so well inside the 30-minute process-group timeout, but that timeout is now the ceiling on evaluator.viz_max_batches. trainHydra therefore refuses a head that has pinned video episodes and an evaluator with viz_max_batches=None. Every recipe that pins video episodes already caps it at 10; the evaluators that leave it unset (eval_hpt, eval_hpt_cartesian, eval_hpt_wrist) pair with configs that pin nothing, so they keep the ungated path where every rank renders together, as before. A desync here would hang a campaign rather than fail it, so it is tested against real DDP rather than argued: tests/fixtures/ddp_video_gate.py runs a 2-process gloo fit with the same collective pattern, at num_workers > 0 so the prefetch is visible. Both ranks finish; rank 0 steps the video loader on the viz epoch, rank 1 never touches the pinned episode. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01M5dTGNydAZbRP1EsXgbp4W

ZarrDataset.__getitem__ issued a separate ZarrEpisode.read for every keymap entry, and 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). Ten reads over six distinct keys, each paying zarr's per-call metadata hop. It now makes one zarr call per key and slices each window out of what it read. Measured on a real mecka episode with that keymap, back to back in one process: 27.5 -> 15.6 ms a sample on the read path (1.76x), one read call instead of ten and six metadata hops instead of ten. read is ~70% of __getitem__, which is also what sets the epoch-boundary stall, since an epoch's first batch is one worker building the whole batch. Windows on one key merge into a single contiguous read only when every row between them lies in a chunk one of them already touches (ZarrEpisode.plan_reads); otherwise the groups are fetched together with one orthogonal index. Numeric keys are chunked 100 rows, so their windows always merge and 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 windows far apart -- the image-history window #678 adds on the front camera is exactly that case. Output is unchanged, except that 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, and a lookup miss here would have escaped as StopIteration, which the DataLoader reads as the end of the epoch. The union is sliced first and padded second, so a chunk running off the end of the episode pads exactly as before, and each slice is copied because a slice of a shared span is a view where a per-entry read was not -- transforms are handed these arrays and may write into them. The test compares against a per-entry oracle over a RAMPED episode ([t, d] = t*1000 + d) rather than the shared synthetic fixture, which holds every numeric key constant in time and so cannot see a wrong frame offset at all. 5 keymaps x 9 indices x {ZarrDataset, ZarrAnnotationCutoffDataset}, plus the two paths nothing else in the repo covers: a key read as both a window and a single frame, and the JPEG-decode retry, which changes idx and so 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 and pin plan_reads on both chunk sizes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01M5dTGNydAZbRP1EsXgbp4W

The 300M HPT train step was launch-bound, not compute-bound: ~14,600 CUDA kernel launches per step, 230 ms of CPU against 111 ms of GPU. Two things dominated, and both are fixed here. PerSampleAugs called torchvision's ColorJitter once per image so each sample drew its own factors -- 64 Python round trips a step, each launching a few dozen kernels and copying its factors from the host. BatchedColorJitter draws the factors as one on-device tensor and runs each adjustment once over the batch; a degenerate factor range reproduces torchvision's scalar adjustment exactly (tests/unit/test_image_augs.py). Anything it cannot batch still goes through the per-image loop. model.compile.enabled=true routes to <algo>.compile_for_training(), which uses nn.Module.compile rather than torch.compile(module) so state_dict keys survive -- the wrapper form renames every parameter _orig_mod.* and would break checkpoints written before the flag. That only routes __call__, so HPT compiles the trunk, the flow head's denoiser and the image encoders (the stems and heads are reached through named methods), and PI's call site moves to __call__. On one H100 at batch 64 and the 360x640 frames mecka decodes to: HPT 245 -> 173 ms from the augs, -> 109 ms compiled (2.24x). Pi at batch 16: 827 -> 588 ms compiled (1.41x). In fp32 the compiled loss is bit-identical to eager (bench/check_compile_parity.py). Also here: the flow head's attentions pass need_weights=False so they dispatch to scaled_dot_product_attention; frames decode to float32 instead of numpy's float64 default, halving the collate and H2D traffic (338 -> 169 MB a camera a batch, bit-identical); a per-camera per-step device->host sync is gone; and the batch is no longer deep-copied every step for an OT loss that is off. The all-zero (absent camera) guard moves into _apply_image_augs and becomes per-sample: the old test was one all() over the batch, so a batch mixing present and absent samples normalized the absent ones to -mean/std. The augs now run on every sample and the selector keeps the absent ones at zero, so RNG use no longer depends on the data. Jitter factors are drawn in float32 (uniform_ into bfloat16 gives ~40 distinct values; into uint8 it raises). enable_grad_norm was never passed to ModelWrapper, so pi0.5_base's false did nothing. It is wired now, and Pi follows openpi instead of opting out: it keeps logging the pre-clip grad norm every step and clips at openpi's 1.0, set as gradient_clip_val in the model config and handed to the Trainer unless trainer.gradient_clip_val is set. Pi recipes trained unclipped until now, so this is a behaviour change for every Pi run. One grad-norm pass remains (the second was removed with the log-only MAD detector), not the two bench_hpt.py models. bench/README.md records the harnesses, the numbers, what was tried and did not pay (persistent_workers, pin_memory, more workers, max-autotune), and which transformers each venv carries. The compile comments note that CUDA-graph modes are where a rank can wedge in cudagraph_trees and keep its GPUs after a cancel. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VQa3o51gPcrK8N1LLoqL95

The speed commit changed three things that are not bit-identical to what they replaced, and it only proved one of them. This adds the other two. need_weights=False moves the flow head's attentions onto scaled_dot_product_attention. check_attn_parity.py forces the old materialized-softmax path back on and compares a full fp32 step: loss 197.009293 vs 197.009033 (1.3e-6 relative), 581 gradient tensors at min cosine 0.999999. Reduction order, nothing else. The batched ColorJitter is the one real behaviour change -- it draws the op order once per call rather than once per image, so a batch shares a permutation. check_aug_distribution.py shows a single sample's marginal is untouched (6400 samples a path, two-sample KS over eight per-sample statistics, worst p = 0.61); only a training run can see the within-batch correlation. ab_augs.sh runs both paths at two seeds -- 2000 steps, 32 train episodes, held-out loss every 100 steps on 15 disjoint episodes, eval-mode augs so the held-out pass measures what was learned rather than what was shown. The batched arm is better at both seeds (-9.0%, -23.2%) and the 1.26x wall-clock gap confirms the control arm really took the per-image path. That rules out a regression; it is NOT evidence the rewrite helps. Two seeds cannot carry that, and the arms consume different amounts of RNG, so each arm/seed pair is effectively its own seed. bench/README.md says so where someone would otherwise quote the number. The control arm reuses the shipped fallback: PerSampleAugs keeps its per-image loop for aug lists _vectorize declines, so stubbing _vectorize out runs the old code rather than a reimplementation of it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VQa3o51gPcrK8N1LLoqL95

_hf_offline monkeypatched transformers.utils.hub._is_offline_mode, a module flag that 5.x replaced with is_offline_mode(). setattr on a name the module no longer has raises, so every test standing on the harness errored out before it reached the network: 12 of them, including the whole PI train-step matrix. Verified failing identically on an unmodified tree, so this is the venv's transformers 5.17 meeting a 4.x assumption, not a regression from this stack. Probe for whichever spelling the installed version has rather than pinning either, since the lockfile still says 4.53.2. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Xh3wMw9pXYPHUSGjU7qmix

An arm for training pi0.5 from the released VLM instead of the pi0.5 base checkpoint, to separate what the VLM pretraining buys from what pi0.5's own robot pretraining buys. The prefix (SigLIP tower, Gemma-2B, the projector) comes from google/paligemma-3b-pt-224; everything pi0.5 adds on top -- the Gemma-300M action expert, action_in/out_proj, the adaRMS time MLPs -- stays at its init. pytorch_weight_path and paligemma_weight_path are mutually exclusive rather than ordered by precedence: silently preferring one would make a typo in the other look like it worked. Setting neither now warns instead of loading nothing quietly. The loader streams shard by shard and then fails if any parameter went untouched or any checkpoint tensor went unused, because a partial load of a 3.6B model is invisible in the loss for a long time. Two places where the release does not line up with what openpi builds, both handled in the key remap: - transformers 5 flattened SiglipVisionModel, so vision_tower.vision_model.* is vision_tower.* there; below 5 (the pin, 4.53.2) the hop is still in the module tree. The loader maps onto whichever the live model has. Mapping to the 5.x spelling unconditionally failed 21 SigLIP tensors on the pin, which the key-string tests could not see; a unit test now loads a tiny real PaliGemmaForConditionalGeneration under the installed transformers, and passes on both 4.53.2 and 5.17.0. - The checkpoint pads the vocabulary to 257216, 64 rows past the 257152 openpi allocates (<image> plus 63 unused). embed_tokens and lm_head are truncated; openpi reaches the image slots through the vision tower, never through an embed_tokens lookup, so the dropped rows are unreachable. google/paligemma-3b-pt-224 is license-gated. With paligemma_allow_mirror: true, resolve_paligemma_dir falls back on a 401/403 (not on a rate limit or an outage) to ungated mirrors whose shards are sha256-checked against the canonical repo's LFS metadata first -- a mirror that has drifted is refused, not used. It is off by default: integrity is checked, but a research arm's weights of record should come from the canonical repo unless a run opts in, and this cluster's HF token is authorized for it. Only *.safetensors are pulled. openpi's own loader source is not an option: its GCS bucket 403s anonymously. pi0.5_paligemma_init is an overlay, not a runnable config, so a keypoint or 6D arm is four lines. It changes the init source only; the constant LR stays, since pi0.5_base documents dropping the cosine schedule on purpose. openpi_compat carries three shims for a transformers 5 venv, which openpi's pin (4.53.2, also this repo's) does not expect; all should be deleted once the fork's transformers_replace patch is ported forward. The third restores openpi's fp32 keep-list, which spells the pre-flattening vision paths, so on 5.x the SigLIP patch and position embeddings ran in bf16. On 5.x only this arm loads: the pi0.5 base checkpoint stores the nested SigLIP names and its strict load_model stops at startup. pi05.md says which setup is supported and what differs. Verified with a full trainHydra run to exit 0 on an H100 (4 epochs, all four validation heads, checkpoint and videos written) and the real 12 GB checkpoint into the real model in tests/integration, both on transformers 5.17, plus the unit tests on 4.53.2 and 5.17.0. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Xh3wMw9pXYPHUSGjU7qmix

…isode fps in every sample Human.get_keymap(image_history_gap_s=...) adds a <front key>_hist camera entry; ZarrDataset reads it round(gap x fps) frames back, clamped at the episode start. Every sample also carries the episode's fps (zarr attr). The past frame names the front camera's zarr key, so the loader's one-call-per- key read covers both windows. JPEG keys are one frame per chunk, and a plain union span would read every frame in between; the read planner keeps them two rows. Per sample on a synthetic 30 fps episode of 360x640 JPEGs (median of 7x40 samples, login node, noisy): no history 11-12 ms; gap 0.1 s 21-24 -> 19-23 ms (the second decode, which history needs); 0.5 s 30-34 -> 20-21 ms; 2 s 53-61 -> 19 ms, flat in the gap. A test counts the JPEG rows read at 0.1, 0.5 and 2 s. Nothing consumes <front>_hist here; the keymap docstring says so. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019DbPWJCKtVg5c2PNHs7egw

RyanPCo · 2026-10-12T20:27:38Z

Folded into #680 to keep the stack under Graphite's 50-branch cap; the changes are unchanged there.

ZarrDataset.__getitem__ issued a separate ZarrEpisode.read for every keymap entry, and 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). Ten reads over six distinct keys, each paying zarr's per-call metadata hop. It now makes one zarr call per key and slices each window out of what it read. Measured on a real mecka episode with that keymap, back to back in one process: 27.5 -> 15.6 ms a sample on the read path (1.76x), one read call instead of ten and six metadata hops instead of ten. read is ~70% of __getitem__, which is also what sets the epoch-boundary stall, since an epoch's first batch is one worker building the whole batch. Windows on one key merge into a single contiguous read only when every row between them lies in a chunk one of them already touches (ZarrEpisode.plan_reads); otherwise the groups are fetched together with one orthogonal index. Numeric keys are chunked 100 rows, so their windows always merge and 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 windows far apart -- the image-history window #678 adds on the front camera is exactly that case. Output is unchanged, except that 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, and a lookup miss here would have escaped as StopIteration, which the DataLoader reads as the end of the epoch. The union is sliced first and padded second, so a chunk running off the end of the episode pads exactly as before, and each slice is copied because a slice of a shared span is a view where a per-entry read was not -- transforms are handed these arrays and may write into them. The test compares against a per-entry oracle over a RAMPED episode ([t, d] = t*1000 + d) rather than the shared synthetic fixture, which holds every numeric key constant in time and so cannot see a wrong frame offset at all. 5 keymaps x 9 indices x {ZarrDataset, ZarrAnnotationCutoffDataset}, plus the two paths nothing else in the repo covers: a key read as both a window and a single frame, and the JPEG-decode retry, which changes idx and so 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 and pin plan_reads on both chunk sizes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01M5dTGNydAZbRP1EsXgbp4W

RyanPCo force-pushed the ryanco/image-history branch from 3688522 to 5a2507c Compare September 20, 2026 19:09

RyanPCo force-pushed the ryanco/pi-scratch-init branch from 211db25 to 6efcdae Compare September 20, 2026 19:09

RyanPCo force-pushed the ryanco/image-history branch from 5a2507c to db5687f Compare September 21, 2026 18:45

RyanPCo force-pushed the ryanco/pi-scratch-init branch from 2fb9f92 to 663b77b Compare September 22, 2026 13:18

RyanPCo force-pushed the ryanco/image-history branch from db5687f to 5eeadea Compare September 22, 2026 13:18

RyanPCo and others added 8 commits September 29, 2026 21:04

aidang3019 force-pushed the ryanco/image-history branch from 5eeadea to 89c4a6a Compare September 30, 2026 01:07

aidang3019 mentioned this pull request Sep 30, 2026

ponytail(u1): drop bench/ and duplicate viz scripts, reuse gpu_orphans in the reaper #703

Draft

RyanPCo force-pushed the ryanco/pi-scratch-init branch from 02b7df5 to c986888 Compare October 12, 2026 20:26

RyanPCo closed this Oct 12, 2026

Navigation Menu

Repository navigation

loader: image history (a past front frame, gap in seconds) and the episode fps in every sample - #678

loader: image history (a past front frame, gap in seconds) and the episode fps in every sample#678
RyanPCo wants to merge 34 commits into
ryanco/pi-scratch-initfrom
ryanco/image-history

RyanPCo commented Sep 20, 2026 •

edited

Loading

Uh oh!

RyanPCo commented Sep 20, 2026 •

edited

Loading

Uh oh!

RyanPCo commented Oct 12, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

RyanPCo commented Sep 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Behaviour changes

Usage

Read cost

Tests

Uh oh!

RyanPCo commented Sep 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

RyanPCo commented Oct 12, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

RyanPCo commented Sep 20, 2026 •

edited

Loading

RyanPCo commented Sep 20, 2026 •

edited

Loading