Repository navigation
Conversation
ecb6aae to
73b5244
Compare
73b5244 to
d5599b8
Compare
8c3b7f1 to
38aaa92
Compare
38aaa92 to
e7eebed
Compare
64bcf03 to
32b73de
Compare
e7eebed to
ce8666e
Compare
32b73de to
6fd89d1
Compare
ce8666e to
2ca1d41
Compare
RyanPCo
left a comment
There was a problem hiding this comment.
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 |
There was a problem hiding this comment.
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.
| ) | ||
| rkl = reverse_kl_from_samples(samples, gt_tensor) | ||
| metrics[f"Valid/{pred_key_name}_reverse_kl_M{M}"] = rkl.item() | ||
| if not do_viz: |
There was a problem hiding this comment.
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.
|
|
||
| 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) |
There was a problem hiding this comment.
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.
|
|
||
| class TrainVizEvalVideo(EvalVideo): | ||
| def __init__( | ||
| self, base: EvalVideo, limit_val_batches: int = 50, prefix: str = "train_viz" |
There was a problem hiding this comment.
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.
| from egomimic.eval.eval_video import EvalVideo | ||
|
|
||
|
|
||
| class TrainVizEvalVideo(EvalVideo): |
There was a problem hiding this comment.
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.
|
|
||
| def on_validation_start(self): | ||
| if self.trainer.is_global_zero: | ||
| if self.trainer.is_global_zero and self._should_viz(): |
There was a problem hiding this comment.
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`. | ||
|
|
||
| # |
There was a problem hiding this comment.
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.
| assert (fd_shift > fd).all(), "mean shift must increase the distance" | ||
|
|
||
|
|
||
| def test_reverse_kl_is_gone(): |
There was a problem hiding this comment.
nit (test-quality): test_reverse_kl_is_gone only checks that a deleted function stays deleted and tests no behavior. Drop it.
| self.base = base | ||
| self.prefix = prefix | ||
| super().__init__( | ||
| limit_val_batches=limit_val_batches, |
There was a problem hiding this comment.
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.
| 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 |
There was a problem hiding this comment.
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]>
6fd89d1 to
7a13410
Compare
2ca1d41 to
05a92e1
Compare
|
Review findings: status after the fixup pass. The fixes are squashed into this PR's commits, and the whole stack was force-pushed. Fixed
Partial
Not fixed
|

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