data/rdt: keypoint RDT recipe, pose-glitch masking, hardened staging

egomimic/algo/rdt.py

@@ -10,6 +10,7 @@
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
 from egomimic.algo.hpt import HPT, HPTModel
 from egomimic.models.rdt_nets import (
@@ -48,6 +49,10 @@ class RDTModel(HPTModel):
     steps are masked out. ``short_dropout`` replaces a sample's whole short
     image stream by a learned null token (training only), so the policy cannot
     lean on the latest frames alone.
+
+    ``lang_pad_to`` right-pads the language tokens (masked) to a multiple of it,
+    so torch.compile sees a fixed length instead of each batch's longest prompt;
+    masked keys get zero attention weight, so outputs do not change.
     """
 
     def __init__(
@@ -61,6 +66,7 @@ def __init__(
         image_history_dropout: float = 0.0,
         memory: dict | None = None,
         short_dropout: float = 0.0,
+        lang_pad_to: int | None = None,
         **kwargs,
     ):
         super().__init__(
@@ -74,6 +80,7 @@ def __init__(
         self.image_history_dropout = float(image_history_dropout)
         self.stem_modality = {}
         self.short_dropout = float(short_dropout)
+        self.lang_pad_to = lang_pad_to
         self.long_memory = (
             None if memory is None else RDTMemory(hidden_dim=embed_dim, **memory)
         )
@@ -137,7 +144,7 @@ def _image_tokens(self, name, modality, stem, x):
             raise ValueError(
                 f"'{modality}' carries {T} frames but trunk.image_history="
                 f"{self.image_history}; set it to match the data config's "
-                "image_history_gap_s (2 frames when set, else 1)"
+                "key_map.image_history (when image_history_gap_s is set, else 1)"
             )
         drop = self._drop(self.image_history_dropout, B, x.device) if T > 1 else None
         if drop is not None:
@@ -222,6 +229,10 @@ def forward_features(self, domain, data):
         if lang:
             cond.lang = torch.cat(lang, dim=1)
             cond.lang_mask = torch.cat(lang_mask, dim=1)
+            if self.lang_pad_to:
+                pad = -cond.lang.shape[1] % self.lang_pad_to
+                cond.lang = F.pad(cond.lang, (0, 0, 0, pad))
+                cond.lang_mask = F.pad(cond.lang_mask, (0, pad), value=False)
         if self.long_memory is not None:
             # after state dropout, so a fused memory loses its proprio with it
             cond.memory, cond.memory_mask = self.long_memory(
@@ -251,7 +262,7 @@ def _robomimic_to_hpt_data(
         domain="",
     ):
         """HPT's layout plus ``fps``; per camera with a ``*_hist`` twin, the
-        (past, current) frame pair on the image input's time axis; and a
+        (past..., current) frames on the image input's time axis; and a
         ``*_mem`` window as ``memory`` / ``memory_mask``, un-augmented (its
         tower is frozen and the stem normalizes)."""
         past = {
@@ -272,13 +283,16 @@ def _robomimic_to_hpt_data(
         )
         for key, past_key in past.items():
             short = key.rsplit(".", 1)[-1]
-            # (B * 2, 3, H, W), each sample's (past, current) adjacent: one
-            # jitter draw per sample covers both frames, and every op -- a
-            # contrast mean, a crop, a resize -- still sees one frame.
-            pair = torch.stack([batch[past_key], batch[key]], dim=1)
-            n = pair.shape[0]
-            pair = self._apply_image_augs(pair.flatten(0, 1), short, frames=2)
-            data[short] = pair.reshape(n, 2, *pair.shape[1:]).unsqueeze(2)
+            # (B * T, 3, H, W), each sample's frames adjacent: one jitter draw
+            # per sample covers all of them, and every op -- a contrast mean,
+            # a crop, a resize -- still sees one frame.
+            past_frames = batch[past_key]
+            if past_frames.dim() == 4:
+                past_frames = past_frames.unsqueeze(1)
+            clip = torch.cat([past_frames, batch[key].unsqueeze(1)], dim=1)
+            n, t = clip.shape[:2]
+            clip = self._apply_image_augs(clip.flatten(0, 1), short, frames=t)
+            data[short] = clip.reshape(n, t, *clip.shape[1:]).unsqueeze(2)
         for key in memory:
             data["memory"] = batch[key]
             data["memory_mask"] = batch[f"{key}{HISTORY_MASK_SUFFIX}"] > 0.5

egomimic/eval/eval_video.py

@@ -35,7 +35,7 @@ def __init__(
         viz_every_n_epochs: int = 1,
         viz_max_batches: int | None = None,
         viz_mode: str = "replay",
-        replay_chunks: int = 5,
+        replay_chunks: int = 17,
         replay_trail: int = 5,
         action_stride: float | None = None,
     ):
@@ -55,7 +55,8 @@ def __init__(
         # "replay" (egomimic/eval/replay_viz.py) plays each chunk as GT then
         # prediction; it needs consecutive frames, so it runs only on a pinned
         # video loader and caps it at ``replay_chunks`` chunks instead of
-        # ``viz_max_batches``. "overlay" draws each frame's whole GT and
+        # ``viz_max_batches`` (17: ~17 s of a 30-frame-chunk episode, 1020 video
+        # frames). "overlay" draws each frame's whole GT and
         # predicted chunk, and is what a head without a video loader gets.
         # ``action_stride`` is frames per chunk step, fractional when the chunk
         # is resampled (None: inferred from the GT keypoints, see

egomimic/hydra_configs/data/mecka_all_kp_pinned_val.yaml

@@ -0,0 +1,45 @@
+# All mecka episodes with the wrist-frame KEYPOINT pipeline, and mecka_20pct_6d's
+# validation heads and video pins so runs are comparable to the 20 % ones:
+#   actions_keypoints              (100, 144)  per hand [wrist xyz+rot6d (9) |
+#                                              21 keypoints (63)], wrist frame
+#   observations.state.keypoints   (144,)      per hand [wrist pose in HEAD
+#                                              frame | keypoints in wrist frame]
+# train = every lab=mecka episode except the bad-egomotion blocklist, the 3 valid
+# pins and all 10 episodes of the 3 unseen_op_valid operators (38,596 episodes,
+# ~77M frames vs 4,078 / 10.5M for the 20 % subset, SQL 2026-09-23).
+defaults:
+  - mecka_20pct_6d
+  - _self_
+
+# SQL app.episodes.operator of the 3 unseen_op_valid pins.
+held_out_operators:
+  - 69475d885c0479a0fb81408e
+  - 69a84eed47e12c9b339208fd
+  - 6974f58bb393af9134c82247
+
+train_datasets:
+  human_bimanual:
+    resolver:
+      key_map:
+        keymap_mode: keypoints
+      transform_list:
+        mode: keypoints_wristframe_6d
+        pad_proprio_gripper: false
+    filters:
+      episode_hashes: null
+      filter_lambdas:
+        - "lambda row: str(row.get('lab', '')).lower() == 'mecka'"
+        - "lambda row: str(row.get('operator', '')) not in ${data.held_out_operators}"
+        - "lambda row: str(row.get('episode_hash', '')) not in ${data.valid_datasets.human_bimanual.filters.episode_hashes}"
+
+# The 20 % heads interpolate the train lambdas, which now exclude their own pins.
+valid_datasets:
+  human_bimanual:
+    filters:
+      filter_lambdas:
+        - "lambda row: str(row.get('lab', '')).lower() == 'mecka'"
+
+unseen_op_valid_datasets:
+  human_bimanual:
+    filters:
+      filter_lambdas: ${data.valid_datasets.human_bimanual.filters.filter_lambdas}

egomimic/hydra_configs/data/mecka_flagship_fold_cup_kp_holdout.yaml

@@ -0,0 +1,21 @@
+# mecka_flagship_fold_kp_holdout plus the flagship cup_on_saucer episodes (5409
+# eps / 905k frames): 7016 eps / 1.22M frames. 3 held out, 1 fold (shared with
+# the parent's valid) and 2 cup, each from a different operator.
+defaults:
+  - mecka_flagship_fold_kp_holdout
+  - _self_
+
+tasks: [fold_clothes, cup_on_saucer]
+
+held_out_hashes:
+  - 6930852199946db9586cf530
+  - 692ed3f8f2e4e72c033b9970
+  - 692eb00f1e78e176ba6cf9b2
+
+metric_frames_per_episode:
+  train_viz: 1
+
+video_episodes:
+  valid:
+    - 6930852199946db9586cf530
+    - 692ed3f8f2e4e72c033b9970

egomimic/hydra_configs/data/mecka_flagship_fold_kp_holdout.yaml

@@ -0,0 +1,88 @@
+# Mecka FLAGSHIP fold_clothes (SQL lab='mecka', task='fold_clothes', 1607 eps /
+# 316k frames under processed_v3/mecka/flagship/) with the wrist-frame KEYPOINT
+# pipeline of mecka_all_kp_pinned_val, and 3 hand-picked held-out episodes as the
+# only validation: a fitting test, not an operator-generalization split.
+#   train = every other episode (mode total), minus the bad-egomotion blocklist
+#   valid = the `held_out_hashes` (3 operators, ~200 frames each, run in full)
+# Children swap `tasks` / `held_out_hashes` (mecka_flagship_fold_cup_kp_holdout,
+# mecka_flagship_fold_kp_overfit1). Norm stats are split-specific.
+_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper
+
+tasks: [fold_clothes]
+
+# 2026-09-25, fold_clothes rows with 40-60th percentile length, one per operator.
+held_out_hashes:
+  - 6930852199946db9586cf530
+  - 692e9f2cba0400c70689fdfd
+  - 692eefe40ba268eb16564ab9
+
+train_datasets:
+  human_bimanual:
+    _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
+    resolver:
+      _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
+      folder_path: ${paths.dataset_dir}
+      key_map:
+        _target_: egomimic.rldb.embodiment.human.Human.get_keymap
+        keymap_mode: keypoints
+        annotation_key: annotations
+      transform_list:
+        _target_: egomimic.rldb.embodiment.human.Human.get_transform_list
+        mode: keypoints_wristframe_6d
+        stride: 1
+        fix_left_wrist_convention: true
+        pad_proprio_gripper: false
+    filters:
+      _target_: egomimic.rldb.filters.DatasetFilter
+      filter_lambdas:
+        - "lambda row: str(row.get('lab', '')).lower() == 'mecka'"
+        - "lambda row: str(row.get('task', '')).lower() in ${data.tasks}"
+        - "lambda row: '/mecka/flagship/' in str(row.get('zarr_processed_path', ''))"
+        - "lambda row: str(row.get('episode_hash', '')) not in ${data.held_out_hashes}"
+      exclude_hashes: mecka_bad_egomotion
+    mode: total
+
+valid_datasets:
+  human_bimanual:
+    _target_: ${data.train_datasets.human_bimanual._target_}
+    resolver: ${data.train_datasets.human_bimanual.resolver}
+    filters:
+      _target_: egomimic.rldb.filters.DatasetFilter
+      filter_lambdas:
+        - "lambda row: str(row.get('lab', '')).lower() == 'mecka'"
+      episode_hashes: ${data.held_out_hashes}
+    mode: total
+
+train_viz_datasets:
+  human_bimanual:
+    _target_: ${data.train_datasets.human_bimanual._target_}
+    resolver: ${data.train_datasets.human_bimanual.resolver}
+    filters: ${data.train_datasets.human_bimanual.filters}
+    mode: total
+
+train_dataloader_params:
+  human_bimanual:
+    batch_size: 32
+    num_workers: 12
+
+valid_dataloader_params:
+  human_bimanual:
+    batch_size: 32
+    num_workers: 8
+    shuffle: false
+
+train_viz_dataloader_params:
+  human_bimanual:
+    batch_size: 32
+    num_workers: 4
+    shuffle: false
+
+# Single-rank K = floor(5120 / episodes); the 3 valid episodes run in full.
+metric_frames_per_episode:
+  train_viz: 3
+
+video_episodes:
+  valid:
+    - 6930852199946db9586cf530
+  train_viz:
+    - 692fe87a3c1dfb13a7a99ee6

egomimic/hydra_configs/data/mecka_flagship_fold_kp_overfit1.yaml

@@ -0,0 +1,22 @@
+# Overfit check: ONE flagship fold_clothes episode (191 frames), trained on and
+# validated on itself (valid runs without dropout or augmentation, so it reads
+# how well the episode was memorized). train_viz would duplicate valid; the
+# train config turns it off.
+defaults:
+  - mecka_flagship_fold_kp_holdout
+  - _self_
+
+train_datasets:
+  human_bimanual:
+    filters:
+      episode_hashes: [692fe87a3c1dfb13a7a99ee6]
+
+valid_datasets:
+  human_bimanual:
+    filters:
+      episode_hashes: [692fe87a3c1dfb13a7a99ee6]
+
+video_episodes:
+  valid:
+    - 692fe87a3c1dfb13a7a99ee6
+  train_viz: []

egomimic/hydra_configs/data/mecka_freeform_fold_kp_all_holdout.yaml

@@ -0,0 +1,18 @@
+# mecka_freeform_fold_kp_top3_holdout on every freeform folding_clothes operator:
+# the same seen_op_valid (top3's 6 held-out episodes) and unseen_op_valid, whose 5
+# operators (`held_out_operators`) are dropped from train. After the blocklist
+# (2026-09-29): 423 of 462 episodes, 9.6 h; train is 171 operators, 407
+# episodes, 9.2 h.
+defaults:
+  - mecka_freeform_fold_kp_top3_holdout
+  - _self_
+
+train_datasets:
+  human_bimanual:
+    filters:
+      filter_lambdas:
+        - "lambda row: str(row.get('lab', '')).lower() == 'mecka'"
+        - "lambda row: str(row.get('task', '')).lower() in ${data.tasks}"
+        - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))"
+        - "lambda row: str(row.get('operator', '')) not in ${data.held_out_operators}"
+        - "lambda row: str(row.get('episode_hash', '')) not in ${data.held_out_hashes}"