From 05a92e177bf490c99615bac40c12bf156ef32c41 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Wed, 16 Sep 2026 18:18:11 -0400 Subject: [PATCH] eval: layout-aware frame-invariant metrics, and a train-split viz head 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) --- egomimic/algo/hpt.py | 2 - egomimic/eval/action_metrics.py | 219 ++++++++++ egomimic/eval/eval_hpt.py | 136 +------ egomimic/eval/eval_pi.py | 97 ++--- egomimic/eval/eval_train_viz.py | 68 ++++ egomimic/eval/eval_video.py | 134 ++++--- .../hydra_configs/evaluator/eval_hpt.yaml | 21 +- .../evaluator/eval_hpt_cartesian.yaml | 14 + .../evaluator/eval_hpt_keypoints_head.yaml | 17 + .../evaluator/eval_hpt_keypoints_wrist.yaml | 25 ++ .../evaluator/eval_hpt_wrist.yaml | 5 +- .../evaluator/eval_hpt_wristframe_6d.yaml | 23 ++ egomimic/hydra_configs/evaluator/eval_pi.yaml | 32 +- .../evaluator/eval_pi_camframe_6d.yaml | 29 ++ .../evaluator/eval_pi_wristframe.yaml | 5 + .../evaluator/eval_pi_wristframe_6d.yaml | 29 ++ .../train_viz_hpt_keypoints_wrist.yaml | 10 + .../train_viz_hpt_wristframe_6d.yaml | 10 + .../hydra_configs/evaluator/train_viz_pi.yaml | 9 + .../evaluator/train_viz_pi_camframe_6d.yaml | 10 + .../evaluator/train_viz_pi_wristframe_6d.yaml | 10 + .../evaluator/viz/cartesian_wrist.yaml | 4 +- .../evaluator/viz/keypoints.yaml | 7 +- .../evaluator/viz/keypoints_wrist.yaml | 6 +- .../viz/pi_keypoints_lang_wrist.yaml | 10 + .../hydra_configs/model/hpt_bc_flow_eva.yaml | 1 - .../model/hpt_bc_flow_human.yaml | 5 +- .../model/hpt_bc_keypoints_base.yaml | 2 - .../model/hpt_cotrain_enc_dec_base.yaml | 1 - .../hydra_configs/train_zarr_cartesian.yaml | 11 + .../train_zarr_keypoint_wrist.yaml | 2 +- .../hydra_configs/train_zarr_keypoints.yaml | 2 +- egomimic/pl_utils/pl_data_utils.py | 65 ++- egomimic/pl_utils/pl_model.py | 50 ++- egomimic/trainHydra.py | 171 +++++++- egomimic/utils/metrics.py | 64 +-- tests/unit/test_action_metrics.py | 378 ++++++++++++++++++ tests/unit/test_eva_wristframe_norm_pool.py | 12 + tests/unit/test_eval_video.py | 80 ++++ tests/unit/test_hydra_configs.py | 38 +- tests/unit/test_infra_port_wiring.py | 190 +++++++++ 41 files changed, 1701 insertions(+), 303 deletions(-) create mode 100644 egomimic/eval/action_metrics.py create mode 100644 egomimic/eval/eval_train_viz.py create mode 100644 egomimic/hydra_configs/evaluator/eval_hpt_cartesian.yaml create mode 100644 egomimic/hydra_configs/evaluator/eval_hpt_keypoints_head.yaml create mode 100644 egomimic/hydra_configs/evaluator/eval_hpt_keypoints_wrist.yaml create mode 100644 egomimic/hydra_configs/evaluator/eval_hpt_wristframe_6d.yaml create mode 100644 egomimic/hydra_configs/evaluator/eval_pi_camframe_6d.yaml create mode 100644 egomimic/hydra_configs/evaluator/eval_pi_wristframe.yaml create mode 100644 egomimic/hydra_configs/evaluator/eval_pi_wristframe_6d.yaml create mode 100644 egomimic/hydra_configs/evaluator/train_viz_hpt_keypoints_wrist.yaml create mode 100644 egomimic/hydra_configs/evaluator/train_viz_hpt_wristframe_6d.yaml create mode 100644 egomimic/hydra_configs/evaluator/train_viz_pi.yaml create mode 100644 egomimic/hydra_configs/evaluator/train_viz_pi_camframe_6d.yaml create mode 100644 egomimic/hydra_configs/evaluator/train_viz_pi_wristframe_6d.yaml create mode 100644 egomimic/hydra_configs/evaluator/viz/pi_keypoints_lang_wrist.yaml create mode 100644 tests/unit/test_action_metrics.py create mode 100644 tests/unit/test_eval_video.py diff --git a/egomimic/algo/hpt.py b/egomimic/algo/hpt.py index 8184f5ac74..979a536b79 100644 --- a/egomimic/algo/hpt.py +++ b/egomimic/algo/hpt.py @@ -912,8 +912,6 @@ def __init__( self.freeze_depth = kwargs.get("freeze_depth", 8) model.depth = self.depth - self.rkl_samples = kwargs.get("reverse_kl_samples", 4) - if self.ot: self.ot_warm_start_steps = kwargs.get("ot_warm_start_steps", 0) self.ot_6dof = kwargs.get("ot_6dof", False) diff --git a/egomimic/eval/action_metrics.py b/egomimic/eval/action_metrics.py new file mode 100644 index 0000000000..0264895277 --- /dev/null +++ b/egomimic/eval/action_metrics.py @@ -0,0 +1,219 @@ +"""Layout-aware action metrics shared by the PI and HPT evaluators. + +Every metric here is computed on the model's native (unnormalized) output and +is invariant to the rigid transform the evaluator's cam-frame revert applies +to prediction and ground truth alike, so it never needs the revert. + +Two layouts are recognised, by action key and then width: + +- bimanual cartesian (12/14 ypr, 18/20 rot6d; see + ``pose_utils.bimanual_cartesian_layout``): position MSE, geodesic rotation + error, raw rotation-channel MSE, DTW + Frechet on xyz; +- wrist-first bimanual hand keypoints (138 ypr / 144 rot6d wrist; see + ``pose_utils.bimanual_keypoint_layout``): mean keypoint L2 (m), wrist + position MSE, wrist geodesic rotation error, DTW on the keypoint block. + +Other keys (e.g. ``actions_joints``, whose 14-D width matches the robot ypr +layout) and unknown widths yield no metrics. +""" + +from __future__ import annotations + +import math + +import torch + +from egomimic.utils.action_utils import _reconstruct_R_from_cols, _ypr_to_matrix +from egomimic.utils.metrics import dtw_distance, frechet_gaussian_over_time +from egomimic.utils.pose_utils import ( + _CARTESIAN_KEYS, + _KEYPOINT_KEYS, + bimanual_cartesian_layout, + bimanual_keypoint_layout, +) + +_RAD2DEG = 180.0 / math.pi + + +def _paired_mse(pred: torch.Tensor, gt: torch.Tensor) -> torch.Tensor: + """Plain elementwise MSE, stateless (no torchmetrics accumulator).""" + return (pred.float() - gt.float()).pow(2).mean() + + +# Horizon segments for the per-segment L2 metrics, as fractions of the chunk: +# early = the first 10 % of timesteps, mid = 10-50 %, late = the second half. +# They separate near-term precision from long-range drift, which one +# chunk-mean error hides (and which per-timestep vs horizon-pooled norm stats +# trade against each other). +_HORIZON_SEGMENTS = (("early", 0.0, 0.1), ("mid", 0.1, 0.5), ("late", 0.5, 1.0)) + + +def _segment_l2(pred: torch.Tensor, gt: torch.Tensor, idx: list, prefix: str) -> dict: + """Mean euclidean error (m) of the xyz triples at ``idx`` per horizon + segment of a ``(B, T, D)`` chunk.""" + T = pred.shape[1] + diff = pred[..., idx].float() - gt[..., idx].float() + l2 = diff.reshape(*diff.shape[:-1], -1, 3).norm(dim=-1) # (B, T, P) + out = {} + for name, lo, hi in _HORIZON_SEGMENTS: + start = min(int(round(lo * T)), T - 1) + end = max(int(round(hi * T)), start + 1) + out[f"{prefix}_{name}"] = l2[:, start:end].mean() + return out + + +def _split_mse(pred_t: torch.Tensor, gt_t: torch.Tensor): + """(translation MSE, rotation-channel MSE) over a bimanual cartesian + vector, so a translation problem reads apart from a rotation one. The + rotation channels are whatever the layout holds: the continuous 6D + rotation-matrix columns for the 18D (human) / 20D (robot) widths, + yaw/pitch/roll for the 12/14D widths. Either way this is a plain + per-channel MSE (pre-Gram-Schmidt, axis-dependent, unitless for 6D; +-pi + wrap and gimbal lock for ypr), i.e. a training-tracking diagnostic, not a + rotation error; see :func:`_rot_geodesic_error` for that. Returns + (None, None) for an unknown width. + """ + layout = bimanual_cartesian_layout(pred_t.shape[-1]) + if layout is None: + return None, None + xyz_idx = list(layout["xyz"]) + rot_idx = list(layout["rot"]) + xyz = _paired_mse(pred_t[..., xyz_idx], gt_t[..., xyz_idx]) + rot = _paired_mse(pred_t[..., rot_idx], gt_t[..., rot_idx]) + return xyz, rot + + +def _geodesic_from_rot_channels(pred: torch.Tensor, gt: torch.Tensor, rot: list): + """Mean geodesic angle (rad) over both arms given the rotation channel + indices of a per-arm-concatenated vector (3 ypr or 6 rot6d per arm).""" + per_arm = len(rot) // 2 + # float64: arccos near 1 loses ~sqrt(eps), which in float32 puts a + # ~1e-4 rad floor under identical rotations; double makes it ~1e-8. + pred = pred.double() + gt = gt.double() + errs = [] + for arm in (rot[:per_arm], rot[per_arm:]): + p, g = pred[..., arm], gt[..., arm] + if per_arm == 3: + Rp, Rg = _ypr_to_matrix(p), _ypr_to_matrix(g) + else: + Rp = _reconstruct_R_from_cols(p[..., 0:3], p[..., 3:6]) + Rg = _reconstruct_R_from_cols(g[..., 0:3], g[..., 3:6]) + tr = (Rp.transpose(-1, -2) @ Rg).diagonal(dim1=-2, dim2=-1).sum(-1) + errs.append(torch.arccos(((tr - 1.0) / 2.0).clamp(-1.0, 1.0))) + return torch.stack(errs, dim=-1).mean().float() + + +def _rot_geodesic_error(pred: torch.Tensor, gt: torch.Tensor): + """Mean geodesic rotation error in radians over batch / time / both arms + of a bimanual cartesian vector. + + Euler-free: per arm, builds proper rotation matrices (``_ypr_to_matrix`` + for the 12/14-dim ypr widths; Gram-Schmidt on the two 6D columns for the + 18/20-dim widths, the same reconstruction the model decode uses) and + takes ``arccos((tr(R_pred^T R_gt) - 1) / 2)``. Unlike a ypr MSE this is + immune to the +-pi wrap AND to the yaw/roll degeneracy at pitch ~ +-pi/2. + Returns None for an unknown width. + """ + layout = bimanual_cartesian_layout(pred.shape[-1]) + if layout is None: + return None + return _geodesic_from_rot_channels(pred, gt, list(layout["rot"])) + + +def cartesian_metrics(pred: torch.Tensor, gt: torch.Tensor, prefix: str) -> dict: + """Metrics for a bimanual cartesian chunk ``(B, T, D)``; ``{}`` if ``D`` + is not a cartesian width. ``pred``/``gt`` may live on any device.""" + pred_cpu, gt_cpu = pred.detach().cpu(), gt.detach().cpu() + layout = bimanual_cartesian_layout(pred_cpu.shape[-1]) + if layout is None: + return {} + metrics = {} + xyz_idx = list(layout["xyz"]) + xyz_p, rot_p = _split_mse(pred_cpu, gt_cpu) + rot_name = "rot6d" if len(layout["rot"]) == 12 else "ypr" + metrics[f"{prefix}_xyz_paired_mse_avg"] = xyz_p + metrics[f"{prefix}_{rot_name}_paired_mse_avg"] = rot_p + metrics[f"{prefix}_xyz_final_mse_avg"] = _paired_mse( + pred_cpu[:, -1, xyz_idx], gt_cpu[:, -1, xyz_idx] + ) + metrics.update(_segment_l2(pred_cpu, gt_cpu, xyz_idx, f"{prefix}_xyz_l2")) + # Distributional / alignment metrics on the position channels only + # (metres), so a unitless 6D column never trades off against a metre. + # DTW forgives temporal misalignment (a correct motion executed early / + # late scores near zero where paired MSE penalizes it); Frechet compares + # the time-distribution shape of the motion. + pred_xyz = pred[..., xyz_idx] + gt_xyz = gt[..., xyz_idx].to(pred_xyz.device) + metrics[f"{prefix}_xyz_dtw_avg"] = dtw_distance(pred_xyz, gt_xyz).mean().item() + metrics[f"{prefix}_xyz_frechet_gauss_avg"] = ( + frechet_gaussian_over_time(pred_xyz, gt_xyz).mean().item() + ) + # Rotation error on the manifold, in degrees: what the decode actually + # produces after Gram-Schmidt, free of wrap / gimbal artefacts. + metrics[f"{prefix}_rot_err_deg_avg"] = ( + _rot_geodesic_error(pred_cpu, gt_cpu) * _RAD2DEG + ) + metrics[f"{prefix}_rot_err_deg_final"] = ( + _rot_geodesic_error(pred_cpu[:, -1], gt_cpu[:, -1]) * _RAD2DEG + ) + return metrics + + +def keypoint_metrics(pred: torch.Tensor, gt: torch.Tensor, prefix: str) -> dict: + """Metrics for a wrist-first bimanual keypoint chunk ``(B, T, D)`` + (D = 138 ypr wrist or 144 rot6d wrist); ``{}`` for other widths. + + ``kp_l2`` is the mean euclidean keypoint error in metres, in whatever + frame the layout holds (the wrist frame for the ``*_wristframe_*`` modes, + head frame otherwise); ``wrist_xyz`` / ``wrist_rot_err_deg`` score the + wrist pose channels; ``kp_dtw`` is DTW over the flattened keypoint block. + """ + pred_cpu, gt_cpu = pred.detach().cpu(), gt.detach().cpu() + layout = bimanual_keypoint_layout(pred_cpu.shape[-1]) + if layout is None: + return {} + kp_idx = list(layout["keypoints"]) + wrist_idx = list(layout["wrist_xyz"]) + rot_idx = list(layout["rot"]) + + def _kp_l2(p, g): + diff = p[..., kp_idx].float() - g[..., kp_idx].float() + return diff.reshape(*diff.shape[:-1], -1, 3).norm(dim=-1).mean() + + metrics = { + f"{prefix}_kp_l2_avg": _kp_l2(pred_cpu, gt_cpu), + f"{prefix}_kp_l2_final": _kp_l2(pred_cpu[:, -1], gt_cpu[:, -1]), + f"{prefix}_wrist_xyz_paired_mse_avg": _paired_mse( + pred_cpu[..., wrist_idx], gt_cpu[..., wrist_idx] + ), + f"{prefix}_wrist_xyz_final_mse_avg": _paired_mse( + pred_cpu[:, -1, wrist_idx], gt_cpu[:, -1, wrist_idx] + ), + f"{prefix}_wrist_rot_err_deg_avg": ( + _geodesic_from_rot_channels(pred_cpu, gt_cpu, rot_idx) * _RAD2DEG + ), + f"{prefix}_wrist_rot_err_deg_final": ( + _geodesic_from_rot_channels(pred_cpu[:, -1], gt_cpu[:, -1], rot_idx) + * _RAD2DEG + ), + } + metrics.update(_segment_l2(pred_cpu, gt_cpu, kp_idx, f"{prefix}_kp_l2")) + metrics.update(_segment_l2(pred_cpu, gt_cpu, wrist_idx, f"{prefix}_wrist_xyz_l2")) + pred_kp = pred[..., kp_idx] + gt_kp = gt[..., kp_idx].to(pred_kp.device) + metrics[f"{prefix}_kp_dtw_avg"] = dtw_distance(pred_kp, gt_kp).mean().item() + return metrics + + +def layout_metrics( + pred: torch.Tensor, gt: torch.Tensor, prefix: str, ac_key: str +) -> dict: + """Cartesian or keypoint metrics for ``ac_key``'s layout, or ``{}`` when + ``ac_key`` is neither: width alone would match e.g. 14-D joint angles to + the robot ypr layout.""" + if ac_key in _CARTESIAN_KEYS: + return cartesian_metrics(pred, gt, prefix) + if ac_key in _KEYPOINT_KEYS: + return keypoint_metrics(pred, gt, prefix) + return {} diff --git a/egomimic/eval/eval_hpt.py b/egomimic/eval/eval_hpt.py index ee7af2a0e2..d7490bcb86 100644 --- a/egomimic/eval/eval_hpt.py +++ b/egomimic/eval/eval_hpt.py @@ -3,9 +3,10 @@ import torch from torchmetrics import MeanSquaredError +from egomimic.eval.action_metrics import layout_metrics from egomimic.eval.eval_video import EvalVideo from egomimic.rldb.embodiment.embodiment import Embodiment, get_embodiment -from egomimic.utils.metrics import frechet_gaussian_over_time, reverse_kl_from_samples +from egomimic.utils.metrics import frechet_gaussian_over_time class HPTEvalVideo(EvalVideo): @@ -13,14 +14,13 @@ class HPTEvalVideo(EvalVideo): Eval class for HPT models. Per embodiment, computes: - val loss (BC loss, same as training; also aggregated as ``Valid/action_loss``) - paired/final MSE + Frechet over time for the main / shared / auxiliary heads - - paired/final MSE in cam frame on the main ``ac_key``, when a - ``transform_lists`` entry is configured - - optional Reverse KL from samples - The revert transform is applied once and reused by both the cam-frame MSE - and the viz video. + - the layout metrics of :mod:`egomimic.eval.action_metrics` on the main + ``ac_key`` (position / rotation / keypoint errors, frame-invariant) + Metrics are frame-invariant, so the ``transform_lists`` revert runs only + when rendering the viz video (``do_viz``), never for metrics. """ - def compute_metrics_and_viz(self, batch): + def compute_metrics_and_viz(self, batch, do_viz=True): algo = self.model preds = algo.forward_eval(batch) @@ -44,24 +44,22 @@ def compute_metrics_and_viz(self, batch): n_loss_embodiments += 1 if f"{embodiment_name}_{ac_key}" in preds and ac_key != algo.shared_ac_key: - metrics[f"Valid/{embodiment_name}_{ac_key}_paired_mse_avg"] = mse( - preds[f"{embodiment_name}_{ac_key}"].cpu(), _batch[ac_key].cpu() + pred_key = f"{embodiment_name}_{ac_key}" + metrics[f"Valid/{pred_key}_paired_mse_avg"] = mse( + preds[pred_key].cpu(), _batch[ac_key].cpu() ) - metrics[f"Valid/{embodiment_name}_{ac_key}_final_mse_avg"] = mse( - preds[f"{embodiment_name}_{ac_key}"][:, -1].cpu(), + metrics[f"Valid/{pred_key}_final_mse_avg"] = mse( + preds[pred_key][:, -1].cpu(), _batch[ac_key][:, -1].cpu(), ) - fd = frechet_gaussian_over_time( - preds[f"{embodiment_name}_{ac_key}"], _batch[ac_key] - ) - metrics[f"Valid/{embodiment_name}_{ac_key}_frechet_gauss_avg"] = ( - fd.mean().item() - ) - metrics[f"Valid/{embodiment_name}_{ac_key}_frechet_gauss_min"] = ( - fd.min().item() - ) - metrics[f"Valid/{embodiment_name}_{ac_key}_frechet_gauss_max"] = ( - fd.max().item() + fd = frechet_gaussian_over_time(preds[pred_key], _batch[ac_key]) + metrics[f"Valid/{pred_key}_frechet_gauss_avg"] = fd.mean().item() + metrics[f"Valid/{pred_key}_frechet_gauss_min"] = fd.min().item() + metrics[f"Valid/{pred_key}_frechet_gauss_max"] = fd.max().item() + metrics.update( + layout_metrics( + preds[pred_key], _batch[ac_key], f"Valid/{pred_key}", ac_key + ) ) if embodiment_name in algo.auxiliary_ac_keys: @@ -102,62 +100,8 @@ def compute_metrics_and_viz(self, batch): metrics[f"Valid/{pred_key}_frechet_gauss_min"] = fd.min().item() metrics[f"Valid/{pred_key}_frechet_gauss_max"] = fd.max().item() - if algo.rkl_samples and algo.rkl_samples > 1: - hpt_batch = { - "domain": embodiment_name, - "data": algo._robomimic_to_hpt_data( - batch[embodiment_id], - algo.camera_keys[embodiment_id], - algo.proprio_keys[embodiment_id], - algo.lang_keys[embodiment_id], - ac_key, - algo.auxiliary_ac_keys.get(embodiment_name, []), - ), - } - rkl_targets = [] - - if ( - f"{embodiment_name}_{ac_key}" in preds - and ac_key != algo.shared_ac_key - ): - rkl_targets.append( - ( - f"{embodiment_name}_{ac_key}", - _batch[ac_key].to(algo.device), - embodiment_name, - ) - ) - - if embodiment_name in algo.auxiliary_ac_keys: - for aux_key in algo.auxiliary_ac_keys[embodiment_name]: - aux_pred_key = f"{embodiment_name}_{aux_key}" - if aux_pred_key in preds: - rkl_targets.append( - ( - aux_pred_key, - _batch[aux_key].to(algo.device), - aux_key, - ) - ) - - if algo.shared_ac_key: - shared_pred_key = f"{embodiment_name}_{algo.shared_ac_key}" - if shared_pred_key in preds: - rkl_targets.append( - ( - shared_pred_key, - _batch[algo.shared_ac_key].to(algo.device), - "shared", - ) - ) - - M = int(algo.rkl_samples) - for pred_key_name, gt_tensor, head_key in rkl_targets: - samples = self._collect_policy_samples( - hpt_batch, ref=gt_tensor, key_name=head_key, M=M - ) - rkl = reverse_kl_from_samples(samples, gt_tensor) - metrics[f"Valid/{pred_key_name}_reverse_kl_M{M}"] = rkl.item() + if not do_viz: + continue transform_list = self.transform_lists.get(embodiment_name) main_pred_key = f"{embodiment_name}_{ac_key}" @@ -172,19 +116,6 @@ def compute_metrics_and_viz(self, batch): # (e.g. ``embodiment``, ``annotations``). Merge to preserve them. gt_batch_viz = {**_batch, **gt_t} pred_batch_viz = {**_batch, **pred_t} - - # ``.contiguous()`` because ``apply_transform`` returns CPU tensors, - # so ``.cpu()`` here is a no-op and ``[:, -1]`` leaves a non-contiguous - # view that torchmetrics' MSE doesn't accept. - metrics[f"Valid/{main_pred_key}_cam_paired_mse_avg"] = mse( - pred_batch_viz[ac_key].cpu().contiguous(), - gt_batch_viz[ac_key].cpu().contiguous(), - ) - metrics[f"Valid/{main_pred_key}_cam_final_mse_avg"] = mse( - pred_batch_viz[ac_key][:, -1].cpu().contiguous(), - gt_batch_viz[ac_key][:, -1].cpu().contiguous(), - ) - preds_for_viz = dict(preds) preds_for_viz[main_pred_key] = pred_batch_viz[ac_key] @@ -202,26 +133,3 @@ def _visualize_preds(self, predictions, batch): embodiment_id = batch["embodiment"][0].item() embodiment_name = get_embodiment(embodiment_id).lower() return self.viz_func[embodiment_name](predictions, batch) - - @torch.no_grad() - def _collect_policy_samples(self, hpt_batch, ref, key_name, M): - """Collect policy samples for Reverse KL.""" - algo = self.model - B, T, D = ref.shape - samples = [] - was_training = algo.nets.training - algo.nets.eval() - for _ in range(M): - out = algo.nets["policy"].forward( - hpt_batch["domain"], algo._clone_batch(hpt_batch["data"]) - ) - if key_name in out: - pred = out[key_name] - else: - pred = out[hpt_batch["domain"]] - - pred = pred[:, :T, :D] - samples.append(pred.unsqueeze(0)) - if was_training: - algo.nets.train() - return torch.cat(samples, dim=0) diff --git a/egomimic/eval/eval_pi.py b/egomimic/eval/eval_pi.py index 7206197701..4a1478609e 100644 --- a/egomimic/eval/eval_pi.py +++ b/egomimic/eval/eval_pi.py @@ -1,29 +1,39 @@ import copy import torch -from torchmetrics import MeanSquaredError +from egomimic.eval.action_metrics import _paired_mse, layout_metrics from egomimic.eval.eval_video import EvalVideo from egomimic.rldb.embodiment.embodiment import Embodiment, get_embodiment class PIEvalVideo(EvalVideo): """ - Eval class for PI models. Per embodiment, computes: - - val loss (flow-matching loss, same as training; also aggregated as ``Valid/action_loss``) - - paired/final MSE in the model's native wrist frame - - paired/final MSE in cam frame, when a ``transform_lists`` entry is configured - The revert transform is applied once and reused for both the cam-frame MSE - and the viz video. + Eval class for PI models. Per embodiment, on the model's native + (unnormalized) output, computes: + - ``action_loss``: flow-matching val loss, same as training (also logged + per embodiment as ``_loss``) + - the layout metrics of :mod:`egomimic.eval.action_metrics`: + cartesian widths get ``xyz_paired_mse_avg`` / ``xyz_final_mse_avg`` + (m^2), ``rot_err_deg_avg`` / ``rot_err_deg_final`` (geodesic, after + Gram-Schmidt), ``rot6d_paired_mse_avg`` / ``ypr_paired_mse_avg`` (raw + rotation channels, a training-tracking diagnostic), ``xyz_dtw_avg`` and + ``xyz_frechet_gauss_avg``; keypoint widths get ``kp_l2_avg`` / + ``kp_l2_final`` (m), ``wrist_xyz_*_mse``, ``wrist_rot_err_deg_*`` and + ``kp_dtw_avg``. Other action keys get a plain ``paired_mse_avg`` / + ``final_mse_avg``. + + Metrics are frame-invariant (the cam-frame revert is a rigid transform + shared by prediction and gt), so the ``transform_lists`` revert runs only + when rendering the viz video (``do_viz``), never for metrics. """ - def compute_metrics_and_viz(self, batch): + def compute_metrics_and_viz(self, batch, do_viz=True): algo = self.model preds = algo.forward_eval(batch) metrics = {} images_dict = {} - mse = MeanSquaredError() total_loss = None n_loss_embodiments = 0 @@ -43,43 +53,38 @@ def compute_metrics_and_viz(self, batch): n_loss_embodiments += 1 if pred_key in preds: - metrics[f"Valid/{pred_key}_paired_mse_avg"] = mse( - preds[pred_key].cpu(), _batch[ac_key].cpu() - ) - metrics[f"Valid/{pred_key}_final_mse_avg"] = mse( - preds[pred_key][:, -1].cpu(), _batch[ac_key][:, -1].cpu() - ) - - transform_list = self.transform_lists.get(embodiment_name) - gt_batch_viz = _batch - preds_for_viz = preds - if transform_list is not None and pred_key in preds: - pred_batch = copy.deepcopy(_batch) - pred_batch[ac_key] = preds[pred_key] - gt_t = Embodiment.apply_transform(_batch, transform_list) - pred_t = Embodiment.apply_transform(pred_batch, transform_list) - # apply_transform drops keys whose shape[0] != batch_size - # (e.g. ``embodiment``, ``annotations``). Merge to preserve them. - gt_batch_viz = {**_batch, **gt_t} - pred_batch_viz = {**_batch, **pred_t} - - # ``.contiguous()`` because ``apply_transform`` returns CPU tensors, - # so ``.cpu()`` here is a no-op and ``[:, -1]`` leaves a non-contiguous - # view that torchmetrics' MSE doesn't accept. - metrics[f"Valid/{pred_key}_cam_paired_mse_avg"] = mse( - pred_batch_viz[ac_key].cpu().contiguous(), - gt_batch_viz[ac_key].cpu().contiguous(), - ) - metrics[f"Valid/{pred_key}_cam_final_mse_avg"] = mse( - pred_batch_viz[ac_key][:, -1].cpu().contiguous(), - gt_batch_viz[ac_key][:, -1].cpu().contiguous(), - ) - - preds_for_viz = dict(preds) - preds_for_viz[pred_key] = pred_batch_viz[ac_key] - - ims = self._visualize_preds(preds_for_viz, gt_batch_viz) - images_dict[embodiment_id] = ims + pred, gt = preds[pred_key], _batch[ac_key] + layout = layout_metrics(pred, gt, f"Valid/{pred_key}", ac_key) + if not layout: + layout = { + f"Valid/{pred_key}_paired_mse_avg": _paired_mse( + pred.cpu(), gt.cpu() + ), + f"Valid/{pred_key}_final_mse_avg": _paired_mse( + pred[:, -1].cpu(), gt[:, -1].cpu() + ), + } + metrics.update(layout) + + if do_viz: + gt_batch_viz = _batch + preds_for_viz = preds + transform_list = self.transform_lists.get(embodiment_name) + if transform_list is not None and pred_key in preds: + # Revert to cam (head) frame + the ypr layout for the + # overlay only; the metrics above are frame-invariant. + pred_batch = copy.deepcopy(_batch) + pred_batch[ac_key] = preds[pred_key] + gt_t = Embodiment.apply_transform(_batch, transform_list) + pred_t = Embodiment.apply_transform(pred_batch, transform_list) + # apply_transform drops keys whose shape[0] != batch_size + # (e.g. ``embodiment``, ``annotations``). Merge to preserve them. + gt_batch_viz = {**_batch, **gt_t} + pred_batch_viz = {**_batch, **pred_t} + preds_for_viz = dict(preds) + preds_for_viz[pred_key] = pred_batch_viz[ac_key] + ims = self._visualize_preds(preds_for_viz, gt_batch_viz) + images_dict[embodiment_id] = ims if total_loss is not None and n_loss_embodiments > 0: metrics["Valid/action_loss"] = total_loss / n_loss_embodiments diff --git a/egomimic/eval/eval_train_viz.py b/egomimic/eval/eval_train_viz.py new file mode 100644 index 0000000000..fa1d6ef9ff --- /dev/null +++ b/egomimic/eval/eval_train_viz.py @@ -0,0 +1,68 @@ +"""Prefixed evaluator heads. + +Wraps a concrete EvalVideo (PIEvalVideo, HPTEvalVideo, ...) so the same +forward/metric/viz logic runs against a second val loader: the train split +(``train_viz``), the unseen-operator split (``unseen_op_valid``) or a renamed +canonical head (``data.valid_prefix``). Metric keys are prefixed with +``/`` and videos go to ``/videos_/``. + +trainHydra builds the train_viz head around the canonical ``evaluator`` by +default; the ``evaluator/train_viz_*.yaml`` configs pin a different ``base`` +via ``+evaluator@train_viz_evaluator=``. That ``base`` is composed +separately from the canonical evaluator, so CLI overrides of +``evaluator.viz_every_n_epochs`` / ``viz_max_batches`` do not reach it; +override ``train_viz_evaluator.base.`` as well. +""" + +from __future__ import annotations + +import os + +from egomimic.eval.eval_video import EvalVideo + + +class PrefixedEvalVideo(EvalVideo): + """Runs ``base`` on another val loader under a ``prefix`` for metric keys + and the ``videos_/`` directory.""" + + def __init__(self, base: EvalVideo, prefix: str = "train_viz"): + # `base` must be set before super().__init__: the trainer/model + # property setters fire on the base attribute during construction. + self.base = base + self.prefix = prefix + super().__init__( + viz_func=base.viz_func, + transform_lists=base.transform_lists, + viz_every_n_epochs=base.viz_every_n_epochs, + viz_max_batches=base.viz_max_batches, + ) + + @property + def trainer(self): + return self._trainer + + @trainer.setter + def trainer(self, value): + self._trainer = value + self.base.trainer = value + + @property + def model(self): + return self._model + + @model.setter + def model(self, value): + self._model = value + self.base.model = value + + def video_dir(self): + return os.path.join(self.root_dir(), f"videos_{self.prefix}") + + def compute_metrics_and_viz(self, batch, do_viz=True): + metrics, images_dict = self.base.compute_metrics_and_viz(batch, do_viz=do_viz) + metrics = {f"{self.prefix}/{k}": v for k, v in metrics.items()} + return metrics, images_dict + + +# Configs and launch scripts written against the original name. +TrainVizEvalVideo = PrefixedEvalVideo diff --git a/egomimic/eval/eval_video.py b/egomimic/eval/eval_video.py index 9cf44d121c..cf662e54f5 100644 --- a/egomimic/eval/eval_video.py +++ b/egomimic/eval/eval_video.py @@ -20,14 +20,24 @@ def __init__( limit_val_batches: int = 400, viz_func: dict = None, transform_lists: dict | None = None, + viz_every_n_epochs: int = 1, + viz_max_batches: int | None = None, ): super().__init__() self.trainer = None self.model = None self.viz_func = viz_func + # Render the overlay video only on validation passes where + # (current_epoch + 1) is a multiple of this; metrics log every pass. + self.viz_every_n_epochs = viz_every_n_epochs + # On viz epochs, render overlay frames for only the first N val + # batches (None = all). Rendering is CPU-bound (~1s/frame) and + # dominates viz-epoch wall time; metrics are still computed on EVERY + # batch regardless. + self.viz_max_batches = viz_max_batches # Per-embodiment list[Transform] applied once during eval to project - # the model's wrist-frame actions back into cam (head) frame. Reused for - # both cam-frame MSE and the viz video so we don't transform twice. + # the model's wrist-frame actions back into cam (head) frame for the + # viz video. self.transform_lists = transform_lists or {} self.val_image_buffer = {} self.val_counter = {} @@ -44,52 +54,77 @@ def __init__( def video_dir(self): return os.path.join(self.root_dir(), "videos") + def _should_viz(self) -> bool: + if (self.viz_every_n_epochs or 0) <= 0: + return False + # Lightning runs validation when (current_epoch + 1) is a multiple of + # check_val_every_n_epoch, with current_epoch still the pre-increment + # value during the val hooks (19, 39, ...). Gate on the same (+1) + # convention: a plain `current_epoch % n` never aligns with a + # validation epoch unless check_val_every_n_epoch == 1. + epoch = self.trainer.current_epoch + 1 + # The last epoch always renders, so short runs (trainer=debug, eval + # mode's single validate) get a video without retuning the interval. + max_epochs = getattr(self.trainer, "max_epochs", None) + if max_epochs is not None and max_epochs > 0 and epoch == max_epochs: + return True + return epoch % self.viz_every_n_epochs == 0 + @abstractmethod - def compute_metrics_and_viz(self, batch): + def compute_metrics_and_viz(self, batch, do_viz=True): """ Run the model's eval forward and compute metrics and visualization frames. Args: batch (dict): processed batch produced by the algo's `process_batch_for_training`. + do_viz (bool): render overlay frames for this batch. When False the + returned ``images_dict`` must be empty. Returns: metrics (dict[str, torch.Tensor | float]) images_dict (dict[embodiment_id, np.ndarray (B, H, W, 3)]) """ raise NotImplementedError + def _video_fps(self, source_fps: int = 30) -> int: + """Playback fps compensating for the DistributedSampler stride. + + Distributed validation deals the (unshuffled) val indices round-robin, + so rank 0's consecutive rendered frames are ``world_size`` source + frames apart. Writing them at the raw source fps produces a + world_size-times timelapse; scale playback down to restore wall-clock + pacing.""" + world = max(1, int(getattr(self.trainer, "world_size", 1) or 1)) + return max(1, round(source_fps / world)) + + def _write_video(self, key, frames) -> None: + path = os.path.join( + self.video_dir(), + f"epoch_{self.trainer.current_epoch}", + str(get_embodiment(key)), + f"validation_video_{self.val_counter[key]}.mp4", + ) + os.makedirs(os.path.dirname(path), exist_ok=True) + tvio.write_video(path, frames, fps=self._video_fps(), video_codec="h264") + def on_validation_start(self): - if self.trainer.is_global_zero: - os.makedirs( - os.path.join(self.video_dir(), f"epoch_{self.trainer.current_epoch}"), - exist_ok=True, - ) + # _write_video creates the epoch directory it writes into. + pass def on_validation_end(self): + if not self._should_viz(): + return for key, buffer in self.val_image_buffer.items(): - os.makedirs( - os.path.join( - self.video_dir(), - f"epoch_{self.trainer.current_epoch}", - str(get_embodiment(key)), - ), - exist_ok=True, - ) if len(buffer) != 0: - frames = torch.stack(buffer) - path = os.path.join( - self.video_dir(), - f"epoch_{self.trainer.current_epoch}", - str(get_embodiment(key)), - f"validation_video_{self.val_counter[key]}.mp4", - ) - tvio.write_video(path, frames, fps=30, video_codec="h264") - + self._write_video(key, torch.stack(buffer)) self.val_counter[key] = 0 self.val_image_buffer[key] = [] def on_validation_step(self, batch, batch_idx, dataloader_idx=0): - metrics, images_dict = self.compute_metrics_and_viz(batch) + do_viz = self._should_viz() and ( + self.viz_max_batches is None or batch_idx < self.viz_max_batches + ) + metrics, images_dict = self.compute_metrics_and_viz(batch, do_viz=do_viz) device = self.trainer.lightning_module.device metrics = { @@ -97,30 +132,23 @@ def on_validation_step(self, batch, batch_idx, dataloader_idx=0): for k, v in metrics.items() } - ## images is now a dict - for key, images in images_dict.items(): - os.makedirs( - os.path.join( - self.video_dir(), - f"epoch_{self.trainer.current_epoch}", - str(get_embodiment(key)), - ), - exist_ok=True, - ) - if key not in self.val_image_buffer or self.val_image_buffer[key] is None: - self.val_image_buffer[key] = [] - self.val_counter[key] = 0 - self.val_image_buffer[key].extend(torch.from_numpy(images)) - if len(self.val_image_buffer[key]) >= 1000: - frames = torch.stack(self.val_image_buffer[key]) - path = os.path.join( - self.video_dir(), - f"epoch_{self.trainer.current_epoch}", - str(get_embodiment(key)), - f"validation_video_{self.val_counter[key]}.mp4", - ) - tvio.write_video(path, frames, fps=30, video_codec="h264") - self.val_image_buffer[key].clear() - self.val_counter[key] += 1 - - self.trainer.lightning_module.log_dict(metrics, sync_dist=True) + if do_viz: + for key, images in images_dict.items(): + if ( + key not in self.val_image_buffer + or self.val_image_buffer[key] is None + ): + self.val_image_buffer[key] = [] + self.val_counter[key] = 0 + self.val_image_buffer[key].extend(torch.from_numpy(images)) + if len(self.val_image_buffer[key]) >= 1000: + self._write_video(key, torch.stack(self.val_image_buffer[key])) + self.val_image_buffer[key].clear() + self.val_counter[key] += 1 + + # add_dataloader_idx=False: with the train_viz loader Lightning would + # otherwise suffix every key with "/dataloader_idx_N"; the train-viz + # wrapper prefixes its keys itself. + self.trainer.lightning_module.log_dict( + metrics, sync_dist=True, add_dataloader_idx=False + ) diff --git a/egomimic/hydra_configs/evaluator/eval_hpt.yaml b/egomimic/hydra_configs/evaluator/eval_hpt.yaml index 901be30749..9eaed9f541 100644 --- a/egomimic/hydra_configs/evaluator/eval_hpt.yaml +++ b/egomimic/hydra_configs/evaluator/eval_hpt.yaml @@ -1,5 +1,24 @@ +# Default HPT evaluator, matching the per-vendor data defaults: eva on the +# wrist-frame cartesian 6D action (reverted to cam-frame ypr for the overlay), +# human on the 144-D wrist-frame hand-keypoint action (reverted to head-frame +# keypoints for the overlay). +# For cartesian human data use eval_hpt_cartesian / eval_hpt_wrist. defaults: - - viz@viz_func: cartesian + - viz@viz_func: keypoints_wrist - _self_ _target_: egomimic.eval.eval_hpt.HPTEvalVideo + +# Render the val viz video only when (current_epoch + 1) is a multiple of +# this; metrics still log on every validation pass. Keep it a multiple of +# trainer.check_val_every_n_epoch or viz never coincides with a validation run. +viz_every_n_epochs: 200 +# On viz epochs, render only the first N val batches per rank; metrics still +# run on every batch. +viz_max_batches: 10 + +transform_lists: + eva_bimanual: + _target_: egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_wristframe_transform_list + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_keypoints_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_hpt_cartesian.yaml b/egomimic/hydra_configs/evaluator/eval_hpt_cartesian.yaml new file mode 100644 index 0000000000..576a384562 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_hpt_cartesian.yaml @@ -0,0 +1,14 @@ +# HPT evaluator for cam/head-frame CARTESIAN ypr data on every embodiment +# (the pre-keypoint default; e.g. model=hpt_bc_flow_human_cartesian with a +# cartesian data override). No revert: the actions are already in the frame +# the overlay draws in. +defaults: + - viz@viz_func: cartesian + - _self_ + +_target_: egomimic.eval.eval_hpt.HPTEvalVideo + +viz_every_n_epochs: 200 +# On viz epochs, render only the first N val batches per rank; metrics still +# run on every batch. +viz_max_batches: 10 diff --git a/egomimic/hydra_configs/evaluator/eval_hpt_keypoints_head.yaml b/egomimic/hydra_configs/evaluator/eval_hpt_keypoints_head.yaml new file mode 100644 index 0000000000..abb99f2f8f --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_hpt_keypoints_head.yaml @@ -0,0 +1,17 @@ +# HPT head-frame keypoint evaluator: use with a data config that applies the +# `keypoints_headframe_6d` transform mode (e.g. data=aria_keypoints). The +# keypoints are already in head frame; the revert only converts the wrist +# rotation rot6d -> ypr (144 -> 138) so the overlay's keypoint splitter sees +# the ypr layout. +defaults: + - viz@viz_func: keypoints + - _self_ + +_target_: egomimic.eval.eval_hpt.HPTEvalVideo + +viz_every_n_epochs: 100 +viz_max_batches: 10 + +transform_lists: + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_keypoints_revert_6d_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_hpt_keypoints_wrist.yaml b/egomimic/hydra_configs/evaluator/eval_hpt_keypoints_wrist.yaml new file mode 100644 index 0000000000..d9d5a34204 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_hpt_keypoints_wrist.yaml @@ -0,0 +1,25 @@ +# HPT wrist-frame keypoint evaluator: use with a data config that applies the +# `keypoints_wristframe_6d` transform mode (e.g. +# mecka_fold_freeform_opsplit_hpt_keypoints, or the per-vendor defaults). The +# model predicts the 144-D wrist-first action [wrist xyz+rot6d (9) | 21 +# keypoints (63)] x 2 in each hand's wrist frame; the revert converts the +# wrist rotation back to ypr and projects the two 63-D keypoint chunks back to +# head (camera) frame via the head-frame wrist pose in +# `observations.state.keypoints`, giving a 126-D head-frame +# `actions_keypoints` for the overlay video. Metrics (paired/final MSE, +# Frechet, the keypoint layout metrics) run on the native 144-D output. +defaults: + - viz@viz_func: keypoints_wrist + - _self_ + +_target_: egomimic.eval.eval_hpt.HPTEvalVideo + +# Render the val viz video only when (current_epoch + 1) is a multiple of this. +viz_every_n_epochs: 100 +# On viz epochs, render only the first N val batches per rank; metrics still +# run on every batch. The train_viz wrapper inherits this cap. +viz_max_batches: 10 + +transform_lists: + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_keypoints_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_hpt_wrist.yaml b/egomimic/hydra_configs/evaluator/eval_hpt_wrist.yaml index fc2099c190..143da6b33f 100644 --- a/egomimic/hydra_configs/evaluator/eval_hpt_wrist.yaml +++ b/egomimic/hydra_configs/evaluator/eval_hpt_wrist.yaml @@ -4,9 +4,8 @@ defaults: _target_: egomimic.eval.eval_hpt.HPTEvalVideo -# Per-embodiment revert transform. Applied once during validation to project -# the model's wrist-frame action chunks back to cam (head) frame, then reused -# for both the cam-frame MSE and the viz video. +# Per-embodiment revert transform: projects the model's wrist-frame action +# chunks back to cam (head) frame for the viz video. transform_lists: eva_bimanual: _target_: egomimic.rldb.embodiment.eva._build_eva_bimanual_revert_eef_frame_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_hpt_wristframe_6d.yaml b/egomimic/hydra_configs/evaluator/eval_hpt_wristframe_6d.yaml new file mode 100644 index 0000000000..00559d38c1 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_hpt_wristframe_6d.yaml @@ -0,0 +1,23 @@ +# HPT wrist-frame 6D CARTESIAN evaluator: HPT twin of eval_pi_wristframe_6d. +# Use with a data config that applies the `cartesian_wristframe_6d` transform +# mode with `keymap_mode: cartesian` (e.g. mecka_fold_flagship_opsplit_hpt_6d). +# The model predicts the 18-D [xyz | 6D rot] x 2 action in each wrist's frame; +# the revert converts 6D -> ypr (Gram-Schmidt), unpads the 20-D proprio, and +# projects predictions + gt back to head (cam) frame for the overlay video. +# Metrics (paired/final MSE, Frechet, xyz / rotation errors) run on the +# native 18-D output. +defaults: + - viz@viz_func: cartesian_wrist_lang + - _self_ + +_target_: egomimic.eval.eval_hpt.HPTEvalVideo + +# Render the val viz video only when (current_epoch + 1) is a multiple of this. +viz_every_n_epochs: 100 +# On viz epochs, render only the first N val batches per rank; metrics still +# run on every batch. The train_viz wrapper inherits this cap. +viz_max_batches: 10 + +transform_lists: + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_cartesian_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_pi.yaml b/egomimic/hydra_configs/evaluator/eval_pi.yaml index 97b9c1a1f3..62fdd2c39b 100644 --- a/egomimic/hydra_configs/evaluator/eval_pi.yaml +++ b/egomimic/hydra_configs/evaluator/eval_pi.yaml @@ -1,18 +1,30 @@ +# Default Pi evaluator, matching the cotrain_pi_base data defaults: eva on the +# wrist-frame cartesian 6D action, human on the 144-D wrist-frame hand-keypoint +# action. Metrics are computed on the native output (frame-invariant; see +# egomimic/eval/action_metrics.py) and never go through the reverts below, +# which only feed the viz overlay: +# eva: xyz+6D -> xyz+ypr, then each arm's chunk wrist -> cam frame via +# the cam-frame proprio ee_pose +# human: wrist rot6d -> ypr, then each hand's keypoints wrist -> head frame +# via the head-frame wrist pose in observations.state.keypoints +# For the all-cartesian recipes use eval_pi_camframe_6d / eval_pi_wristframe_6d. defaults: - - viz@viz_func: pi_cartesian_lang + - viz@viz_func: pi_keypoints_lang_wrist - _self_ _target_: egomimic.eval.eval_pi.PIEvalVideo -# Per-embodiment revert transform. Applied once during validation to project -# the model's wrist-frame action chunks back to cam (head) frame, then reused -# for both the cam-frame MSE and the viz video. Each value resolves to a -# list[Transform] via its ``_target_``. Must match whichever viz config is -# mounted above. +# Render the val viz video only when (current_epoch + 1) is a multiple of +# this; metrics still log on every validation pass. Keep it a multiple of +# trainer.check_val_every_n_epoch or viz never coincides with a validation +# run. Without it the EvalVideo default of 1 renders video on every pass. +viz_every_n_epochs: 100 +# On viz epochs, render only the first N val batches per rank (~1s/frame CPU +# draw); metrics still run on every batch. +viz_max_batches: 10 + transform_lists: eva_bimanual: - _target_: egomimic.rldb.embodiment.eva._build_eva_bimanual_revert_eef_frame_transform_list - is_quat: false + _target_: egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_wristframe_transform_list human_bimanual: - _target_: egomimic.rldb.embodiment.human._build_human_cartesian_revert_eef_frame_transform_list - is_quat: false + _target_: egomimic.rldb.embodiment.human._build_human_keypoints_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_pi_camframe_6d.yaml b/egomimic/hydra_configs/evaluator/eval_pi_camframe_6d.yaml new file mode 100644 index 0000000000..eceaf649de --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_pi_camframe_6d.yaml @@ -0,0 +1,29 @@ +defaults: + - viz@viz_func: pi_cartesian_lang + - _self_ + +_target_: egomimic.eval.eval_pi.PIEvalVideo + +# Render the val viz video only when (current_epoch + 1) is a multiple of +# this; metrics still log on every validation pass. Keep it a multiple of +# trainer.check_val_every_n_epoch or viz never coincides with a validation +# run. +viz_every_n_epochs: 100 +# On viz epochs, render only the first N val batches per rank; metrics still +# run on every batch. The train_viz wrapper inherits this cap. +viz_max_batches: 10 + +# All-cartesian cam-frame evaluator: use when every embodiment expresses its +# cartesian action in head/camera frame with the continuous 6D rotation +# (`transform_list: mode: cartesian_6d`, e.g. data=mecka_fold_freeform_6d +# with data.train_datasets.human_bimanual.resolver.transform_list.mode=cartesian_6d). +# Actions are already in cam (head) frame, so +# no frame change is needed; the revert only converts the rotation back from +# xyz+6D (9/arm) to xyz+ypr (6/arm) for the viz overlay, which draws in ypr. +# Metrics are computed on the native 6D output and never go through this +# revert. Each value resolves to a list[Transform] via its ``_target_``. +transform_lists: + eva_bimanual: + _target_: egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_transform_list + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_cartesian_revert_6d_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_pi_wristframe.yaml b/egomimic/hydra_configs/evaluator/eval_pi_wristframe.yaml new file mode 100644 index 0000000000..7f83600b7a --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_pi_wristframe.yaml @@ -0,0 +1,5 @@ +# Alias of eval_pi (data=cotrain_pi_lang_wrist): the default Pi evaluator is +# already the wrist-frame one. +defaults: + - eval_pi + - _self_ diff --git a/egomimic/hydra_configs/evaluator/eval_pi_wristframe_6d.yaml b/egomimic/hydra_configs/evaluator/eval_pi_wristframe_6d.yaml new file mode 100644 index 0000000000..581d454fc0 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_pi_wristframe_6d.yaml @@ -0,0 +1,29 @@ +defaults: + - viz@viz_func: pi_cartesian_lang + - _self_ + +_target_: egomimic.eval.eval_pi.PIEvalVideo + +# Render the val viz video only when (current_epoch + 1) is a multiple of +# this (see eval_pi_camframe_6d.yaml); without it the EvalVideo default of 1 +# renders video on every validation pass. +viz_every_n_epochs: 100 +# On viz epochs, render only the first N val batches per rank (~1s/frame +# CPU draw; unbounded rendering added ~2h per viz epoch on the combined +# pretrain). Metrics still run on every batch. The train_viz wrapper +# inherits this cap. +viz_max_batches: 10 + +# All-cartesian wrist-frame evaluator: use with a data config that applies the +# `cartesian_wristframe_6d` transform mode (e.g. mecka_all_6d, +# mecka_fold_freeform*_6d). The model predicts in each wrist's local frame +# using the continuous 6D rotation; these revert transforms first convert the +# rotation xyz+6D -> xyz+ypr and then project predictions + gt back to cam +# (head) frame for the viz video only (metrics are frame-invariant and +# computed on the native output). Each value resolves to a list[Transform] +# via its ``_target_``. Must match the viz config mounted above. +transform_lists: + eva_bimanual: + _target_: egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_wristframe_transform_list + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_cartesian_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/train_viz_hpt_keypoints_wrist.yaml b/egomimic/hydra_configs/evaluator/train_viz_hpt_keypoints_wrist.yaml new file mode 100644 index 0000000000..cf58f65da4 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/train_viz_hpt_keypoints_wrist.yaml @@ -0,0 +1,10 @@ +# Keypoint HPT train_viz head around eval_hpt_keypoints_wrist. Optional: +# train_viz wraps the canonical evaluator by default; select this to pin a +# different base with +# `+evaluator@train_viz_evaluator=train_viz_hpt_keypoints_wrist`. See +# egomimic/eval/eval_train_viz.py. +defaults: + - /evaluator@base: eval_hpt_keypoints_wrist + - _self_ + +_target_: egomimic.eval.eval_train_viz.PrefixedEvalVideo diff --git a/egomimic/hydra_configs/evaluator/train_viz_hpt_wristframe_6d.yaml b/egomimic/hydra_configs/evaluator/train_viz_hpt_wristframe_6d.yaml new file mode 100644 index 0000000000..cbc2432523 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/train_viz_hpt_wristframe_6d.yaml @@ -0,0 +1,10 @@ +# Wrist-frame 6D cartesian HPT train_viz head around eval_hpt_wristframe_6d. +# Optional: train_viz wraps the canonical evaluator by default; select this to +# pin a different base with +# `+evaluator@train_viz_evaluator=train_viz_hpt_wristframe_6d`. See +# egomimic/eval/eval_train_viz.py. +defaults: + - /evaluator@base: eval_hpt_wristframe_6d + - _self_ + +_target_: egomimic.eval.eval_train_viz.PrefixedEvalVideo diff --git a/egomimic/hydra_configs/evaluator/train_viz_pi.yaml b/egomimic/hydra_configs/evaluator/train_viz_pi.yaml new file mode 100644 index 0000000000..e73f496140 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/train_viz_pi.yaml @@ -0,0 +1,9 @@ +# Default Pi train_viz head around eval_pi. Optional: train_viz wraps the +# canonical evaluator by default; select this to pin a different base with +# `+evaluator@train_viz_evaluator=train_viz_pi`. See +# egomimic/eval/eval_train_viz.py. +defaults: + - /evaluator@base: eval_pi + - _self_ + +_target_: egomimic.eval.eval_train_viz.PrefixedEvalVideo diff --git a/egomimic/hydra_configs/evaluator/train_viz_pi_camframe_6d.yaml b/egomimic/hydra_configs/evaluator/train_viz_pi_camframe_6d.yaml new file mode 100644 index 0000000000..882474d7cc --- /dev/null +++ b/egomimic/hydra_configs/evaluator/train_viz_pi_camframe_6d.yaml @@ -0,0 +1,10 @@ +# Cam-frame all-cartesian Pi train_viz head around eval_pi_camframe_6d (needs +# the human transform mode overridden to `cartesian_6d`). Optional: train_viz +# wraps the canonical evaluator by default; select this to pin a different base +# with `+evaluator@train_viz_evaluator=train_viz_pi_camframe_6d`. See +# egomimic/eval/eval_train_viz.py. +defaults: + - /evaluator@base: eval_pi_camframe_6d + - _self_ + +_target_: egomimic.eval.eval_train_viz.PrefixedEvalVideo diff --git a/egomimic/hydra_configs/evaluator/train_viz_pi_wristframe_6d.yaml b/egomimic/hydra_configs/evaluator/train_viz_pi_wristframe_6d.yaml new file mode 100644 index 0000000000..991123ada7 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/train_viz_pi_wristframe_6d.yaml @@ -0,0 +1,10 @@ +# Wrist-frame all-cartesian Pi train_viz head around eval_pi_wristframe_6d. +# Optional: train_viz wraps the canonical evaluator by default; select this to +# pin a different base with +# `+evaluator@train_viz_evaluator=train_viz_pi_wristframe_6d`. See +# egomimic/eval/eval_train_viz.py. +defaults: + - /evaluator@base: eval_pi_wristframe_6d + - _self_ + +_target_: egomimic.eval.eval_train_viz.PrefixedEvalVideo diff --git a/egomimic/hydra_configs/evaluator/viz/cartesian_wrist.yaml b/egomimic/hydra_configs/evaluator/viz/cartesian_wrist.yaml index e4aae52a72..9863646876 100644 --- a/egomimic/hydra_configs/evaluator/viz/cartesian_wrist.yaml +++ b/egomimic/hydra_configs/evaluator/viz/cartesian_wrist.yaml @@ -1,12 +1,12 @@ eva_bimanual: _target_: egomimic.rldb.embodiment.eva.Eva.viz_gt_preds _partial_: true - image_key: front_img_1 + image_key: observations.images.front_img_1 action_key: actions_cartesian mode: traj human_bimanual: _target_: egomimic.rldb.embodiment.human.Human.viz_gt_preds _partial_: true - image_key: front_img_1 + image_key: observations.images.front_img_1 action_key: actions_cartesian mode: traj diff --git a/egomimic/hydra_configs/evaluator/viz/keypoints.yaml b/egomimic/hydra_configs/evaluator/viz/keypoints.yaml index cf5d1d218c..8f000e5ca7 100644 --- a/egomimic/hydra_configs/evaluator/viz/keypoints.yaml +++ b/egomimic/hydra_configs/evaluator/viz/keypoints.yaml @@ -1,12 +1,15 @@ +# Head-frame keypoint overlay: eva draws its cartesian trajectory, human its +# MANO keypoints (the 138-D ypr layout after the evaluator's 6d -> ypr revert, +# or 126-D head-frame keypoints after the wrist-frame revert). eva_bimanual: _target_: egomimic.rldb.embodiment.eva.Eva.viz_gt_preds _partial_: true - image_key: front_img_1 + image_key: observations.images.front_img_1 action_key: actions_cartesian mode: traj human_bimanual: _target_: egomimic.rldb.embodiment.human.Human.viz_gt_preds _partial_: true - image_key: front_img_1 + image_key: observations.images.front_img_1 action_key: actions_keypoints mode: keypoints diff --git a/egomimic/hydra_configs/evaluator/viz/keypoints_wrist.yaml b/egomimic/hydra_configs/evaluator/viz/keypoints_wrist.yaml index cf5d1d218c..2f92baf889 100644 --- a/egomimic/hydra_configs/evaluator/viz/keypoints_wrist.yaml +++ b/egomimic/hydra_configs/evaluator/viz/keypoints_wrist.yaml @@ -1,12 +1,14 @@ +# Wrist-frame keypoint overlay (same drawing as keypoints.yaml; the evaluator's +# transform_lists revert the wrist-frame prediction to head frame first). eva_bimanual: _target_: egomimic.rldb.embodiment.eva.Eva.viz_gt_preds _partial_: true - image_key: front_img_1 + image_key: observations.images.front_img_1 action_key: actions_cartesian mode: traj human_bimanual: _target_: egomimic.rldb.embodiment.human.Human.viz_gt_preds _partial_: true - image_key: front_img_1 + image_key: observations.images.front_img_1 action_key: actions_keypoints mode: keypoints diff --git a/egomimic/hydra_configs/evaluator/viz/pi_keypoints_lang_wrist.yaml b/egomimic/hydra_configs/evaluator/viz/pi_keypoints_lang_wrist.yaml new file mode 100644 index 0000000000..65d79d73a4 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/viz/pi_keypoints_lang_wrist.yaml @@ -0,0 +1,10 @@ +# Pi twin of keypoints_wrist: eva cartesian trajectory + human MANO keypoint +# overlay with the sampled prompt drawn on the frame. +defaults: + - keypoints_wrist + - _self_ + +eva_bimanual: + annotation_key: sampled_prompt +human_bimanual: + annotation_key: sampled_prompt diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_eva.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_eva.yaml index 196dc7a716..6440e4dc08 100644 --- a/egomimic/hydra_configs/model/hpt_bc_flow_eva.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_flow_eva.yaml @@ -32,7 +32,6 @@ robomimic_model: multitask: false pretrained: false pretrained_checkpoint: "" # TODO - reverse_kl_samples: 8 domains: ["eva_bimanual"] shared_obs_keys: ["front_img_1"] diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml index 2248f562b0..25b467bfe6 100644 --- a/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml @@ -40,9 +40,8 @@ robomimic_model: multitask: false pretrained: false - pretrained_checkpoint: "" # TODO - reverse_kl_samples: 8 - + pretrained_checkpoint: "" + domains: ["human_bimanual"] shared_obs_keys: ["front_img_1"] diff --git a/egomimic/hydra_configs/model/hpt_bc_keypoints_base.yaml b/egomimic/hydra_configs/model/hpt_bc_keypoints_base.yaml index 2112c52b55..02bfbf7ab3 100644 --- a/egomimic/hydra_configs/model/hpt_bc_keypoints_base.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_keypoints_base.yaml @@ -8,12 +8,10 @@ robomimic_model: dims: human_bimanual: {proprio: 144, action: 144} - ac_keys: human_bimanual: "actions_keypoints" diffusion: true - reverse_kl_samples: 8 trunk: embed_dim: 256 diff --git a/egomimic/hydra_configs/model/hpt_cotrain_enc_dec_base.yaml b/egomimic/hydra_configs/model/hpt_cotrain_enc_dec_base.yaml index 1c10c3480e..7edd9e4904 100644 --- a/egomimic/hydra_configs/model/hpt_cotrain_enc_dec_base.yaml +++ b/egomimic/hydra_configs/model/hpt_cotrain_enc_dec_base.yaml @@ -13,7 +13,6 @@ robomimic_model: human_bimanual: "actions_cartesian" eva_bimanual: "actions_joints" - reverse_kl_samples: 8 6dof: true trunk: diff --git a/egomimic/hydra_configs/train_zarr_cartesian.yaml b/egomimic/hydra_configs/train_zarr_cartesian.yaml index a23d9fe03a..4676f6f728 100644 --- a/egomimic/hydra_configs/train_zarr_cartesian.yaml +++ b/egomimic/hydra_configs/train_zarr_cartesian.yaml @@ -28,6 +28,17 @@ launch_params: seed: 42 +# Optional pre-fit baseline validation (metrics at epoch 0; the overlay video +# still follows evaluator.viz_every_n_epochs). Off by default; enable with +# `val_at_start=true`. Skipped automatically on requeues/checkpoint resumes. +val_at_start: false + +# Second eval head on TRAIN data at every validation (metrics `train_viz/`, +# videos under `videos_train_viz/`). Wraps the canonical evaluator and reuses +# the train split unless the data config declares `train_viz_datasets` / +# the top-level config a `train_viz_evaluator`. `train_viz=false` turns it off. +train_viz: true + # Normalization and norm-stat cache settings (override per-run as needed). norm_stats: norm_mode: quantile diff --git a/egomimic/hydra_configs/train_zarr_keypoint_wrist.yaml b/egomimic/hydra_configs/train_zarr_keypoint_wrist.yaml index e45cfeea0a..b1571252c1 100644 --- a/egomimic/hydra_configs/train_zarr_keypoint_wrist.yaml +++ b/egomimic/hydra_configs/train_zarr_keypoint_wrist.yaml @@ -2,5 +2,5 @@ defaults: - train_zarr_cartesian - override model: hpt_bc_keypoints_base - override data: aria_keypoints_wrist - - override evaluator/viz@evaluator.viz_func: keypoints_wrist + - override evaluator: eval_hpt_keypoints_wrist - _self_ diff --git a/egomimic/hydra_configs/train_zarr_keypoints.yaml b/egomimic/hydra_configs/train_zarr_keypoints.yaml index 170bffd287..3cbe07a386 100644 --- a/egomimic/hydra_configs/train_zarr_keypoints.yaml +++ b/egomimic/hydra_configs/train_zarr_keypoints.yaml @@ -2,5 +2,5 @@ defaults: - train_zarr_cartesian - override model: hpt_bc_keypoints_base - override data: aria_keypoints - - override evaluator/viz@evaluator.viz_func: keypoints + - override evaluator: eval_hpt_keypoints_head - _self_ diff --git a/egomimic/pl_utils/pl_data_utils.py b/egomimic/pl_utils/pl_data_utils.py index e16a1a1b13..e4840115db 100644 --- a/egomimic/pl_utils/pl_data_utils.py +++ b/egomimic/pl_utils/pl_data_utils.py @@ -20,6 +20,12 @@ def __init__( valid_datasets: dict, train_dataloader_params: dict, valid_dataloader_params: dict, + train_viz_datasets: dict | None = None, + train_viz_dataloader_params: dict | None = None, + unseen_op_valid_datasets: dict | None = None, + unseen_op_valid_dataloader_params: dict | None = None, + valid_prefix: str | None = None, + held_out_operators: list | None = None, ): """ Args: @@ -27,10 +33,21 @@ def __init__( valid_datasets: dictionary of valid datasets train_dataloader_params: dictionary of train dataloader parameters valid_dataloader_params: dictionary of valid dataloader parameters + train_viz_datasets: optional dict of datasets iterated as a second + val loader, evaluated by a prefixed copy of the canonical + evaluator (metrics ``train_viz/``). + train_viz_dataloader_params: per-dataset DataLoader kwargs for the + train_viz loader. + unseen_op_valid_datasets: optional dict of datasets iterated as a + third val loader (metrics ``unseen_op_valid/``). + unseen_op_valid_dataloader_params: per-dataset DataLoader kwargs + for the unseen_op_valid loader. + valid_prefix, held_out_operators: config-only data-config keys that + hydra.instantiate(cfg.data) forwards; ignored here. Tokenization (sampling a prompt from per-sample annotation lists, splicing in embodiment / control-mode / proprio blocks, and running - the HF tokenizer) lives on the algo side now — see + the HF tokenizer) lives on the algo side now; see ``PI.process_batch_for_training``. The collate here only stacks tensors and preserves variable-length list-valued keys (e.g. raw ``annotations``) so the algo can consume them downstream. @@ -43,6 +60,14 @@ def __init__( self.valid_datasets = {k: v for k, v in valid_datasets.items() if v is not None} self.train_dataloader_params = train_dataloader_params self.valid_dataloader_params = valid_dataloader_params + self.train_viz_datasets = { + k: v for k, v in (train_viz_datasets or {}).items() if v is not None + } + self.train_viz_dataloader_params = train_viz_dataloader_params or {} + self.unseen_op_valid_datasets = { + k: v for k, v in (unseen_op_valid_datasets or {}).items() if v is not None + } + self.unseen_op_valid_dataloader_params = unseen_op_valid_dataloader_params or {} self.collate_fn = annotation_collate def train_dataloader(self): @@ -62,13 +87,13 @@ def train_dataloader(self): return CombinedLoader(iterables, "max_size_cycle") - def val_dataloader(self): + def _build_val_style_loader(self, datasets: dict, params: dict, kind: str): iterables = dict() - for dataset_name, dataset in self.valid_datasets.items(): - dataset_params = self.valid_dataloader_params.get(dataset_name) + for dataset_name, dataset in datasets.items(): + dataset_params = params.get(dataset_name) if dataset_params is None or len(dataset_params) == 0: raise ValueError( - f"No dataloader params found for dataset {dataset_name}. Please add {dataset_name} into your data config valid_dataloader_params." + f"No dataloader params found for dataset {dataset_name}. Please add {dataset_name} into your data config {kind}_dataloader_params." ) dataset_params = dict(dataset_params) shuffle = dataset_params.pop("shuffle", False) @@ -78,9 +103,37 @@ def val_dataloader(self): collate_fn=self.collate_fn, **dataset_params, ) - return CombinedLoader(iterables, "max_size_cycle") + def val_loader_names(self) -> list[str]: + """Names of the val loaders in ``val_dataloader()`` order; position i is + Lightning's ``dataloader_idx`` i. ModelWrapper dispatches on these.""" + names = ["valid"] + if self.train_viz_datasets: + names.append("train_viz") + if self.unseen_op_valid_datasets: + names.append("unseen_op_valid") + return names + + def val_dataloader(self): + sources = { + "valid": (self.valid_datasets, self.valid_dataloader_params), + "train_viz": (self.train_viz_datasets, self.train_viz_dataloader_params), + "unseen_op_valid": ( + self.unseen_op_valid_datasets, + self.unseen_op_valid_dataloader_params, + ), + } + names = self.val_loader_names() + loaders = [ + self._build_val_style_loader(*sources[name], kind=name) for name in names + ] + if len(loaders) == 1: + 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 + def _extract_list_keys(batch): """Pop all list-valued keys from *batch* samples and return them separately. diff --git a/egomimic/pl_utils/pl_model.py b/egomimic/pl_utils/pl_model.py index b5657baf7c..c1b9034153 100644 --- a/egomimic/pl_utils/pl_model.py +++ b/egomimic/pl_utils/pl_model.py @@ -75,6 +75,16 @@ def __init__( self.epoch_memory_stats = [] # Store memory stats per epoch self.evaluator = evaluator + # Optional second evaluator that runs against a train-sampled + # dataloader during the same validation pass (dataloader_idx=1). + # See egomimic/eval/eval_train_viz.py. + self.train_viz_evaluator = None + # Optional third evaluator for the unseen_op_valid loader (held-out + # operators in the opsplit data configs). + self.unseen_op_valid_evaluator = None + # Val loader names in dataloader_idx order; trainHydra sets + # MultiDataModuleWrapper.val_loader_names(). + self.val_loader_names = ["valid"] @staticmethod def _as_config(cfg): @@ -205,30 +215,56 @@ def on_before_optimizer_step(self, optimizer): sync_dist=True, ) + def _val_heads(self) -> dict: + return { + "valid": self.evaluator, + "train_viz": self.train_viz_evaluator, + "unseen_op_valid": self.unseen_op_valid_evaluator, + } + def on_validation_start(self): - if self.evaluator is None: + heads = [h for h in self._val_heads().values() if h is not None] + if not heads: return self.model.device = self.device - self.evaluator.on_validation_start() + for head in heads: + head.on_validation_start() def validation_step(self, batch, batch_idx, dataloader_idx=0): """ Run a validation step on the batch, and save that batch of images into the val_image_buffer. Once the buffer hits 1000 images, save that as a 30fps video using torchvision.io.write_video. + + ``dataloader_idx`` is populated by Lightning when ``val_dataloader()`` + returns a list; ``val_loader_names[dataloader_idx]`` picks the head + (valid -> evaluator, train_viz -> train_viz_evaluator, unseen_op_valid -> + unseen_op_valid_evaluator). """ - if self.evaluator is None: + names = self.val_loader_names + name = names[dataloader_idx] if dataloader_idx < len(names) else None + active = self._val_heads().get(name) + if active is None: return + # When val_dataloader returns a list of CombinedLoaders (valid + + # train_viz), Lightning wraps it in an outer sequential CombinedLoader. + # The outer iterator calls next() on each inner CombinedLoader, which + # itself yields a (batch_dict, batch_idx, dataloader_idx) triple; that + # triple lands here as `batch`. Unwrap to recover the dict. + if isinstance(batch, tuple) and len(batch) == 3 and isinstance(batch[0], dict): + batch = batch[0] batch = self.model.process_batch_for_training(batch) print( - f"[VAL_STEP] rank={self.global_rank}, batch_idx={batch_idx}", + f"[VAL_STEP] rank={self.global_rank}, batch_idx={batch_idx}, " + f"dataloader_idx={dataloader_idx}", flush=True, ) - self.evaluator.on_validation_step(batch, batch_idx, dataloader_idx) + active.on_validation_step(batch, batch_idx, dataloader_idx) def on_validation_end(self): print(f"[ON_VALIDATION_END] rank={self.global_rank}", flush=True) - if self.evaluator is not None: - self.evaluator.on_validation_end() + for head in self._val_heads().values(): + if head is not None: + head.on_validation_end() print( f"Rank {self.global_rank} on validation end, waiting for all ranks to synchronize", diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index b76d648a61..6f3be5a888 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -18,6 +18,7 @@ import egomimic.utils.hydra_resolvers # noqa: F401 -- registers OmegaConf resolvers from egomimic.eval.eval import Eval +from egomimic.eval.eval_train_viz import PrefixedEvalVideo from egomimic.pl_utils.pl_model import ModelWrapper from egomimic.rldb.embodiment.embodiment import get_embodiment_id from egomimic.rldb.resolve_memo import resolve_once @@ -120,7 +121,11 @@ def _weights_from_checkpoint(cfg: DictConfig) -> bool: return bool(ckpt_path) and os.path.isfile(ckpt_path) and not cfg.get("pretrained") -def _log_dataset_frame_counts(train_datasets: dict, valid_datasets: dict) -> None: +def _log_dataset_frame_counts( + train_datasets: dict, + valid_datasets: dict, + unseen_op_valid_datasets: dict | None = None, +) -> None: rows = [] for name, ds in train_datasets.items(): rows.append(("train", name, len(ds))) @@ -134,6 +139,8 @@ def _log_dataset_frame_counts(train_datasets: dict, valid_datasets: dict) -> Non rows.append( ("TOTAL", "(valid)", sum(len(ds) for ds in valid_datasets.values())) ) + for name, ds in (unseen_op_valid_datasets or {}).items(): + rows.append(("unseen_op_valid", name, len(ds))) table = tabulate( rows, headers=["Split", "Dataset", "Frames"], @@ -196,6 +203,93 @@ def load_checkpoint( return pl_load(path, map_location=map_location, weights_only=weights_only) +def _resolve_mode(cfg: DictConfig) -> str: + """``mode`` key, else the legacy ``train`` / ``eval`` booleans.""" + if cfg.get("mode") is not None: + return cfg.mode + if cfg.get("train", False): + return "train" + if cfg.get("eval", False): + return "eval" + raise ValueError("Config must specify either `mode` or `train`/`eval` booleans") + + +def _train_viz_enabled(cfg: DictConfig) -> bool: + """The train_viz head (metrics + video on TRAIN data each validation) is on + by default for training runs with an evaluator; ``train_viz=false`` turns + it off, including the heads a data/top-level config declares itself.""" + return ( + bool(cfg.get("train_viz", True)) + and _resolve_mode(cfg) == "train" + and cfg.get("evaluator") is not None + ) + + +def _train_viz_datasets(cfg: DictConfig, train_datasets: dict, instantiate): + """Datasets for the train_viz loader and, when they are derived here, its + loader params (None = leave ``data.train_viz_dataloader_params`` to Hydra). + + A data config's own ``train_viz_datasets`` win. Otherwise the train split + itself is reused (same dataset objects, so no second resolve) through an + unshuffled loader with the valid loader's params.""" + if not _train_viz_enabled(cfg): + return {}, None + explicit = cfg.data.get("train_viz_datasets") + if explicit is not None: + return { + name: None if node is None else instantiate(node, dataset_name=name) + for name, node in explicit.items() + }, None + valid_params = cfg.data.get("valid_dataloader_params") or {} + params = { + name: {**OmegaConf.to_container(valid_params[name]), "shuffle": False} + for name in train_datasets + if name in valid_params + } + return {name: train_datasets[name] for name in params}, params + + +def _build_train_viz_evaluator(cfg: DictConfig): + """``train_viz_evaluator`` if configured, else a PrefixedEvalVideo around a + second instance of the canonical evaluator (own frame buffers).""" + if not _train_viz_enabled(cfg): + return None + if cfg.get("train_viz_evaluator") is not None: + return hydra.utils.instantiate(cfg.train_viz_evaluator) + return PrefixedEvalVideo(hydra.utils.instantiate(cfg.evaluator)) + + +def _unseen_op_valid_datasets(cfg: DictConfig, instantiate) -> dict: + """Datasets for the third (unseen_op_valid) val loader: a data config's own + ``unseen_op_valid_datasets``, in training runs with an evaluator only (eval + mode validates ``valid_datasets`` alone).""" + explicit = cfg.data.get("unseen_op_valid_datasets") + if ( + explicit is None + or _resolve_mode(cfg) != "train" + or cfg.get("evaluator") is None + ): + return {} + return { + name: None if node is None else instantiate(node, dataset_name=name) + for name, node in explicit.items() + } + + +def _build_unseen_op_valid_evaluator(cfg: DictConfig): + """The canonical evaluator (fresh instance, own frame buffers) wrapped to + log ``unseen_op_valid/`` and write ``videos_unseen_op_valid/``.""" + return PrefixedEvalVideo( + hydra.utils.instantiate(cfg.evaluator), prefix="unseen_op_valid" + ) + + +def _attach(eval_obj, trainer: Trainer, model: LightningModule): + eval_obj.trainer = trainer + eval_obj.model = model.model + return eval_obj + + @task_wrapper def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Trains the model. Can additionally evaluate on a testset, using best weights obtained during @@ -217,8 +311,9 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: load_env() - # One SQL pull / path resolution per dataset spec across train, valid and - # the norm-stat copies; dropped on exit so nothing outlives this run. + # One SQL pull / path resolution per dataset spec across train, valid, + # train_viz and the norm-stat copies; dropped on exit so nothing outlives + # this run. with resolve_once(): train_datasets = {} for dataset_name in cfg.data.train_datasets: @@ -232,12 +327,30 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: cfg.data.valid_datasets[dataset_name], dataset_name=dataset_name ) + train_viz_datasets, train_viz_params = _train_viz_datasets( + cfg, train_datasets, instantiate=_instantiate_dataset + ) + unseen_op_valid_datasets = _unseen_op_valid_datasets( + cfg, instantiate=_instantiate_dataset + ) + log.info(f"Instantiating datamodule <{cfg.data._target_}>") assert ( "MultiDataModuleWrapper" in cfg.data._target_ ), "cfg.data._target_ must be 'MultiDataModuleWrapper'" + datamodule_kwargs = dict( + train_datasets=train_datasets, valid_datasets=valid_datasets + ) + # Always passed, so train_viz=false also drops a data config's own + # train_viz_datasets. + datamodule_kwargs["train_viz_datasets"] = train_viz_datasets + if train_viz_params is not None: + datamodule_kwargs["train_viz_dataloader_params"] = train_viz_params + # Always passed too: hydra would otherwise instantiate the raw + # config nodes itself (eval mode included). + datamodule_kwargs["unseen_op_valid_datasets"] = unseen_op_valid_datasets datamodule: LightningDataModule = hydra.utils.instantiate( - cfg.data, train_datasets=train_datasets, valid_datasets=valid_datasets + cfg.data, **datamodule_kwargs ) # Stats-only MultiDataset (no graph of its own; explicitly populated from @@ -345,6 +458,10 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: ds.set_norm_stats_from(norm_stats) for ds in datamodule.valid_datasets.values(): ds.set_norm_stats_from(norm_stats) + for ds in datamodule.train_viz_datasets.values(): + ds.set_norm_stats_from(norm_stats) + for ds in datamodule.unseen_op_valid_datasets.values(): + ds.set_norm_stats_from(norm_stats) log.info(f"Instantiating model <{cfg.model._target_}>") model: LightningModule = ModelWrapper( @@ -353,20 +470,16 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: scheduler_interval=cfg.model.get("scheduler_interval", "step"), ) - _log_dataset_frame_counts(datamodule.train_datasets, datamodule.valid_datasets) + _log_dataset_frame_counts( + datamodule.train_datasets, + datamodule.valid_datasets, + datamodule.unseen_op_valid_datasets, + ) log.info("Instantiating callbacks...") callbacks: List[Callback] = instantiate_callbacks(cfg.get("callbacks")) - # Resolve mode: support both new `mode` key and legacy `train`/`eval` booleans - if cfg.get("mode") is not None: - mode = cfg.mode - elif cfg.get("train", False): - mode = "train" - elif cfg.get("eval", False): - mode = "eval" - else: - raise ValueError("Config must specify either `mode` or `train`/`eval` booleans") + mode = _resolve_mode(cfg) # In eval mode, apply trainer overrides from the eval object and disable logger if mode == "eval": @@ -416,9 +529,33 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: if mode == "train": if cfg.get("evaluator") is not None: eval_obj: Eval = hydra.utils.instantiate(cfg.evaluator) - eval_obj.trainer = trainer - eval_obj.model = model.model - model.evaluator = eval_obj + # data.valid_prefix (e.g. seen_op_valid in the opsplit configs) + # renames the canonical head: `/Valid/...` metrics and + # `videos_/` instead of `Valid/...` and `videos/`. + valid_prefix = cfg.data.get("valid_prefix") + if valid_prefix: + eval_obj = PrefixedEvalVideo(eval_obj, prefix=valid_prefix) + model.evaluator = _attach(eval_obj, trainer, model) + train_viz_eval_obj = ( + _build_train_viz_evaluator(cfg) if datamodule.train_viz_datasets else None + ) + if train_viz_eval_obj is not None: + model.train_viz_evaluator = _attach(train_viz_eval_obj, trainer, model) + if datamodule.unseen_op_valid_datasets: + model.unseen_op_valid_evaluator = _attach( + _build_unseen_op_valid_evaluator(cfg), trainer, model + ) + model.val_loader_names = datamodule.val_loader_names() + # Pre-fit baseline val. Skipped on requeues AND checkpoint resumes: + # trainer.validate here runs BEFORE fit restores ckpt_path weights, so + # on a resume it would score the un-resumed base model. + # (_prepare_checkpoint_resume already folded a requeue into ckpt_path.) + if cfg.get("val_at_start", False) and not cfg.get("ckpt_path"): + log.info( + "val_at_start: running validation at epoch 0 (pre-fit baseline; " + "the overlay video follows evaluator.viz_every_n_epochs)" + ) + trainer.validate(model=model, datamodule=datamodule) log.info("Starting training!") trainer.fit( model=model, diff --git a/egomimic/utils/metrics.py b/egomimic/utils/metrics.py index aa6c04756c..8b1ace86c0 100644 --- a/egomimic/utils/metrics.py +++ b/egomimic/utils/metrics.py @@ -1,33 +1,7 @@ -"""Distribution-distance metrics used by evaluation. - -Moved out of egomimicUtils; code unchanged. -""" - -import math +"""Distribution-distance metrics used by evaluation.""" import torch -# ---- moved from egomimicUtils.py (code unchanged) ---- - - -def reverse_kl_from_samples(pred_samples, targets): - M, B, T, D = pred_samples.shape - - TD = T * D - const = -0.5 * TD * math.log(2.0 * math.pi) - - A = pred_samples.permute(1, 0, 2, 3).reshape(B, M, TD) # (B,M,TD) - MU = targets.reshape(B, 1, TD) # (B,1,TD) - - d2 = torch.cdist(A, A).pow(2) # (B,M,M) - log_q_each = torch.logsumexp(const - 0.5 * d2, dim=-1) - math.log(M) # (B,M) - - d2p = ((A - MU) ** 2).sum(dim=-1) # (B,M) - log_p_each = const - 0.5 * d2p # (B,M) - - rkl_each = (log_q_each - log_p_each).mean(dim=-1) # (B,) - return rkl_each.mean() - def frechet_gaussian_over_time( pred: torch.Tensor, @@ -112,3 +86,39 @@ def frechet_gaussian_over_time( "max": dist.max().item(), } return dist + + +def dtw_distance(pred: torch.Tensor, tgt: torch.Tensor, normalize: bool = True): + """Batched dynamic-time-warping distance between (B, T1, D) and (B, T2, D) + trajectories. + + Local cost is the euclidean distance between D-dim frames; the classic + DTW recurrence runs as a wavefront over the T1+T2-1 anti-diagonals of the + cost grid so the whole batch advances in vectorized steps (no per-cell + python loop). Unlike paired MSE, DTW forgives temporal misalignment: a + correct trajectory executed slightly early/late scores near zero. + + Args: + pred/tgt: (B, T, D) action chunks (any device). + normalize: divide the accumulated path cost by (T1 + T2), giving a + per-step average frame distance comparable across chunk lengths. + Returns: + (B,) distances. + """ + B, T1, _ = pred.shape + _, T2, _ = tgt.shape + C = torch.cdist(pred.float(), tgt.float()) # (B, T1, T2) + acc = torch.full((B, T1 + 1, T2 + 1), float("inf"), device=C.device, dtype=C.dtype) + acc[:, 0, 0] = 0.0 + for k in range(2, T1 + T2 + 1): + i = torch.arange(max(1, k - T2), min(T1, k - 1) + 1, device=C.device) + j = k - i + prev = torch.minimum( + torch.minimum(acc[:, i - 1, j], acc[:, i, j - 1]), + acc[:, i - 1, j - 1], + ) + acc[:, i, j] = C[:, i - 1, j - 1] + prev + out = acc[:, T1, T2] + if normalize: + out = out / (T1 + T2) + return out diff --git a/tests/unit/test_action_metrics.py b/tests/unit/test_action_metrics.py new file mode 100644 index 0000000000..35e55a639c --- /dev/null +++ b/tests/unit/test_action_metrics.py @@ -0,0 +1,378 @@ +"""egomimic.eval.action_metrics, the trajectory metrics it uses, and the PI +evaluator's metric set.""" + +import math + +import numpy as np +import pytest +import torch + +from egomimic.eval.action_metrics import ( + cartesian_metrics, + keypoint_metrics, + layout_metrics, +) +from egomimic.utils.pose_utils import _ypr_to_rot6d + + +def test_layout_metrics_dispatch_on_the_action_key_not_just_the_width(): + pred, gt = torch.zeros(2, 5, 14), torch.ones(2, 5, 14) + assert layout_metrics(pred, gt, "p", "actions_joints") == {} + cart = layout_metrics(pred, gt, "p", "actions_cartesian") + assert "p_xyz_paired_mse_avg" in cart + assert "p_ypr_paired_mse_avg" in cart and "p_rot6d_paired_mse_avg" not in cart + six = layout_metrics( + torch.zeros(2, 5, 20), torch.ones(2, 5, 20), "p", "actions_cartesian" + ) + assert "p_rot6d_paired_mse_avg" in six + kp = layout_metrics( + torch.zeros(2, 5, 144), torch.ones(2, 5, 144), "p", "actions_keypoints" + ) + assert "p_kp_l2_avg" in kp + + +def test_pi_eval_falls_back_to_plain_mse_for_other_action_keys(): + from egomimic.eval.eval_pi import PIEvalVideo + from egomimic.rldb.embodiment.embodiment import EMBODIMENT, get_embodiment + + eid = EMBODIMENT.EVA_BIMANUAL.value + name = get_embodiment(eid).lower() + ac_key = "actions_joints" + gt = torch.zeros(2, 5, 14) + pred = torch.ones(2, 5, 14) + + class _NormStats: + def unnormalize(self, batch, embodiment_id): + return batch + + class _Algo: + ac_keys = {eid: ac_key} + norm_stats = _NormStats() + + def forward_eval(self, batch): + return {f"{name}_{ac_key}": pred} + + ev = PIEvalVideo(viz_func={}, transform_lists={}, viz_every_n_epochs=1) + ev.model = _Algo() + batch = {eid: {ac_key: gt, "embodiment": torch.tensor([eid])}} + metrics, _ = ev.compute_metrics_and_viz(batch, do_viz=False) + prefix = f"Valid/{name}_{ac_key}_" + assert set(metrics) == {prefix + "paired_mse_avg", prefix + "final_mse_avg"} + assert float(metrics[prefix + "paired_mse_avg"]) == 1.0 + + +def test_frechet_helper(): + from egomimic.utils.metrics import frechet_gaussian_over_time + + B, T, D = 4, 10, 6 + torch.manual_seed(0) + pred = torch.randn(B, T, D) + fd = frechet_gaussian_over_time(pred, pred.clone()) + assert fd.shape == (B,) + assert fd.max() < 1e-2, "identical distributions must score ~0" + fd_shift = frechet_gaussian_over_time(pred + 5.0, pred) + assert (fd_shift > fd).all(), "mean shift must increase the distance" + + +def test_dtw_distance_matches_bruteforce_and_tolerates_shift(): + from egomimic.utils.metrics import dtw_distance + + def _dtw_ref(x, y): + t1, t2 = len(x), len(y) + cost = np.linalg.norm(x[:, None, :] - y[None, :, :], axis=-1) + acc = np.full((t1 + 1, t2 + 1), np.inf) + acc[0, 0] = 0.0 + for i in range(1, t1 + 1): + for j in range(1, t2 + 1): + acc[i, j] = cost[i - 1, j - 1] + min( + acc[i - 1, j], acc[i, j - 1], acc[i - 1, j - 1] + ) + return acc[t1, t2] + + torch.manual_seed(3) + pred = torch.randn(3, 9, 4) + tgt = torch.randn(3, 9, 4) + got = dtw_distance(pred, tgt, normalize=False) + for b in range(3): + ref = _dtw_ref(pred[b].numpy(), tgt[b].numpy()) + assert abs(got[b].item() - ref) < 1e-4, (b, got[b].item(), ref) + + # identical trajectories -> 0 + assert dtw_distance(pred, pred.clone()).max() < 1e-6 + + # a time-shifted copy of the same smooth trajectory: DTW must forgive the + # shift (score far below the paired per-step distance of the shifted pair) + t = torch.linspace(0, 6.28, 40) + traj = torch.stack([torch.sin(t), torch.cos(t)], dim=-1)[None] # (1, 40, 2) + shifted = torch.roll(traj, shifts=3, dims=1) + dtw_shift = dtw_distance(traj, shifted, normalize=False).item() + paired = (traj - shifted).norm(dim=-1).sum().item() + assert dtw_shift < 0.5 * paired, (dtw_shift, paired) + + +def test_split_mse_is_stateless_and_matches_manual(): + from egomimic.eval.action_metrics import _paired_mse, _split_mse + + rng = np.random.default_rng(21) + pred = torch.from_numpy(rng.normal(size=(4, 10, 18))).float() + gt = torch.from_numpy(rng.normal(size=(4, 10, 18))).float() + xyz_idx = [0, 1, 2, 9, 10, 11] + rot_idx = [i for i in range(18) if i not in xyz_idx] + xyz, rot = _split_mse(pred, gt) + torch.testing.assert_close( + xyz, (pred[..., xyz_idx] - gt[..., xyz_idx]).pow(2).mean() + ) + torch.testing.assert_close( + rot, (pred[..., rot_idx] - gt[..., rot_idx]).pow(2).mean() + ) + # stateless: a second, unrelated call is unaffected by the first + a, b = torch.zeros(2, 3, 18), torch.ones(2, 3, 18) + torch.testing.assert_close(_split_mse(a, b)[0], torch.tensor(1.0)) + torch.testing.assert_close(_paired_mse(a, b), torch.tensor(1.0)) + assert _split_mse(torch.zeros(2, 3, 7), torch.zeros(2, 3, 7)) == (None, None) + + +def test_rot_geodesic_error_matches_angle_and_survives_gimbal_lock(): + from scipy.spatial.transform import Rotation as R + + from egomimic.eval.action_metrics import _rot_geodesic_error + + rng = np.random.default_rng(22) + ypr = rng.uniform(-1.0, 1.0, size=(6, 5, 3)) + theta = 0.3 + # rotate every pose by theta about a random axis -> geodesic error == theta + axes = rng.normal(size=(6, 5, 3)) + axes /= np.linalg.norm(axes, axis=-1, keepdims=True) + Rg = R.from_euler("ZYX", ypr.reshape(-1, 3)) + Rp = R.from_rotvec(theta * axes.reshape(-1, 3)) * Rg + ypr_p = Rp.as_euler("ZYX").reshape(6, 5, 3) + + # 12-dim ypr layout (both arms the same pose) + gt12 = torch.from_numpy(np.concatenate([ypr, ypr], -1)).float() + gt12 = torch.cat( + [torch.zeros(6, 5, 3), gt12[..., :3], torch.zeros(6, 5, 3), gt12[..., 3:]], -1 + ) + pr12 = torch.from_numpy(np.concatenate([ypr_p, ypr_p], -1)).float() + pr12 = torch.cat( + [torch.zeros(6, 5, 3), pr12[..., :3], torch.zeros(6, 5, 3), pr12[..., 3:]], -1 + ) + assert abs(_rot_geodesic_error(pr12, gt12).item() - theta) < 1e-4 + assert _rot_geodesic_error(gt12, gt12).item() < 1e-5 + + # 18-dim 6D layout, same rotations -> same answer + def to18(y): + six = _ypr_to_rot6d(y) + arm = np.concatenate([np.zeros(y.shape[:-1] + (3,)), six], -1) + return torch.from_numpy(np.concatenate([arm, arm], -1)).float() + + assert abs(_rot_geodesic_error(to18(ypr_p), to18(ypr)).item() - theta) < 1e-4 + assert _rot_geodesic_error(torch.zeros(2, 3, 7), torch.zeros(2, 3, 7)) is None + + # Gimbal lock: pitch ~ pi/2, pred = gt rotated 0.004 rad about local y. + # A ypr MSE (even wrap-aware) explodes because yaw and roll trade off; + # the geodesic error reports the true 0.004. + gt = np.array([0.3, np.pi / 2 - 0.002, 0.2]) + Rgt = R.from_euler("ZYX", gt) + Rpr = Rgt * R.from_rotvec([0.0, 0.004, 0.0]) + pr = Rpr.as_euler("ZYX") + v_gt = torch.tensor([[0, 0, 0, *gt, 0, 0, 0, *gt]], dtype=torch.float32) + v_pr = torch.tensor([[0, 0, 0, *pr, 0, 0, 0, *pr]], dtype=torch.float32) + geo = _rot_geodesic_error(v_pr, v_gt).item() + assert abs(geo - 0.004) < 1e-3, geo + euler_mse = (v_pr - v_gt).pow(2).mean().item() + assert euler_mse > 0.1, euler_mse # the Euler metric is fooled here + + +# ------------------------------------------------ PIEvalVideo metric contract + + +def _rand_6d_chunks(seed, B=3, T=6): + """(pred, gt) native 18D human chunks with a small, known perturbation, + plus the per-(batch, arm) rotation matrices so a test can re-express both + in another frame. 6D columns are built directly from rotation matrices + (col 0 | col 1), matching ``_reconstruct_R_from_cols``.""" + from scipy.spatial.transform import Rotation as R + + rng = np.random.default_rng(seed) + xyz_g = rng.uniform(-1.0, 1.0, size=(B, T, 2, 3)) + R_g = R.random(B * T * 2, random_state=seed).as_matrix().reshape(B, T, 2, 3, 3) + xyz_p = xyz_g + 0.05 * rng.normal(size=xyz_g.shape) + R_p = ( + ( + R.from_rotvec(0.1 * rng.normal(size=(B * T * 2, 3))) + * R.from_matrix(R_g.reshape(-1, 3, 3)) + ) + .as_matrix() + .reshape(B, T, 2, 3, 3) + ) + return (xyz_p, R_p), (xyz_g, R_g) + + +def _chunk18(xyz, Rm): + six = np.concatenate([Rm[..., :, 0], Rm[..., :, 1]], -1) # (B,T,2,6) + arm = np.concatenate([xyz, six], -1) # (B,T,2,9) + return torch.from_numpy(arm.reshape(arm.shape[0], arm.shape[1], 18)).float() + + +def _run_pi_eval(pred, gt, ac_key="actions_cartesian"): + from egomimic.eval.eval_pi import PIEvalVideo + from egomimic.rldb.embodiment.embodiment import EMBODIMENT, get_embodiment + + eid = EMBODIMENT.HUMAN_BIMANUAL.value + name = get_embodiment(eid).lower() + + class _NormStats: + def unnormalize(self, batch, embodiment_id): + return batch + + class _Algo: + ac_keys = {eid: ac_key} + norm_stats = _NormStats() + + def forward_eval(self, batch): + return {f"{name}_loss": torch.tensor(0.5), f"{name}_{ac_key}": pred} + + ev = PIEvalVideo(viz_func={}, transform_lists={}, viz_every_n_epochs=1) + ev.model = _Algo() + batch = {eid: {ac_key: gt, "embodiment": torch.tensor([eid])}} + metrics, images = ev.compute_metrics_and_viz(batch, do_viz=False) + return f"Valid/{name}_{ac_key}_", f"Valid/{name}_loss", metrics, images + + +def test_pi_eval_metric_set_and_identity(): + (xyz_p, R_p), (xyz_g, R_g) = _rand_6d_chunks(40) + gt = _chunk18(xyz_g, R_g) + prefix, loss_key, metrics, images = _run_pi_eval(gt.clone(), gt) + expected = { + prefix + k + for k in ( + "xyz_paired_mse_avg", + "xyz_final_mse_avg", + "rot6d_paired_mse_avg", + "rot_err_deg_avg", + "rot_err_deg_final", + "xyz_dtw_avg", + "xyz_frechet_gauss_avg", + "xyz_l2_early", + "xyz_l2_mid", + "xyz_l2_late", + ) + } | {loss_key, "Valid/action_loss"} + assert set(metrics) == expected, sorted(metrics) + assert images == {}, "do_viz=False must not render" + assert float(metrics[prefix + "xyz_paired_mse_avg"]) < 1e-10 + assert float(metrics[prefix + "rot_err_deg_avg"]) < 1e-3 + assert float(metrics[prefix + "xyz_dtw_avg"]) < 1e-6 + assert float(metrics[prefix + "xyz_l2_late"]) < 1e-5 + assert float(metrics[prefix + "xyz_frechet_gauss_avg"]) < 1e-2 + + +def test_pi_eval_metrics_are_invariant_to_a_shared_rigid_transform(): + # The evaluator logs no cam-frame twins of the native metrics: the revert + # applies the SAME rigid transform (the observed wrist pose) to prediction + # and gt, and every kept metric is invariant to that. Re-express pred and + # gt in a random frame per (batch, arm) and check. + from scipy.spatial.transform import Rotation as R + + (xyz_p, R_p), (xyz_g, R_g) = _rand_6d_chunks(41) + B = xyz_g.shape[0] + R_w = R.random(B * 2, random_state=7).as_matrix().reshape(B, 1, 2, 3, 3) + t_w = np.random.default_rng(8).uniform(-2, 2, size=(B, 1, 2, 3)) + + def move(xyz, Rm): + return np.einsum("btaij,btaj->btai", R_w, xyz) + t_w, R_w @ Rm + + prefix, _, native, _ = _run_pi_eval(_chunk18(xyz_p, R_p), _chunk18(xyz_g, R_g)) + prefix, _, moved, _ = _run_pi_eval( + _chunk18(*move(xyz_p, R_p)), _chunk18(*move(xyz_g, R_g)) + ) + for key, atol in ( + ("xyz_paired_mse_avg", 1e-5), + ("xyz_final_mse_avg", 1e-5), + ("rot_err_deg_avg", 1e-3), + ("rot_err_deg_final", 1e-3), + ("xyz_dtw_avg", 1e-4), + ("xyz_frechet_gauss_avg", 1e-3), + ("xyz_l2_early", 1e-5), + ("xyz_l2_late", 1e-5), + ): + a, b = float(native[prefix + key]), float(moved[prefix + key]) + assert abs(a - b) < atol, (key, a, b) + # sanity: the perturbation is visible in the headline metrics + assert float(native[prefix + "rot_err_deg_avg"]) > 1.0 + assert float(native[prefix + "xyz_paired_mse_avg"]) > 1e-4 + + +def test_pi_eval_keypoint_metrics(): + from scipy.spatial.transform import Rotation as R + + rng = np.random.default_rng(50) + B, T = 2, 5 + wrist_xyz = rng.uniform(-1, 1, size=(B, T, 2, 3)) + ypr = rng.uniform(-1, 1, size=(B, T, 2, 3)) + kp = rng.uniform(-0.1, 0.1, size=(B, T, 2, 63)) + + def chunk(wx, y, k): + hand = np.concatenate([wx, _ypr_to_rot6d(y), k], -1) # (B,T,2,72) + return torch.from_numpy(hand.reshape(B, T, 144)).float() + + gt = chunk(wrist_xyz, ypr, kp) + theta = 0.2 + axes = rng.normal(size=(B * T * 2, 3)) + axes /= np.linalg.norm(axes, axis=-1, keepdims=True) + ypr_p = ( + (R.from_rotvec(theta * axes) * R.from_euler("ZYX", ypr.reshape(-1, 3))) + .as_euler("ZYX") + .reshape(B, T, 2, 3) + ) + pred = chunk(wrist_xyz + 0.1, ypr_p, kp + 0.02) + prefix, loss_key, metrics, _ = _run_pi_eval(pred, gt, ac_key="actions_keypoints") + expected = { + prefix + k + for k in ( + "kp_l2_avg", + "kp_l2_final", + "wrist_xyz_paired_mse_avg", + "wrist_xyz_final_mse_avg", + "wrist_rot_err_deg_avg", + "wrist_rot_err_deg_final", + "kp_dtw_avg", + *( + f"{m}_l2_{seg}" + for m in ("kp", "wrist_xyz") + for seg in ("early", "mid", "late") + ), + ) + } | {loss_key, "Valid/action_loss"} + assert set(metrics) == expected, sorted(metrics) + # every keypoint is off by 0.02 in each axis -> L2 = 0.02 * sqrt(3) + assert abs(float(metrics[prefix + "kp_l2_avg"]) - 0.02 * math.sqrt(3)) < 1e-5 + assert abs(float(metrics[prefix + "wrist_xyz_paired_mse_avg"]) - 0.01) < 1e-6 + assert ( + abs(float(metrics[prefix + "wrist_rot_err_deg_avg"]) - math.degrees(theta)) + < 1e-2 + ) + + +def test_segment_l2_separates_early_from_late_error(): + B, T = 2, 100 + gt = torch.zeros(B, T, 20) + pred = gt.clone() + pred[:, :10, 0] = 0.03 # 3 cm on the left arm x, first 10 % only + m = cartesian_metrics(pred, gt, "p") + # mean over both arms: 3 cm on one of two + assert m["p_xyz_l2_early"].item() == pytest.approx(0.015) + assert m["p_xyz_l2_mid"].item() == pytest.approx(0.0) + assert m["p_xyz_l2_late"].item() == pytest.approx(0.0) + + +def test_segment_l2_on_keypoints(): + B, T = 1, 100 + gt = torch.zeros(B, T, 144) + pred = gt.clone() + pred[:, 50:, 9:72] = 0.02 # every left-hand keypoint off by (2,2,2) cm late + m = keypoint_metrics(pred, gt, "k") + expected = 0.5 * float(np.linalg.norm([0.02] * 3)) # one hand of two + assert m["k_kp_l2_late"].item() == pytest.approx(expected, rel=1e-5) + assert m["k_kp_l2_early"].item() == pytest.approx(0.0) + assert m["k_wrist_xyz_l2_late"].item() == pytest.approx(0.0) diff --git a/tests/unit/test_eva_wristframe_norm_pool.py b/tests/unit/test_eva_wristframe_norm_pool.py index 0ceba0b3ec..fc88a913f7 100644 --- a/tests/unit/test_eva_wristframe_norm_pool.py +++ b/tests/unit/test_eva_wristframe_norm_pool.py @@ -254,3 +254,15 @@ def test_bounds_check_ignores_constant_cells(): assert md._check_bounds({"embodiment": 0, key: arr}, None, 0, "ep") is None arr[1, 0] = 50.0 # corrupt value at a regular cell: still rejected assert md._check_bounds({"embodiment": 0, key: arr}, None, 0, "ep") is not None + + +EVA_WRIST_REVERT = "egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_wristframe_transform_list" + + +@pytest.mark.parametrize("top", ["train_zarr_cartesian", "train_zarr_cartesian_pi"]) +def test_eva_defaults_are_wrist_frame_6d(top): + cfg = _compose(top) + eva = cfg.data.train_datasets.eva_bimanual + assert eva.resolver.transform_list.mode == "cartesian_wristframe_6d" + assert cfg.evaluator.transform_lists.eva_bimanual._target_ == EVA_WRIST_REVERT + assert cfg.norm_stats.pool_horizon is False diff --git a/tests/unit/test_eval_video.py b/tests/unit/test_eval_video.py new file mode 100644 index 0000000000..d4a22a92c5 --- /dev/null +++ b/tests/unit/test_eval_video.py @@ -0,0 +1,80 @@ +"""EvalVideo viz gating / fps and the prefixed evaluator wrapper.""" + +from types import SimpleNamespace + +from egomimic.eval.eval_train_viz import PrefixedEvalVideo +from egomimic.eval.eval_video import EvalVideo + + +class _StubEvalVideo(EvalVideo): + def compute_metrics_and_viz(self, batch, do_viz=True): + raise NotImplementedError + + +def test_video_fps_compensates_for_world_size(): + # Distributed val strides an episode's frames by world_size on each rank; + # playback fps must scale down to keep videos wall-clock real-time. + ev = _StubEvalVideo.__new__(_StubEvalVideo) + for world, expected in [(1, 30), (2, 15), (4, 8), (8, 4)]: + ev.trainer = SimpleNamespace(world_size=world) + assert ev._video_fps() == expected, (world, ev._video_fps()) + ev.trainer = SimpleNamespace() # no world_size attr -> assume 1 + assert ev._video_fps() == 30 + + +def test_viz_gate_follows_lightning_epoch_convention(): + ev = _StubEvalVideo(viz_every_n_epochs=20) + # Lightning validates at current_epoch 19, 39, ... for check_val_every_n_epoch=20 + for epoch, expected in [(0, False), (19, True), (20, False), (39, True)]: + ev.trainer = SimpleNamespace(current_epoch=epoch) + assert ev._should_viz() is expected, epoch + ev.viz_every_n_epochs = 0 + assert ev._should_viz() is False + + +def test_viz_gate_always_renders_the_final_epoch(): + ev = _StubEvalVideo(viz_every_n_epochs=100) + # trainer=debug: 4 epochs, val at current_epoch 1 and 3 -> video only at 3 + for epoch, expected in [(1, False), (3, True)]: + ev.trainer = SimpleNamespace(current_epoch=epoch, max_epochs=4) + assert ev._should_viz() is expected, epoch + # eval mode validates once at epoch 0 with max_epochs=1 + ev.trainer = SimpleNamespace(current_epoch=0, max_epochs=1) + assert ev._should_viz() is True + # unbounded runs (max_epochs None / -1) only follow the interval + for max_epochs in (None, -1): + ev.trainer = SimpleNamespace(current_epoch=0, max_epochs=max_epochs) + assert ev._should_viz() is False + # viz_every_n_epochs=0 still turns video off entirely + ev.viz_every_n_epochs = 0 + ev.trainer = SimpleNamespace(current_epoch=3, max_epochs=4) + assert ev._should_viz() is False + + +def test_train_viz_wrapper_prefixes_and_forwards(): + class _Base(EvalVideo): + def __init__(self): + super().__init__(viz_func={}, transform_lists={}, viz_every_n_epochs=7) + self.seen = None + + def compute_metrics_and_viz(self, batch, do_viz=True): + self.seen = (self.model, do_viz) + return {"Valid/x": 1.0}, {} + + base = _Base() + tv = PrefixedEvalVideo(base) + algo = object() + tv.model = algo # property setter forwards to base too + metrics, _ = tv.compute_metrics_and_viz({}, do_viz=False) + assert set(metrics) == {"train_viz/Valid/x"}, metrics + assert base.seen == (algo, False) + assert tv.viz_every_n_epochs == 7, "wrapper inherits the base viz gate" + tv.trainer = SimpleNamespace(default_root_dir="/tmp/run") + assert base.trainer is tv.trainer + assert tv.video_dir().endswith("videos_train_viz") + + op = PrefixedEvalVideo(base, prefix="unseen_op_valid") + metrics, _ = op.compute_metrics_and_viz({}, do_viz=False) + assert set(metrics) == {"unseen_op_valid/Valid/x"}, metrics + op.trainer = SimpleNamespace(default_root_dir="/tmp/run") + assert op.video_dir().endswith("videos_unseen_op_valid") diff --git a/tests/unit/test_hydra_configs.py b/tests/unit/test_hydra_configs.py index 0cdf0de78b..5a24abb428 100644 --- a/tests/unit/test_hydra_configs.py +++ b/tests/unit/test_hydra_configs.py @@ -31,13 +31,20 @@ def _options(group: str) -> list[str]: TRAIN_TOP_LEVEL = [c for c in TOP_LEVEL if c.startswith("train_")] -@pytest.mark.parametrize("config_name", TOP_LEVEL) +# train_* configs are composed (with more checks) by +# test_infra_port_wiring.py::test_human_action_width_matches_data_mode. +@pytest.mark.parametrize( + "config_name", [c for c in TOP_LEVEL if c not in TRAIN_TOP_LEVEL] +) def test_top_level_config_resolves(config_name, compose_resolve): compose_resolve(config_name, []) @pytest.mark.parametrize("data", _options("data")) -@pytest.mark.parametrize("top", TRAIN_TOP_LEVEL) +# train_zarr_cartesian x data is composed by test_data_filters_instantiate_to_dataset_filter. +@pytest.mark.parametrize( + "top", [t for t in TRAIN_TOP_LEVEL if t != "train_zarr_cartesian"] +) def test_data_option_resolves_under_every_train_config(top, data, compose_resolve): compose_resolve(top, [f"data={data}"]) @@ -92,3 +99,30 @@ def test_submitit_launcher_config_resolves(launcher, compose_resolve): launcher_cfg = OmegaConf.to_container(cfg.hydra.launcher, resolve=True) if launcher.startswith("submitit"): assert launcher_cfg["timeout_min"] > 0 + + +KEYPOINT_REVERT = "egomimic.rldb.embodiment.human._build_human_keypoints_revert_6d_wristframe_transform_list" + + +@pytest.mark.parametrize( + "top,model,data", + [ + ("train_zarr_cartesian", "hpt_bc_flow_mecka", "mecka"), + ("train_zarr_cartesian", "hpt_bc_flow_aria", "aria"), + ("train_zarr_cartesian", "hpt_bc_flow_scale", "scale"), + ("train_zarr_cartesian_pi", "pi0.5_bc_mecka", "mecka"), + ("train_zarr_cartesian_pi", "pi0.5_bc_aria", "aria"), + ("train_zarr_cartesian_pi", "pi0.5_bc_scale", "scale"), + ], +) +def test_default_evaluator_matches_the_human_keypoint_action( + top, model, data, compose_resolve +): + """Validation reads predictions[f"{emb}_{viz action_key}"], so the default + evaluator must visualize and revert the action the human recipes train.""" + cfg = compose_resolve(top, [f"model={model}", f"data={data}"]) + ac_key = cfg.model.robomimic_model.ac_keys.human_bimanual + assert ac_key == "actions_keypoints" + assert cfg.evaluator.viz_func.human_bimanual.action_key == ac_key + revert = cfg.evaluator.transform_lists.human_bimanual._target_ + assert revert == KEYPOINT_REVERT diff --git a/tests/unit/test_infra_port_wiring.py b/tests/unit/test_infra_port_wiring.py index 9b59cbfa2a..6665502921 100644 --- a/tests/unit/test_infra_port_wiring.py +++ b/tests/unit/test_infra_port_wiring.py @@ -9,6 +9,7 @@ import numpy as np import pytest import torch +from lightning import LightningModule from scipy.spatial.transform import Rotation as R import egomimic @@ -343,3 +344,192 @@ def check(v): assert check(1.0 + 0.4 * 2.0) is None # within 50 % of the [-1, 1] range assert check(1.0 + 0.6 * 2.0) is not None assert check(1e9) is not None + + +# ------------------------------------------------------ train_viz second loader +class _Dicts: + def __init__(self, n, tag): + self.n, self.tag = n, tag + + def __len__(self): + return self.n + + def __getitem__(self, i): + return {"x": np.float32(i), "tag": [self.tag]} + + +def test_val_dataloader_returns_two_loaders_with_train_viz(): + from lightning.pytorch.utilities.combined_loader import CombinedLoader + + from egomimic.pl_utils.pl_data_utils import MultiDataModuleWrapper + + params = {"human_bimanual": {"batch_size": 2, "num_workers": 0}} + dm = MultiDataModuleWrapper( + train_datasets={"human_bimanual": _Dicts(4, "train")}, + valid_datasets={"human_bimanual": _Dicts(4, "valid")}, + train_dataloader_params=params, + valid_dataloader_params=params, + held_out_operators=["op1"], + ) + assert isinstance(dm.val_dataloader(), CombinedLoader) + dm = MultiDataModuleWrapper( + train_datasets={"human_bimanual": _Dicts(4, "train")}, + valid_datasets={"human_bimanual": _Dicts(4, "valid")}, + train_dataloader_params=params, + valid_dataloader_params=params, + train_viz_datasets={"human_bimanual": _Dicts(4, "viz")}, + train_viz_dataloader_params={ + "human_bimanual": {**params["human_bimanual"], "shuffle": False} + }, + ) + loaders = dm.val_dataloader() + assert isinstance(loaders, list) and len(loaders) == 2 + batch, _, _ = next(iter(loaders[1])) + assert batch["human_bimanual"]["tag"] == [["viz"], ["viz"]] + + +def test_validation_step_routes_by_dataloader_idx(): + from egomimic.pl_utils.pl_model import ModelWrapper + + calls = [] + + class _Ev: + def __init__(self, name): + self.name = name + + def on_validation_step(self, batch, batch_idx, dataloader_idx=0): + calls.append((self.name, batch, batch_idx, dataloader_idx)) + + w = ModelWrapper.__new__(ModelWrapper) + LightningModule.__init__(w) # global_rank etc. without a trainer + w.model = SimpleNamespace(process_batch_for_training=lambda b: {"processed": b}) + w.evaluator = _Ev("valid") + w.train_viz_evaluator = _Ev("train_viz") + w.unseen_op_valid_evaluator = None + w.val_loader_names = ["valid", "train_viz"] + # Lightning hands the outer sequential loader's inner triple as `batch` + w.validation_step(({"k": 1}, 0, 0), 0, dataloader_idx=0) + w.validation_step(({"k": 2}, 0, 1), 3, dataloader_idx=1) + w.validation_step({"k": 3}, 1, dataloader_idx=0) + assert calls == [ + ("valid", {"processed": {"k": 1}}, 0, 0), + ("train_viz", {"processed": {"k": 2}}, 3, 1), + ("valid", {"processed": {"k": 3}}, 1, 0), + ] + w.train_viz_evaluator = None + w.validation_step({"k": 4}, 0, dataloader_idx=1) # no second head: ignored + assert len(calls) == 3 + + +def test_unseen_op_valid_third_loader_and_routing(): + """unseen_op_valid_datasets adds a loader after train_viz; without train_viz it + takes idx 1, and ModelWrapper routes by val_loader_names, not position.""" + from egomimic.pl_utils.pl_data_utils import MultiDataModuleWrapper + from egomimic.pl_utils.pl_model import ModelWrapper + + params = {"human_bimanual": {"batch_size": 2, "num_workers": 0}} + common = dict( + train_datasets={"human_bimanual": _Dicts(4, "train")}, + valid_datasets={"human_bimanual": _Dicts(4, "valid")}, + train_dataloader_params=params, + valid_dataloader_params=params, + unseen_op_valid_datasets={"human_bimanual": _Dicts(4, "unseen")}, + unseen_op_valid_dataloader_params=params, + ) + dm = MultiDataModuleWrapper( + **common, + train_viz_datasets={"human_bimanual": _Dicts(4, "viz")}, + train_viz_dataloader_params=params, + ) + assert dm.val_loader_names() == ["valid", "train_viz", "unseen_op_valid"] + loaders = dm.val_dataloader() + tags = [next(iter(ld))[0]["human_bimanual"]["tag"][0][0] for ld in loaders] + assert tags == ["valid", "viz", "unseen"] + + dm = MultiDataModuleWrapper(**common) + assert dm.val_loader_names() == ["valid", "unseen_op_valid"] + + calls = [] + + class _Ev: + def __init__(self, name): + self.name = name + + def on_validation_step(self, batch, batch_idx, dataloader_idx=0): + calls.append((self.name, dataloader_idx)) + + w = ModelWrapper.__new__(ModelWrapper) + LightningModule.__init__(w) + w.model = SimpleNamespace(process_batch_for_training=lambda b: b) + w.evaluator = _Ev("valid") + w.train_viz_evaluator = None + w.unseen_op_valid_evaluator = _Ev("unseen_op_valid") + w.val_loader_names = dm.val_loader_names() + w.validation_step({"k": 1}, 0, dataloader_idx=0) + w.validation_step({"k": 2}, 0, dataloader_idx=1) + assert calls == [("valid", 0), ("unseen_op_valid", 1)] + + +# ------------------------------------------------------ train_viz on by default +def _viz_cfg(**top): + from omegaconf import OmegaConf + + data = top.pop("data", {}) + cfg = { + "mode": "train", + "evaluator": { + "_target_": "egomimic.eval.eval_hpt.HPTEvalVideo", + "viz_every_n_epochs": 7, + }, + "data": { + "valid_dataloader_params": {HUMAN: {"batch_size": 4, "num_workers": 3}}, + **data, + }, + **top, + } + return OmegaConf.create(cfg) + + +def test_train_viz_defaults_to_the_train_datasets(): + import egomimic.trainHydra as th + + train = {HUMAN: object()} + viz, params = th._train_viz_datasets(_viz_cfg(), train, instantiate=None) + assert viz[HUMAN] is train[HUMAN], "reuses the train split, no second resolve" + # valid loader params, never shuffled (same leading slice every val) + assert params == {HUMAN: {"batch_size": 4, "num_workers": 3, "shuffle": False}} + + +def test_train_viz_explicit_datasets_and_opt_outs(): + import egomimic.trainHydra as th + + train = {HUMAN: object()} + explicit = _viz_cfg( + data={ + "train_viz_datasets": {HUMAN: {"tag": "explicit"}}, + "train_viz_dataloader_params": {HUMAN: {"batch_size": 32}}, + } + ) + viz, params = th._train_viz_datasets( + explicit, train, instantiate=lambda node, dataset_name: ("built", node.tag) + ) + assert viz == {HUMAN: ("built", "explicit")} + assert params is None, "the data config's own params are left to Hydra" + + for off in ( + _viz_cfg(train_viz=False), + _viz_cfg(mode="eval"), + _viz_cfg(evaluator=None), + ): + assert th._train_viz_datasets(off, train, instantiate=None) == ({}, None) + + +def test_train_viz_evaluator_wraps_the_canonical_evaluator(): + import egomimic.trainHydra as th + from egomimic.eval.eval_train_viz import TrainVizEvalVideo + + cfg = _viz_cfg() + ev = th._build_train_viz_evaluator(cfg) + assert isinstance(ev, TrainVizEvalVideo) + assert ev.viz_every_n_epochs == 7 + assert th._build_train_viz_evaluator(_viz_cfg(train_viz=False)) is None