data: mecka_fold_flagship_top3_hpt_6d, the 3-operator flagship train split

.gitignore

@@ -45,10 +45,20 @@ annotations_test/**
 log_conversion/**
 debug_model_inputs/**
 temp_dir/**
+/datasets/
+apikey.txt
+slurm-*.out
+slurmoutputs/
+*.log
+.inductor_cache/
+
 scratch/
 sample_data/
 external_ckpts/
 external/MANO/
+ann_backup/
+aria_ann_backup/
+sort_annotations/
 
 # setuptools build output
 build/

egomimic/algo/hpt.py

@@ -145,6 +145,9 @@ def init_domain_stem(self, domain_name, stem_spec):
             A specification containing configurations for each modality's stem.
         """
 
+        # A ``null`` entry lets a child config drop a stem its base declares
+        # (hydra merges cannot delete keys).
+        stem_spec = {k: v for k, v in stem_spec.items() if v is not None}
         self.stem_spec[domain_name] = stem_spec
         self.modalities[domain_name] = list(stem_spec.keys())
 
@@ -909,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)

egomimic/algo/pi.py

@@ -25,7 +25,7 @@
     gather_pi_images,
 )
 from egomimic.rldb.embodiment.embodiment import get_embodiment, get_embodiment_id
-from egomimic.utils.action_utils import ConverterRegistry
+from egomimic.utils.action_utils import ConverterRegistry, pad_to_width
 
 logger = logging.getLogger(__name__)
 # Ensure logger propagates to root logger and has appropriate level
@@ -142,6 +142,7 @@ def __init__(
         self.action_converters = action_converters
 
         self.action_registry = ConverterRegistry()
+        self._packed_widths: dict[tuple, int] = {}
 
         arcfg = self.action_converters
         default_ac_key = getattr(arcfg, "ac_key", "actions_cartesian")
@@ -195,9 +196,45 @@ def __init__(
                 "No pytorch_weight_path: no base weights loaded. The weights must "
                 "come from a checkpoint; otherwise this trains from scratch."
             )
+        # openpi's PyTorch port hard-codes 32-wide action projections
+        # regardless of ``Pi0Config.action_dim``. Wider action spaces (the
+        # 144-D hand-keypoint action, or a cotrain that pads eva's 20-D into
+        # the same vector) need ``action_dim``-wide ones, swapped in AFTER the
+        # (strict) base-weight load. Only these two layers start fresh;
+        # ``sample_actions`` already draws its noise at ``config.action_dim``.
+        self.action_dim = int(self.config.model.action_dim)
+        if self.action_dim != 32:
+            self._resize_action_projections(self.action_dim)
         self.nets = nn.ModuleDict()
         self.nets["policy"] = self.model
 
+    def _resize_action_projections(self, action_dim: int) -> None:
+        target = (
+            self.model.module
+            if isinstance(self.model, torch.nn.parallel.DistributedDataParallel)
+            else self.model
+        )
+        old_in, old_out = target.action_in_proj, target.action_out_proj
+        width = old_in.out_features
+        new_in = nn.Linear(action_dim, width).to(
+            dtype=old_in.weight.dtype, device=old_in.weight.device
+        )
+        new_out = nn.Linear(width, action_dim).to(
+            dtype=old_out.weight.dtype, device=old_out.weight.device
+        )
+        target.action_in_proj = new_in
+        target.action_out_proj = new_out
+        logger.warning(
+            "action_dim=%d != 32: re-initialized action_in_proj (%d->%d) and "
+            "action_out_proj (%d->%d); the pretrained 32-wide projections are "
+            "discarded, everything else comes from the base checkpoint.",
+            action_dim,
+            action_dim,
+            width,
+            width,
+            action_dim,
+        )
+
     def _control_mode_for(self, emb_name: str | None) -> str:
         if self.control_mode and emb_name is not None:
             for key, val in self.control_mode.items():
@@ -479,19 +516,40 @@ def forward_eval(self, batch):
         return unnorm_preds
 
     def _reduce_loss(self, losses, action, embodiment_id, ac_key) -> torch.Tensor:
-        """Mean flow-matching loss for the embodiment."""
+        """Mean flow-matching loss over the embodiment's packed action width.
+
+        openpi returns the per-element loss ``(B, H, action_dim)``. With
+        ``action_dim`` widened for the 144-D keypoint action, a narrower
+        embodiment (eva's 32-slot cartesian layout) would otherwise average
+        its loss over 112 zero-padded "predict the noise" dims and contribute
+        a fraction of the gradient signal it had at width 32; restrict the
+        mean to the slots the converter actually packs.
+        """
         if isinstance(losses, (list, tuple)):
             losses = torch.stack(losses)
         elif not isinstance(losses, torch.Tensor):
             losses = torch.tensor(losses, device=action.device, dtype=torch.float32)
+        if losses.ndim == 3:
+            losses = losses[..., : self._packed_width(action, embodiment_id, ac_key)]
         return losses.mean()
 
+    def _packed_width(self, action, embodiment_id, ac_key) -> int:
+        """Width the embodiment's converter packs into, cached per pairing: it
+        is a property of the converter, not of the batch."""
+        key = (embodiment_id, ac_key)
+        if key not in self._packed_widths:
+            converter = self.action_registry.get(embodiment_id, ac_key)
+            self._packed_widths[key] = int(
+                converter.to32_norm_6d(action[:1, :1]).shape[-1]
+            )
+        return self._packed_widths[key]
+
     def _postprocess_sampled_actions(self, pred_actions, _batch, embodiment_id, ac_key):
         """Raw ``sample_actions`` output -> unnormalized native action dict.
 
-        Unpack the normalized native action (xyz+6D(+gripper)) from the
-        model's action vector, then unnormalize via the standard pipeline
-        (stats were computed on the native layout)."""
+        Unpack the normalized native action (xyz+6D(+gripper) or the 144-D
+        keypoint vector) from the model's action vector, then unnormalize via
+        the standard pipeline (stats were computed on the native layout)."""
         ref = _batch[ac_key]
         _, T, D = ref.shape
         converter = self.action_registry.get(embodiment_id, ac_key)
@@ -549,7 +607,7 @@ def log_info(self, info):
         return log
 
     def _robomimic_to_pi_data(self, batch, proprio_keys, lang_keys, ac_key, embodiment):
-        """Dataset batch -> (openpi observation, 32-D action chunk)."""
+        """Dataset batch -> (openpi observation, ``action_dim``-wide action chunk)."""
         if ac_key not in batch:
             raise KeyError(f"Missing action key '{ac_key}' in batch")
 
@@ -561,8 +619,8 @@ def _robomimic_to_pi_data(self, batch, proprio_keys, lang_keys, ac_key, embodime
         converter = self.action_registry.get(emb_id, ac_key)
         # The action is already normalized and in its native layout (the
         # ypr->6D conversion happened in the data transforms). Pack it into
-        # the canonical 32-slot block layout.
-        action32 = converter.to32_norm_6d(action)
+        # the canonical block layout and zero-pad up to the model width.
+        action32 = pad_to_width(converter.to32_norm_6d(action), self.action_dim)
 
         # OpenPI expects a fixed camera tuple under its own names. Human data
         # only has the front camera; the missing wrist slots get a copy of it

egomimic/eval/action_metrics.py

@@ -0,0 +1,210 @@
+"""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 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.
+
+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 (
+    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)
+    metrics[f"{prefix}_xyz_paired_mse_avg"] = xyz_p
+    metrics[f"{prefix}_rot6d_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) -> dict:
+    """Dispatch on the last-dim width: cartesian, keypoint, or ``{}``."""
+    metrics = cartesian_metrics(pred, gt, prefix)
+    if metrics:
+        return metrics
+    return keypoint_metrics(pred, gt, prefix)

egomimic/eval/eval.py

@@ -18,5 +18,8 @@ def on_validation_end(self):
         pass
 
     @abstractmethod
-    def on_validation_step(self, batch, batch_idx, dataloader_idx=0):
+    def on_validation_step(self, batch, batch_idx, dataloader_idx=0, mode="both"):
+        """``mode``: "metrics" / "video" / "both". An evaluator that renders
+        nothing may ignore it, but it must accept it -- ModelWrapper passes it
+        for every loader of a config that declares ``video_episodes``."""
         pass