Skip to content

eval: layout-aware frame-invariant metrics, and a train-split viz head - #659

Draft
RyanPCo wants to merge 1 commit into
ryanco/dataset-robustnessfrom
ryanco/eval-metrics-train-viz
Draft

RyanPCo wants to merge 1 commit into
ryanco/dataset-robustnessfrom
ryanco/eval-metrics-train-viz

Conversation

@RyanPCo

@RyanPCo RyanPCo commented Sep 17, 2026 •

Copy link
Copy Markdown
Contributor

Branch 7 of 8 in the split of #640, which is the tip of the stack and carries the map of all eight.
Unit suite on this branch: 634 passed, 0 failed.

Metrics were paired MSE over whatever the action vector happened to be, which
mixes metres and rotation-matrix entries into one number that means nothing in
either unit. action_metrics dispatches on the layout instead: cartesian widths
report xyz MSE (m^2), geodesic rotation error in degrees (after Gram-Schmidt,
so it survives a non-orthonormal prediction), raw rot6d MSE as a
training-tracking diagnostic, DTW and Frechet; keypoint widths report keypoint
L2 in metres, wrist xyz / rotation error and keypoint DTW. Each also splits
early / mid / late (0-10 / 10-50 / 50-100 % of the chunk), because a chunk
model can be right at t=0 and drift, and one average hides that. DTW forgives
temporal misalignment that paired MSE punishes: a correct trajectory executed
slightly early scores near zero.

These are all frame-invariant -- the cam-frame revert is a rigid transform
applied to prediction and ground truth alike -- so the revert no longer runs
for metrics at all, only when rendering the overlay. That also removes the
cam-frame MSE, which was the same number as the native one.

Rendering is CPU-bound at roughly a second a frame and dominated viz-epoch wall
time (unbounded rendering added ~2h per viz epoch on the combined pretrain), so
viz_every_n_epochs gates it on Lightning's (current_epoch + 1) convention, the
one validation actually runs on, with the final epoch always rendering so short
runs get a video; viz_max_batches caps frames per rank. Metrics still run on
every batch of every pass. Video fps is divided by world_size, since
distributed validation deals indices round-robin and rank 0's consecutive
frames are world_size source frames apart -- they were being written as a
world_size-times timelapse.

train_viz adds a second val loader over the train split (metrics train_viz/,
videos videos_train_viz/), on by default for training runs with an evaluator,
so train and val curves come from the same evaluator on the same axes. A third
unseen_op_valid loader and data.valid_prefix let the held-out-operator splits
log seen_op_valid/ and unseen_op_valid/ separately, instead of the seen
operators landing in the train_viz slot and the unseen ones in Valid/.
ModelWrapper dispatches on val_loader_names by dataloader_idx. val_at_start
runs a pre-fit baseline validation, skipped on resumes because trainer.validate
there would score the un-resumed base model.

Reverse KL is removed: M forward passes per val batch for a number nothing was
read off.

Co-Authored-By: Claude Opus 5 (1M context) [email protected]

🤖 Generated with Claude Code

https://claude.ai/code/session_01HjCQBqQCr2igFEMR7yUovw

RyanPCo commented Sep 17, 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/dataset-robustness branch from ecb6aae to 73b5244 Compare September 30, 2026 01:07
@RyanPCo
RyanPCo force-pushed the ryanco/dataset-robustness branch from 73b5244 to d5599b8 Compare October 12, 2026 21:50
@RyanPCo
RyanPCo force-pushed the ryanco/eval-metrics-train-viz branch from 8c3b7f1 to 38aaa92 Compare October 12, 2026 21:50
@RyanPCo
RyanPCo force-pushed the ryanco/eval-metrics-train-viz branch from 38aaa92 to e7eebed Compare October 12, 2026 23:19
@RyanPCo RyanPCo mentioned this pull request Oct 19, 2026
@RyanPCo
RyanPCo force-pushed the ryanco/dataset-robustness branch from 64bcf03 to 32b73de Compare October 19, 2026 19:37
@RyanPCo
RyanPCo force-pushed the ryanco/eval-metrics-train-viz branch from e7eebed to ce8666e Compare October 19, 2026 19:37
@RyanPCo
RyanPCo force-pushed the ryanco/dataset-robustness branch from 32b73de to 6fd89d1 Compare October 30, 2026 18:27
@RyanPCo
RyanPCo force-pushed the ryanco/eval-metrics-train-viz branch from ce8666e to 2ca1d41 Compare October 30, 2026 18:27

@RyanPCo RyanPCo left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Automated adversarial review (read-only; findings were checked by a second skeptical pass before posting): 1 high, 2 medium, 8 low, 9 nit.

return loaders[0]
# Several heads: return a list so Lightning populates dataloader_idx
# (the position in val_loader_names()) and ModelWrapper can dispatch.
return loaders

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high (correctness): Returning a list of CombinedLoaders breaks DDP sharding. Lightning only injects a DistributedSampler into the leaves of combined_loader.flattened, and here those leaves are the inner CombinedLoaders. _prepare_dataloader returns early for anything that isn't a DataLoader, so the inner loaders are never sharded. Because train_viz: true is the default, every multi-GPU run now has each rank evaluating the same first limit_val_batches batches of valid and train_viz. That cuts distinct-sample coverage by world_size and makes the _video_fps stride compensation wrong. Fix: return one flat sequential CombinedLoader over plain DataLoaders keyed so head and dataset can be recovered, or attach an explicit DistributedSampler(shuffle=False) when distributed is initialized.

Comment thread egomimic/eval/eval_hpt.py
)
rkl = reverse_kl_from_samples(samples, gt_tensor)
metrics[f"Valid/{pred_key_name}_reverse_kl_M{M}"] = rkl.item()
if not do_viz:

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium (correctness): Because of if not do_viz: continue, the cam-frame MSE keys (Valid/<key>_cam_paired_mse_avg / _cam_final_mse_avg) are now logged only on viz epochs and only for the first viz_max_batches batches. Every other metric averages over all val batches, so this curve is sparse and can't be compared with them. PIEvalVideo removed cam-frame MSE because the layout metrics are frame-invariant. Suggest doing the same here: keep the revert only for preds_for_viz, then update the class docstring and the eval_hpt_keypoints_wrist / eval_hpt_wristframe_6d yaml comments that still list these keys.

Comment thread egomimic/eval/action_metrics.py Outdated

def layout_metrics(pred: torch.Tensor, gt: torch.Tensor, prefix: str) -> dict:
"""Dispatch on the last-dim width: cartesian, keypoint, or ``{}``."""
metrics = cartesian_metrics(pred, gt, prefix)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium (correctness): layout_metrics picks the layout from width alone. The cotrain HPT configs use eva_bimanual: actions_joints with 14-D actions (7 joints/arm), and 14 matches the robot ypr cartesian layout. Those runs therefore log xyz_*, xyz_dtw_avg and rot_err_deg_* computed on joint radians as if they were metres and Euler angles. Pass ac_key in and dispatch only when it is in pose_utils._CARTESIAN_KEYS / _KEYPOINT_KEYS, as rot6d_channels does. In eval_pi, keep a generic paired/final MSE fallback for when layout_metrics returns {}; right now a PI action at an unrecognized width logs no action metric at all.

Comment thread egomimic/eval/eval_train_viz.py Outdated

class TrainVizEvalVideo(EvalVideo):
def __init__(
self, base: EvalVideo, limit_val_batches: int = 50, prefix: str = "train_viz"

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

low (dead-code): limit_val_batches on TrainVizEvalVideo is never read. It only goes into override_dict, which trainHydra applies only in eval mode and only from the canonical cfg.evaluator, while this wrapper is built only in train mode. Drop the parameter, and remove limit_val_batches: 50 and its "eval-only mode" comment from the train_viz_*.yaml files.

Comment thread egomimic/eval/eval_train_viz.py Outdated
from egomimic.eval.eval_video import EvalVideo


class TrainVizEvalVideo(EvalVideo):

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

low (naming): This wrapper now also prefixes the canonical valid head (data.valid_prefix) and the unseen_op_valid head, so TrainVizEvalVideo wrapping validation evaluators in trainHydra reads as a mistake. Consider renaming it to something like PrefixedEvalVideo, generalizing the module docstring, and moving the docstring from __init__ to the class.

Comment thread egomimic/eval/eval_video.py Outdated

def on_validation_start(self):
if self.trainer.is_global_zero:
if self.trainer.is_global_zero and self._should_viz():

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit (dead-code): _write_video now creates epoch_<n>/<embodiment>/ itself, so this on_validation_start makedirs only leaves an empty epoch_<n> directory. The override can be removed, since the base Eval.on_validation_start is a no-op.

# `+evaluator@train_viz_evaluator=train_viz_pi_camframe_6d`
# Needs the human transform mode overridden to `cartesian_6d`.

#

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit (duplication): The blank line followed by a lone # splits this comment block, and the same leftover appears in train_viz_pi_wristframe_6d and train_viz_hpt_keypoints_wrist. Join the block. Consider keeping the shared "base is composed separately" explanation in one place, such as the eval_train_viz.py docstring.

Comment thread tests/unit/test_pi05_norm_rot6d.py Outdated
assert (fd_shift > fd).all(), "mean shift must increase the distance"


def test_reverse_kl_is_gone():

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit (test-quality): test_reverse_kl_is_gone only checks that a deleted function stays deleted and tests no behavior. Drop it.

Comment thread egomimic/eval/eval_train_viz.py Outdated
self.base = base
self.prefix = prefix
super().__init__(
limit_val_batches=limit_val_batches,

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit (dead-code): limit_val_batches on TrainVizEvalVideo only feeds override_dict, which trainHydra applies in eval mode. The train_viz head is only built when mode == 'train', so the value never applies. The train_viz_*.yaml comment ("Cap on batches ... in eval-only mode") describes behaviour that can't happen. Remove the parameter and the yaml keys, or fix the comment.

Comment thread egomimic/eval/action_metrics.py Outdated
xyz_idx = list(layout["xyz"])
xyz_p, rot_p = _split_mse(pred_cpu, gt_cpu)
metrics[f"{prefix}_xyz_paired_mse_avg"] = xyz_p
metrics[f"{prefix}_rot6d_paired_mse_avg"] = rot_p

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit (naming): For 12/14-D (ypr) widths this logs the yaw/pitch/roll channel MSE under <prefix>_rot6d_paired_mse_avg. Use a format-neutral name such as rot_paired_mse_avg, or pick rot6d/ypr from the layout width, before dashboards start keying on it.

Metrics were paired MSE over whatever the action vector happened to be, which
mixes metres and rotation-matrix entries into one number that means nothing in
either unit. action_metrics dispatches on the layout instead: cartesian widths
report xyz MSE (m^2), geodesic rotation error in degrees (after Gram-Schmidt,
so it survives a non-orthonormal prediction), raw rot6d MSE as a
training-tracking diagnostic, DTW and Frechet; keypoint widths report keypoint
L2 in metres, wrist xyz / rotation error and keypoint DTW. Each also splits
early / mid / late (0-10 / 10-50 / 50-100 % of the chunk), because a chunk
model can be right at t=0 and drift, and one average hides that. DTW forgives
temporal misalignment that paired MSE punishes: a correct trajectory executed
slightly early scores near zero.

These are all frame-invariant -- the cam-frame revert is a rigid transform
applied to prediction and ground truth alike -- so the revert no longer runs
for metrics at all, only when rendering the overlay. That also removes the
cam-frame MSE, which was the same number as the native one.

Rendering is CPU-bound at roughly a second a frame and dominated viz-epoch wall
time (unbounded rendering added ~2h per viz epoch on the combined pretrain), so
viz_every_n_epochs gates it on Lightning's (current_epoch + 1) convention, the
one validation actually runs on, with the final epoch always rendering so short
runs get a video; viz_max_batches caps frames per rank. Metrics still run on
every batch of every pass. Video fps is divided by world_size, since
distributed validation deals indices round-robin and rank 0's consecutive
frames are world_size source frames apart -- they were being written as a
world_size-times timelapse.

train_viz adds a second val loader over the train split (metrics train_viz/,
videos videos_train_viz/), on by default for training runs with an evaluator,
so train and val curves come from the same evaluator on the same axes. A third
unseen_op_valid loader and data.valid_prefix let the held-out-operator splits
log seen_op_valid/ and unseen_op_valid/ separately, instead of the seen
operators landing in the train_viz slot and the unseen ones in Valid/.
ModelWrapper dispatches on val_loader_names by dataloader_idx. val_at_start
runs a pre-fit baseline validation, skipped on resumes because trainer.validate
there would score the un-resumed base model.

Reverse KL is removed: M forward passes per val batch for a number nothing was
read off.

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
@RyanPCo
RyanPCo force-pushed the ryanco/dataset-robustness branch from 6fd89d1 to 7a13410 Compare October 30, 2026 23:05
@RyanPCo
RyanPCo force-pushed the ryanco/eval-metrics-train-viz branch from 2ca1d41 to 05a92e1 Compare October 30, 2026 23:05
@RyanPCo

RyanPCo commented Oct 30, 2026

Copy link
Copy Markdown
Contributor Author

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

Fixed

  • HPT cam-frame MSE is now computed only on viz batches: Removed the HPT cam-frame MSE, matching PI. Updated the docstring and the yaml comments, including main's eval_hpt_wrist.yaml, which would otherwise go stale. obs: proprio history (K proprio steps, one token each, learned time embedding) #650's proprio-history evaluator test no longer expects cam-frame keys.
  • Re-export with noqa 'for tests' is unused: Now imports only _paired_mse (used by the new fallback) and layout_metrics.
  • TrainVizEvalVideo.limit_val_batches and every limit_val_batches: 50 in train_viz_*.yaml are dead: Removed the parameter and every yaml entry.
  • TrainVizEvalVideo is now a generic prefixing wrapper but is still named and documented as train-viz only: Renamed to PrefixedEvalVideo, with TrainVizEvalVideo kept as an alias. Docstring is on the class; module docstring is neutral.
  • Three copies of wrap + trainer/model attach and three function-local imports: Import moved to module level; one _attach helper used for all three heads.
  • getattr(datamodule, ..., {}) is unnecessary defensiveness: Reads datamodule.train_viz_datasets / unseen_op_valid_datasets directly.
  • getattr fallbacks and the hard-coded ['valid','train_viz'] default exist only for tests that skip init: val_loader_names defaults to ['valid'] and both attributes are read directly. The routing test sets them explicitly.
  • valid_prefix and held_out_operators are stored on the datamodule but never read: Still accepted but no longer stored; docstring reduced to one line.
  • rot6d_paired_mse_avg is logged for ypr widths too: Logged as ypr_paired_mse_avg for the 12/14-D widths; still rot6d_paired_mse_avg for 18/20-D.
  • on_validation_start directory creation is redundant now that _write_video makedirs: Reduced to a no-op (Eval declares it abstract, so it can't be deleted); guard simplified. Touches main's makedirs lines.
  • eval_pi_wristframe.yaml is a byte-for-byte copy of eval_pi.yaml's settings: Now just defaults: [eval_pi, self], kept as a name alias.
  • Five near-identical train_viz_*.yaml files with broken comment blocks: Each file is now defaults plus a short pointer; the shared notes moved to the eval_train_viz module docstring; stray '#' lines gone. All five files kept.
  • test_reverse_kl_is_gone asserts that deleted code stays deleted: Deleted.
  • Eval-metric/viz tests placed in an unrelated rot6d test module: Moved to new test_action_metrics.py (including the segment_l2 tests) and test_eval_video.py, with one module-level _StubEvalVideo. test_pi05_norm_rot6d.py is back to its pre-eval: layout-aware frame-invariant metrics, and a train-split viz head #659 layout and docstring. Later PRs' edits to the moved tests (infra: mecka 6D / keypoint recipes, 300M models, defaults, and the rollout revert #640, eval: overlay videos play at the source rate under DDP #666) follow them into the new files.
  • layout_metrics chooses the layout from width alone: layout_metrics takes ac_key and dispatches only for the cartesian/keypoint keys. PIEvalVideo falls back to plain paired/final MSE for other keys. Tests added.
  • TrainVizEvalVideo.limit_val_batches / train_viz_*.yaml limit_val_batches: 50 never take effect: Same fix as the dead limit_val_batches finding.
  • rot6d_paired_mse_avg is logged for ypr widths too: Same fix as the other ypr-naming finding.

Partial

  • Default HPT evaluator uses interval 200 and no viz_max_batches: Added viz_max_batches: 10 to eval_hpt and eval_hpt_cartesian. Left the 200-epoch interval as is, since the reason for it isn't known.

Not fixed

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