Skip to content

algo: DeterministicEvalMixin, the seeded eval helpers out of hpt.py - #665

Closed
RyanPCo wants to merge 20 commits into
ryanco/hpt-devicefrom
ryanco/deterministic-eval-mixin
Closed

RyanPCo wants to merge 20 commits into
ryanco/hpt-devicefrom
ryanco/deterministic-eval-mixin

Conversation

@RyanPCo

@RyanPCo RyanPCo commented Sep 19, 2026 •

Copy link
Copy Markdown
Contributor

The eval-time seeding (EVAL_BASE_SEED + rank * EVAL_RANK_STRIDE + pass
counter, the prompt sampling and the seeded generator) was 96 lines inside
HPT, though nothing in it is HPT-specific: it reads only nets, device,
annotation_key, annotation_sampling_mode and default_prompt. Moving it to
egomimic/algo/eval_determinism.py as a mixin leaves hpt.py with the algo and
gives the seeding contract one documented home.

Pure code motion -- no behaviour change, and the existing determinism tests
cover it unchanged.

Co-Authored-By: Claude Fable 5.1 [email protected]
Co-Authored-By: Claude Opus 5 (1M context) [email protected]
Claude-Session: https://claude.ai/code/session_01THTuWdgk6M6pZ5Ya7ihZEF

🤖 Generated with Claude Code

https://claude.ai/code/session_01THTuWdgk6M6pZ5Ya7ihZEF

RyanPCo commented Sep 19, 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.

@RyanPCo
RyanPCo force-pushed the ryanco/deterministic-eval-mixin branch from 472e022 to ad1ccef Compare September 20, 2026 19:09
@RyanPCo
RyanPCo force-pushed the ryanco/hpt-device branch 2 times, most recently from 7848422 to b7a6793 Compare September 21, 2026 18:45
@RyanPCo
RyanPCo force-pushed the ryanco/deterministic-eval-mixin branch 2 times, most recently from 2dee1f0 to 249e508 Compare September 22, 2026 13:18
RyanPCo and others added 8 commits September 29, 2026 21:04
torchvision draws one random parameter set per call, so applying the
configured train_image_augs to the stacked (B, 3, H, W) camera batch gave
every image in the batch the same jitter. Wrap train_image_augs in a new
PerSampleAugs module so each image draws its own parameters, and move the
inline aug block of _robomimic_to_hpt_data into _apply_image_augs with the
same gating. Deterministic eval augs stay batched.

_apply_image_augs also checks train_image_augs before calling it: __init__
accepts None and train mode would have raised TypeError.

The loop is not free: measured on an A40 at batch 64 with the shipped
ColorJitter + Normalize, 13.8 -> 92.2 ms at 224x224 and 62.2 -> 105.5 ms at
the camera's 360x640, per camera per training step. Correctness first here;
the cost is removed further up the stack, where the jitter itself is
vectorised over the batch.

Co-Authored-By: Claude Fable 5.1 <[email protected]>
Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
… LR group

Two defects that would confound any backbone comparison:

* a frozen `ResNet` backbone kept its BatchNorm layers in train mode, so it
  normalized with batch statistics and kept updating its running stats.
  `requires_grad = False` does not touch either. `train()` now puts a frozen
  backbone in eval mode. Latent today -- no shipped config freezes the ResNet
  -- so no run to date was affected.
* the pretrained backbone and the randomly initialised trunk shared one LR.
  The new `model.backbone_lr_scale` key gives the encoders' backbones
  (`backbone_parameters()`) their own param group at `lr * scale`; at the
  default 1.0 the optimizer is built exactly as before, as a single group over
  `trainer.model.parameters()`. Nothing shipped sets it to anything else, so
  the two-group path is exercised only by its tests for now.

Co-Authored-By: Claude Fable 5.1 <[email protected]>
Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
ZarrDataset emits a real action_pad_mask (1.0 per real frame, 0.0 per
repeat-last pad, honouring the annotation-cutoff EOS clamp) next to every
horizoned action chunk. A new InterpolatePadMask transform carries it
through the 30 -> 100 upsampling on the interpolators' own grid: an output
step is real only if both source frames it interpolates between are real.

HPT gains a use_pad_mask algo key (default false, declared in the two flagship
model yamls; every other HPT config takes the __init__ default). With it on,
process_batch_for_training builds pad_mask (B, S, 1) from the dataset mask
instead of ones; DenoisingPolicy.loss_fn takes the mask and means the squared
error over the valid steps only. With it off nothing changes: the mask is
dropped from the batch and pad_mask stays all ones, which is exactly
F.mse_loss.

It ships off to match openpi, which asks LeRobotDataset for action_is_pad and
then ignores it (Pi0.compute_loss means over the whole horizon). LeRobot's own
policies do mask, but reduce with .mean() over all elements including the
zeroed ones, so their loss shrinks as padding grows; this divides by the valid
count instead, which keeps the gradient magnitude constant. Measured on the
three pinned flagship episodes, 6.4-7.6 % of supervised action steps are
repeat-last padding and 12.8-15.2 % of samples contain some, so turning it on
changes the loss: a run before and a run after are not comparable, and the
layout-aware eval metrics would have to be masked in the same change or train
and val would be measuring different things.

Padding is still repeat-last and no tail samples are dropped.

Co-Authored-By: Claude Fable 5.1 <[email protected]>
Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
…mbedding)

The proprio keys the model consumes read K steps instead of one. Every change
is gated on K > 1 (data) or history_len > 1 (model), and K = 1 / stride 1 is
bit-for-bit today's pipeline.

Data: Human.get_keymap gains proprio_history and history_stride; for K > 1 it
marks the proprio entries with "history": K plus the stride, and
ZarrDataset.__getitem__ reads the window idx - s*(K-1) ... idx in steps of s,
clamped at the episode start and front-padded by repeating the first real
frame, so the key is always (K, D) with the current frame last. A
proprio_history_mask (1 = real, 0 = front padding) travels with it.

The stride is why this ships on rather than at K = 1. Measured on a real
flagship episode (30 fps, fold_clothes), the median wrist translation is
6.8 mm over 1 frame, 25.7 mm over 3, 42.9 mm over 5 and 80.8 mm over 10: an
unstrided K = 3 window is 67 ms of past, close enough to a duplicate of the
current pose to be mostly sensor noise, while the action chunk it conditions
spans a full second. The flagship recipe therefore ships proprio_history: 3 /
history_stride: 5 -- 0.33 s of past -- with the stem at history_len: 3. That
changes the observation, so runs before and after this commit are not
comparable, and history_dropout: 0.2 stops being inert: it is now a real
hyperparameter, and the mitigation for the copycat/causal-confusion risk that
proprio history classically brings.

Transforms: the cartesian builders already treat a leading axis as batch;
ActionChunkCoordinateFrameTransform / PoseCoordinateFrameTransform gain an
opt-in target_history so a proprio target frame collapses to the CURRENT step
(the action chunk must not move when history is added), and Reshape /
QuaternionPoseToYPR tolerate the extra axis. observations.state.ee_pose is
(K, 20).

Eval: every human revert transform list starts with a SelectCurrentStep on its
proprio key, so the val metrics, the cam-frame revert and the overlay video see
the single current pose they always saw. No-op at K = 1 and idempotent.

Norm stats are the baseline's: per-channel (D,) stats broadcast over (K, D),
and both the training and the precompute stats pass force proprio_history=1 so
K > 1 and K = 1 runs share one set.

Model: MLPPolicyStem gains history_len / history_dropout. For history_len > 1
it owns a zero-init learned time embedding added after the projection, and
HPT.stem_process skips the input-side sinusoid for such a stem (a constant
[0, 1, 0, 1, ...] on normalized values is amplitude-1 noise on [-1, 1] data).
history_dropout replaces the whole past by the current step per sample in
training, exactly what an episode-start sample looks like after front padding.

The K = 1 baseline stays one override away (proprio_history=1,
history_stride=1, history_len=1) and the tests pin that it is bit-identical.
Every other shipped config is untouched at K = 1, including the keypoint arm,
whose wrist-frame proprio would need per-step framing before it can carry
history.

A compose test pins every train_*.yaml recipe's data K against its stem's
history_len, because a recipe that inherits one model's data config and
overrides the model is exactly how the two drift apart, and nothing else
composes both halves.

_stem_history_len takes the domain rather than matching a stem name by suffix:
a cotrain run registers one state_ee_pose stem per domain, and a suffix match
returns whichever comes first, so the human side was checked against eva's
history_len. Latent while everything shipped K = 1; a blocker now.

Co-Authored-By: Claude Fable 5.1 <[email protected]>
Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
… padding

`HPTModel.finalize_modules` ran `self.apply(self._init_weights)` over the whole
model, which xavier-randomised the pretrained HF encoders inside the text stems
(the frozen Qwen3-Embedding weights were overwritten in place and their biases
zeroed). Modules that own pretrained weights now declare them through the
`PretrainedWeights` mixin, and the init pass walks the tree without ever
entering such a subtree, so the ResNet backbone, `T5Encoder` and any future
pretrained stem are covered by the same mechanism.

`HPT.__init__` also called `finalize_modules` twice; the first call is gone.

Once the model is built, moved and upcast, `verify_pretrained_weights` hashes
every declared pretrained submodule and compares it with the checkpoint read
straight off the local HF snapshot, failing the run on a mismatch. It costs a
second full copy of the encoder (~1.1 GB fp16 at 0.6B) and ~3 s of hashing per
rank, so `verify_pretrained: false` turns it off; it is on by default, because
a silently re-initialised encoder trains to a plausible wrong result. Only
stems that own an HF snapshot are checked -- `ResNet` returns no reference, so
it is protected from the init pass but not hash-verified, which the mixin
docstring now says.

Finally, `QwenPerTokenEncoder.compute_latent` now passes the tokenizer
attention mask into its cross-attention. Padded positions are zeroed before the
biased `proj`, so they were the learned constant `proj.bias`, not zero, and got
real attention weight: text-conditioned results from before this commit are not
comparable with results after it.

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
The annotation vocabulary is tiny next to the number of training steps, so a
frozen text encoder re-runs the same 600M-parameter forward every step.
`_Qwen3BaseEncoder` now keeps a lazy in-memory cache of the per-token hidden
states keyed by the exact prompt string: misses go through the tokenizer +
encoder in ONE sub-batch of unique strings, their un-padded (L_i, D) rows are
stored detached, and the batch is re-assembled with the same LEFT padding the
tokenizer produces, so `forward_with_mask` / the pooled path are unchanged.

Only legal when the encoder is frozen - a trainable encoder must stay online.
Knobs on the stem: `cache_text_features` (default true) and
`text_cache_max_entries` (default 8192, LRU); both added to the model yamls
with a Qwen stem. The cap counts entries, and an entry is the prompt's
un-padded (L, D) states -- 4-25 KB for real annotations, 262 KB at
max_length 128 -- so the default is ~100-200 MB of device memory in practice
and 2 GB in the worst case; the docstring says so next to the knob. A batch
carrying more unique prompts than the cap reads its rows from the miss pass
rather than from the cache, because the LRU trim runs first and would
otherwise evict entries the very next line looks up. `_apply` (to/cuda/float) clears the cache, and
`clear_text_feature_cache()` / `text_feature_cache_stats()` are exposed.

Padded positions come back as zeros instead of the encoder's pad-token output;
every consumer masks them out. Cached-vs-online on 50 real prompts (fp16, each
warmed alone: worst-case padding mismatch): max |diff| 3.6e-3 on the latents
and 9.4e-2 on raw hidden states, exactly equal to the online path's own
batch-padding noise (9.4e-2); in fp32 the same comparison is 1.1e-6 / 1.6e-4.

Co-Authored-By: Claude Fable 5.1 <[email protected]>
Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
HPT.__init__ defaulted self.device to torch.device('cuda') and moved the
freshly built policy there. Under submitit every rank sees all 8 GPUs, so all
8 ranks resolved that to cuda:0 and built their policy on the same device:
8 x 3.34B, which OOMs during construction. The 1B arm fit and survived, which
made the failure look config-specific.

PI already builds on CPU and leaves self.device unset for Lightning to fill in
through ModelWrapper.on_fit_start / on_validation_start. HPT now does the same:
no device is chosen in __init__ and nothing leaves CPU there.

Lightning's write then has to reach the HPTModel the algo holds. That model
reads self.device to build the loss accumulators in compute_loss, while
HPT.compute_losses accumulates on the algo's, so a write that stopped at the
algo would leave the two on different devices -- 'Expected all tensors to be
on the same device, but found at least two devices, cuda:0 and cpu!'.
HPT.device is therefore a property whose setter also writes
nets['policy'].device, so the two cannot diverge through the path the trainer
uses.

Verified on 8xH100 with the 3.34B hpt_3b arch-comparison arm: every rank
reports CUDA_VISIBLE_DEVICES [0..7], all 8 reach fit start and train an epoch
without an OOM.

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01EAGuZKPVfr2J9L7W9gtxET
The eval-time seeding (eval_base_seed + rank * EVAL_RANK_STRIDE + pass
counter, the prompt sampling and the seeded generator) was 96 lines inside
HPT, though nothing in it is HPT-specific: it reads only nets, device,
eval_base_seed, annotation_key, annotation_sampling_mode and default_prompt.
Moving it to egomimic/algo/eval_determinism.py as a mixin leaves hpt.py with
the algo and gives the seeding contract one documented home.

The attributes it reads are annotated, not defaulted: a class-level default
would turn "the algo forgot to set this" from an AttributeError into a silent
wrong value. The module docstring says plainly that this is HPT-only today --
PI mixes in nothing and has no eval seeding at all, so pi0.5 validation is
still not reproducible -- and what wiring PI up would take.

Pure code motion -- no behaviour change, the eight method bodies are
byte-identical to the originals, and the existing determinism tests cover it
unchanged.

Co-Authored-By: Claude Fable 5.1 <[email protected]>
Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
@RyanPCo

RyanPCo commented Oct 12, 2026

Copy link
Copy Markdown
Contributor Author

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

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