Skip to content

pi: keep per-sample intrinsics so overlays use the episode camera - #690

Draft
RyanPCo wants to merge 55 commits into
ryanco/rdt-memoryfrom
ryanco/pi-overlay-intrinsics
Draft

RyanPCo wants to merge 55 commits into
ryanco/rdt-memoryfrom
ryanco/pi-overlay-intrinsics

Conversation

@RyanPCo

@RyanPCo RyanPCo commented Sep 24, 2026 •

Copy link
Copy Markdown
Contributor

PI overlay videos on non-Aria data drew GT and predictions with the wrong camera.

Problem. PI.process_batch_for_training keeps only keymap keys plus episode_hash. intrinsics is not a keymap key, so it was dropped, viz_gt_preds got no per-sample camera, and Human.viz fell back to ARIA_INTRINSICS. On Mecka that puts every overlay point about 58 px too low (cy 240 vs 182) and pushes points outward (fx about 7% larger). HPT and RDT keep unknown keys (zarr_key_to_keyname(...) or key), so their overlays were already right.

Verified by rebuilding one episode's GT: with the episode's intrinsics it matches the RDT video; with the Aria intrinsics it matches the pi05 video.

Fix. Carry intrinsics forward alongside episode_hash.

Scope. Only the videos were wrong. Neither training nor validation metrics read intrinsics: actions, state and loss are all 3D.

Test. test_pi_process_batch_keeps_the_per_sample_intrinsics.

🤖 Generated with Claude Code

https://claude.ai/code/session_01NxUsbpS65QBRpBiauTEtTs


Folded in (to keep the stack under Graphite's 50-branch cap):

RyanPCo commented Sep 24, 2026 •

Copy link
Copy Markdown
Contributor Author

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.

@aidang3019
aidang3019 force-pushed the ryanco/pi-overlay-intrinsics branch from a5a8948 to 70140ee Compare September 30, 2026 01:07
@RyanPCo
RyanPCo force-pushed the ryanco/pi-overlay-intrinsics branch from 70140ee to ecf46e5 Compare October 12, 2026 20:26
@RyanPCo
RyanPCo force-pushed the ryanco/pi-overlay-intrinsics branch from ecf46e5 to f600e0d Compare October 12, 2026 21:50
@RyanPCo
RyanPCo force-pushed the ryanco/pi-overlay-intrinsics branch from f600e0d to 651e44b Compare October 12, 2026 23:19
@RyanPCo RyanPCo mentioned this pull request Oct 19, 2026
RyanPCo and others added 3 commits October 19, 2026 19:36
…an default

Human data trains on hand keypoints instead of a palm-origin cartesian pose.
The action is 144-D wrist-first, per hand
[wrist xyz (3) | wrist rot6d (6) | 21 MANO keypoints (63)] x {left, right},
built by the keypoints_{wrist,head}frame_6d modes: the existing ypr keypoint
builders produce the 138-D layout and KeypointsYPRToRot6D re-encodes just the
wrist rotation, leaving the keypoint blocks untouched. Both the action and the
observations.state.keypoints proprio are converted, so neither carries a
wrapping euler angle. MANO keypoint 0 IS the wrist, so per hand the wrist
position rides at both ch 0-2 and ch 9-11; that is kept deliberately (21
keypoints is the MANO interface every consumer reshapes to) and written down
in the bimanual_keypoint_layout docstring.

include_ee_pose additionally builds the cartesian head-frame ee_pose next to
the keypoint action (keymap + transform list), so a keypoint pi0.5 run
discretizes the same State: prompt block as the cartesian runs and differs
from them only in the action space. The cartesian builder runs with
delete_target_world=False to keep obs_head_pose alive for the keypoint
builder, and its actions_cartesian by-product is dropped.

144 does not fit openpi's action vector, whose PyTorch port hard-codes 32-wide
action projections regardless of Pi0Config.action_dim. PI now reads action_dim
and, when it is not 32, swaps in action_dim-wide action_in_proj /
action_out_proj AFTER the strict base-weight load: only those two layers start
fresh, everything else comes from the checkpoint, and sample_actions already
draws noise at config.action_dim. pad_to_width zero-pads narrower embodiments
(eva's 20-D) into the same vector and raises rather than truncating a wider
one. _reduce_loss then restricts the mean to the slots the converter actually
packs, so eva does not average its loss over 112 zero-padded dims and
contribute a fraction of the gradient signal it had at width 32.
HumanBimanualKeypoints is identity-in, slice-out.

HPT human configs move to a 144-D state_keypoints stem. hpt_bc_flow_human
holds the body -- there is no vendor notion at the embodiment level, so the
canonical config is the embodiment one -- and hpt_bc_flow_{mecka,aria,scale}
are two-line aliases of it. HPTModel.add_stem drops null stem entries so a
child config can delete a stem its base declares, which hydra merges cannot
express, and hpt_cotrain_flow_seperate_head uses it to drop state_ee_pose.

With the keypoints_*_6d modes in place, the guard from the previous commit
widens to the keypoints ypr and quat modes, so every remaining discontinuous
rotation target needs an explicit allow_legacy_rotation=True (the aria->MANO
and mecka keypoint viz tutorials, which read pre-6D data). The cartesian
configs that were still pinned to the ypr widths move with it:
hpt_bc_flow_human_cartesian and the cotrain base to human 20/18, the
shared-head cotrains to a 20-wide shared action.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
Wrist-relative actions made a latent division safe-looking: the t=0 cell of a
wrist-frame chunk is exactly the identity pose, so its quantile range is 0 and
the historical `(q99 - q1) + 1e-6` denominator scaled any off-convention value
by 1e6. A channel whose range is below NORM_MIN_RANGE (1e-4) is now treated as
constant -- it normalizes to 0 and unnormalizes to its centre rather than
dividing by epsilon. Diffusion Policy / robomimic ignore ranges under 1e-4 and
GR00T maps min == max to 0; this is the same rule, applied in the shared
_apply_norm_one / _apply_unnorm_one so the dataset and the tests agree.

norm_stats.pool_horizon (default false, so nothing changes unless asked)
computes one set of stats per channel over every timestep of the chunk, the way
openpi does, instead of one per (timestep, channel) cell. The pooled stats are
tiled back to the (T, D) shape, so normalize, the bounds check and checkpoints
see the same layout either way. It is part of the norm-stat cache key, so
flipping it cannot silently reuse stats computed the other way.

precompute_norm_stats.py computes and caches the stats for a recipe outside a
training run, keyed by the same file + content hash.

Which of the two is better is unsettled: the literature is split (LBM
per-timestep with a stated rationale, openpi / RDT / GR00T / Diffusion Policy /
UMI / lerobot / Octo pooled) and nothing here has compared them on a real
recipe, so the default stays the inherited per-timestep one and the config
comment says what would make it worth flipping.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
A 6D rotation is two columns of a rotation matrix: already bounded in [-1, 1],
and per-channel centre-and-scale pulls the pair off the manifold without buying
any resolution. The two projects that anchor an action chunk at near-identity
the way this repo does -- TRI's LBM and UMI -- both exclude rotation from
normalization for exactly that reason (LBM: "We avoid this normalization
procedure for the 6D rotation in the poses to avoid corrupting the rotation
matrix"; UMI sends its rotation channels through an identity normalizer). Per
timestep AND normalized was the one combination with no precedent, and it is
what manufactures the t=0 zero-range rotation cells NORM_MIN_RANGE exists to
catch -- this removes them at the source.

rot6d_channels() reads the indices off the existing layout tables, so the rule
follows the layout rather than a second hardcoded list: the cartesian 18/20
widths and the 144-D keypoint layout, action and proprio alike. The ypr widths
keep the old behaviour; an Euler angle is not a matrix column and those modes
are legacy-only now. _apply_norm_one / _apply_unnorm_one take the channels and
pass them through untouched, so the dataset path, the evaluator revert and the
rollout revert all follow from one place, and the degenerate-channel guard is
left to do what it is actually good at: the structural kp0 / kp9 position
zeros.

Norm-stat files stay valid -- what changes is how they are applied -- but a run
before this commit and a run after it are not comparable, and the loss is now
weighted differently between translation (unit scale) and rotation (raw).

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
RyanPCo and others added 16 commits October 19, 2026 15:44
…e-rank

trainer/ddp.yaml and ddp_pi.yaml set devices to gpus_per_node * nodes next to
num_nodes, but Lightning reads devices per node and multiplies by num_nodes
itself, so a 2-node run asked each node for twice the GPUs it has. The submitit
launchers made the same mistake in gres, which Slurm also reads per node. Both
now use gpus_per_node; single-node runs are unchanged.

trainHydra scales each val head's frames-per-episode K by devices * num_nodes,
but eval mode forces devices=1 only after the loaders are built. An eval run of
a multi-GPU recipe therefore kept W times the frames per episode that its single
rank's limit_val_batches window holds, and the window ran out before the last
episodes. _trainer_world_size now returns 1 in eval mode.

Both predate this stack (the ddp devices expression is from 2025-01).

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UkAZwERXcFKjueZQKfRF6J
Each mp4 EvalVideo writes is also logged to any WandbLogger under
<video dir>/<embodiment> (videos/, videos_train_viz/, videos_unseen_op_valid/),
captioned with the epoch, so runs can be reviewed without copying files
off the cluster.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01415567AQYskZLeSpSMsgwL
viz_mode=replay plays each action chunk twice over the frames it spans:
first the ground-truth hand at each step, then the predicted one. Each step
is drawn on the frame it belongs to, with the chunk mapped into that frame's
camera by a rigid fit of the GT keypoints (the same hand at the same instant
in both frames), so head motion needs no head pose. The 126-D keypoint
layout draws both hands' 21-keypoint skeletons; other layouts fall back to
the embodiment's own viz over a short trail.

Chunks don't overlap and span horizon * action_stride frames; frames of an
incomplete chunk carry over to the next flush. viz_max_batches must cover at
least one chunk (a 100-step horizon needs ~13 batches of 8).

_visualize_preds moves from HPTEvalVideo / PIEvalVideo (identical copies)
into EvalVideo.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01516MaELbiPNmJzkjPGebhu
Replay becomes EvalVideo's default. It needs consecutive frames, so it runs
only on a pinned video loader; a head without one keeps the overlay video,
and EvalLatent stays on overlay. The replay video loader is capped at
replay_chunks (5) chunks instead of viz_max_batches, which at 8 frames a
batch could not cover one 100-step chunk; viz_max_batches still bounds a
loader that yields no clips.

action_stride defaults to None and is inferred from the GT keypoints (step k
of a chunk rigidly matches frame k*stride's own step 0): data configs use 1
or 3 and cotrain mixes them under one embodiment. TrainVizEvalVideo passes
the replay settings through and tells its base which mode each step is.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01516MaELbiPNmJzkjPGebhu
2,866 of 41,607 local Mecka episodes carry one head pose for the whole
episode (no tracking) and 142 have single-frame head jumps of 0.3-3.7 m.
Both corrupt the head-frame action chunks: with a frozen pose, head motion
leaks into the hand trajectory. The vendor's egomotion is passed through
unchanged by mecka_to_zarr, so the fault is upstream.

- DatasetFilter(exclude_hashes=...): a list, a file, or a blocklist name
  under egomimic/resources/blocklists/. Excluded episodes never match; an
  excluded pin is dropped with a warning (an all-excluded pin list raises,
  since no pins means no filter).
- scan_egomotion.py classifies every local episode's obs_head_pose and
  writes the blocklist; mecka_bad_egomotion.txt holds its 3,008 hashes.
- mecka_all_6d applies it (this reshuffles its seeded train/valid split).
- mecka_20pct_6d: 393 of 4,471 train pins drop out. Its valid and
  unseen_op_valid video episodes were two of the frozen-pose episodes;
  both are repinned to packaging_food episodes with tracked egomotion
  (69b5691e... seen operator, 69b525ff... held-out operator).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01516MaELbiPNmJzkjPGebhu
Norm stats were cached per split, so a new episode set (20% -> full, an
operator hold-out, newly collected data) resampled every episode. Each
episode's sampled proprio/action rows are now cached on their own under
<cache_dir>/episodes/, keyed by the resolver recipe and reader code but not by
the split, and a split's stats are computed from the union of its episodes'
rows. Frames sit on a power-of-two stride with a per-episode phase, so a file
cached at stride s serves every coarser stride.

Adds norm_stats.max_samples (now part of the split key; KEY_VERSION 3).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0118RhoJmzcWC75B3vtc3zkF
…andas

Startup on the 41.6k-episode full-mecka split spent ~20 min before training:
populate_from_datasets read a full sample from every episode to learn the
post-transform keys, though a resolver hands every episode the same key_map
and transform objects. It now probes one episode per (embodiment, key_map,
transform) group, trying up to 8 before falling back to the raw key_map.

Resolving pinned filters was ~50 s each: a row-wise df.apply(row.to_dict())
over the ~458k-row episode table plus a per-pin .loc in _check_pins. Pinned
filters now narrow the table to their hashes first, rows are filtered as
records, and _check_pins builds one dict of pinned rows. Same paths on the real
table; 41.6k pins 2.5 s, val pins ~0 s, unpinned lab filter 32 s -> 13 s.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0118RhoJmzcWC75B3vtc3zkF
…ide in seconds

Human.get_keymap(image_memory=N, image_memory_stride_s=s) adds <front>_mem:
N front frames s seconds apart, the newest s seconds back, plus a (N,)
<front>_mem_mask of the steps inside the episode. history_stride_s spaces
proprio history in seconds, so windows span the same time at any fps.

History windows are now point reads instead of a contiguous span thinned
afterwards: on JPEG keys (one frame per chunk) a span would fetch every
frame between steps. plan_reads still merges them on numeric keys.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UL4UaiqSDNbhpatNwdvJbe
…-step eval MSE

_stem_history_len called .get on the policy's stems, an nn.ModuleDict once
built, so every proprio-history (K > 1) run crashed at its first batch. The
unit tests used a plain dict.

The last-step MSE handed torchmetrics a strided [:, -1] slice, which it
.view(-1)s; on GPU .cpu() copied it contiguous, on a CPU run it raised.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UL4UaiqSDNbhpatNwdvJbe
…lation arms

RDTMemory: N past frames of one camera through a frozen DINOv3 tower (only its
proj trains), a Perceiver resampler shared across frames pooling each to M
latents, a sin-cos code of each frame's age in seconds, and optionally that
step's proprio token (fuse_proprio). The N x M tokens join the DiT's
self-attention prefix [t; freq; state; memory; actions] under a key mask, so
steps before the episode start and dropped memories (trunk.memory.dropout)
are never attended. Without trunk.memory the model is unchanged.

Four train configs on mecka 20% 6D, a 2x2 over proprio history and visual
memory (10 steps, 1 s apart): _base, _proprio, _vision, _both.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UL4UaiqSDNbhpatNwdvJbe
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UL4UaiqSDNbhpatNwdvJbe
Follows the dual-encoder design of diffusion-policy-experiments:
- the memory tower is a second DINOv3, fine-tuned at the backbone LR scale
  (the param-group walk now covers every module with backbone_parameters),
  pooling each frame to its CLS token; the Perceiver resampler is gone
- the window ends at the current frame (t-9 s .. t)
- short_dropout (0.3 in the memory arms) replaces a sample's short image
  stream by a learned null token, replacing whole-memory dropout

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UL4UaiqSDNbhpatNwdvJbe
Backprop through 320 frames / GPU OOMed an H100; with the tower's blocks
recomputed the both arm trains at 1.19 s / step on 8xH100.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UL4UaiqSDNbhpatNwdvJbe
PI.process_batch_for_training kept only keymap keys plus episode_hash, so
batch['intrinsics'] never reached viz_gt_preds and Human.viz fell back to
ARIA_INTRINSICS. Every PI overlay video on non-Aria data was drawn with the
wrong camera (about 58 px low on Mecka). Metrics were unaffected.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NxUsbpS65QBRpBiauTEtTs
The keypoint pipeline resamples 30 raw frames to a 100-step chunk, so a
mecka step is 29/99 of a frame, but the replay drew step k on frame
a + k*stride with an integer stride (inferred as 1). GT and prediction
played ~3.4x slow and the head-motion fit paired different hand poses:
15-50 mm between the drawn GT and the frame's own label on the pinned
valid episode, 1.6 mm mean after this change.

render_replay takes a fractional stride and draws each real frame with
the step that falls on it; infer_stride searches D/(H-1) for whole
spans D. Integer strides render as before.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015Kf5s9frd6T6cbBvapfjFy
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant