diff --git a/egomimic/algo/hpt.py b/egomimic/algo/hpt.py index 7565b44c4e..b66886afbf 100644 --- a/egomimic/algo/hpt.py +++ b/egomimic/algo/hpt.py @@ -407,16 +407,24 @@ def stem_process(self, domain, data): horizon = np.random.randint(1, data_horizon + 1) data[modality] = data[modality][:, data_horizon - horizon :] - positional_embedding = get_sinusoid_encoding_table( - 0, horizon * int(np.prod(data_shape[2:-1])), data_shape[-1] - ).to(data[modality]) - positional_embedding = einops.repeat( - positional_embedding, "b h w -> (repeat b) h w", repeat=data_shape[0] - ) + # An observation-history stem owns its own LEARNED time-step + # embedding (added after the projection); the input-side sinusoid + # here is a constant [0, 1, 0, 1, ...] added to the NORMALIZED + # values, which would be amplitude-1 noise on [-1, 1] data once + # axis 1 carries more than the current step. + if getattr(stem, "history_len", 1) == 1: + positional_embedding = get_sinusoid_encoding_table( + 0, horizon * int(np.prod(data_shape[2:-1])), data_shape[-1] + ).to(data[modality]) + positional_embedding = einops.repeat( + positional_embedding, + "b h w -> (repeat b) h w", + repeat=data_shape[0], + ) - data[modality] = data[modality] + positional_embedding.view( - data[modality].shape - ) + data[modality] = data[modality] + positional_embedding.view( + data[modality].shape + ) stem_token = stem.compute_latent(data[modality]) feats.append(stem_token) feat_dict[modality] = stem_token @@ -1169,7 +1177,13 @@ def forward_training(self, batch): ac_key = self.ac_keys[embodiment_id] aux_ac_keys = self.auxiliary_ac_keys.get(embodiment_name, []) data = self._robomimic_to_hpt_data( - _batch, cam_keys, proprio_keys, lang_keys, ac_key, aux_ac_keys + _batch, + cam_keys, + proprio_keys, + lang_keys, + ac_key, + aux_ac_keys, + domain=embodiment_name, ) hpt_batch = { "domain": embodiment_name, # readability on config side @@ -1227,7 +1241,13 @@ def forward_eval(self, batch): ac_key = self.ac_keys[embodiment_id] aux_ac_keys = self.auxiliary_ac_keys.get(embodiment_name, []) data = self._robomimic_to_hpt_data( - _batch, cam_keys, proprio_keys, lang_keys, ac_key, aux_ac_keys + _batch, + cam_keys, + proprio_keys, + lang_keys, + ac_key, + aux_ac_keys, + domain=embodiment_name, ) hpt_batch = { "domain": embodiment_name, # readability on config side @@ -1357,8 +1377,25 @@ def _apply_image_augs(self, images, short): augs = self.train_image_augs if self.nets.training else self.eval_image_augs return augs(images) if augs is not None else images + def _stem_history_len(self, domain: str, modality: str) -> int: + """``history_len`` of ``domain``'s ``modality`` stem (1 when the stem + declares none). Stems are registered as ``f"{domain}_{modality}"``; a + cotrain run has one per domain, so the domain has to be named -- a + suffix match would return whichever came first in the dict. + """ + stem = self.nets["policy"].stems[f"{domain}_{modality}"] + return int(getattr(stem, "history_len", 1)) + def _robomimic_to_hpt_data( - self, batch, cam_keys, proprio_keys, lang_keys, ac_key, aux_ac_keys=[] + self, + batch, + cam_keys, + proprio_keys, + lang_keys, + ac_key, + aux_ac_keys=[], + *, + domain: str, ): """ helper method that returns data in the format required for the HPT model @@ -1371,7 +1408,22 @@ def _robomimic_to_hpt_data( for key in proprio_keys: if key in batch: short = key.rsplit(".", 1)[-1] - data[f"state_{short}"] = batch[key].unsqueeze(1) + modality = f"state_{short}" + value = batch[key] + if value.ndim == 2: + data[modality] = value.unsqueeze(1) # (B, D) -> (B, 1, D) + else: + # Observation history: already (B, K, D), one token per step. + history_len = self._stem_history_len(domain, modality) + if value.shape[1] != history_len: + raise ValueError( + f"proprio '{key}' carries {value.shape[1]} history " + f"steps but the '{modality}' stem has " + f"history_len={history_len}; set model...stem_specs." + f"{domain}.{modality}.history_len and the keymap's " + "proprio_history to the same K" + ) + data[modality] = value for key in cam_keys: if key in batch: diff --git a/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml index 08615cc94f..ffc891e2a1 100644 --- a/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml +++ b/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml @@ -76,6 +76,18 @@ train_datasets: _target_: egomimic.rldb.embodiment.human.Human.get_keymap keymap_mode: cartesian annotation_key: annotations + # Proprio observation history: the ee_pose keys read K steps spaced + # `history_stride` frames apart, (K, 20) with the current frame last. + # At 30 fps the median wrist translation is 6.8 mm over 1 frame but + # 42.9 mm over 5, so an unstrided window is mostly sensor noise while + # the action chunk it conditions spans a full second; K = 3, s = 5 is + # 0.33 s of past. Must match the stem's `history_len` (the run fails + # loudly if not), and the K = 1 baseline needs + # ...key_map.proprio_history=1 ...key_map.history_stride=1 + # model.robomimic_model.stem_specs.human_bimanual.state_ee_pose.history_len=1 + # The two val loaders interpolate this resolver, so they follow. + proprio_history: 3 + history_stride: 5 transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list mode: cartesian_wristframe_6d diff --git a/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml index 66ca79a60a..84293e83ff 100644 --- a/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml @@ -123,6 +123,16 @@ robomimic_model: input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} output_dim: 840 widths: [840] + # Proprio observation history: K tokens (one per step, current step + # last) plus a learned zero-init time-step embedding, instead of the + # single current-step token. Must equal the data config's + # `proprio_history` (which also sets how far apart the steps are). + # `history_dropout` replaces the whole past by the current step for a + # sample (train only) so the policy still works when history is + # missing at episode start -- it is live at history_len 3, and it is + # what keeps the policy from reading its own past as the answer. + history_len: 3 + history_dropout: 0.2 specs: random_horizon_masking: false cross_attn: diff --git a/egomimic/models/hpt_nets.py b/egomimic/models/hpt_nets.py index 119b205902..4eea4ee131 100644 --- a/egomimic/models/hpt_nets.py +++ b/egomimic/models/hpt_nets.py @@ -541,10 +541,35 @@ def __init__( tanh_end: bool = False, ln: bool = True, num_of_copy: int = 1, + history_len: int = 1, + history_dropout: float = 0.0, **kwargs, ) -> None: - """vanilla MLP class""" + """vanilla MLP class + + ``history_len`` (K) > 1 makes this an observation-history stem: the + input carries the last K steps on axis 1 (current step last), one token + per step, and a LEARNED per-step embedding (zero-init, so at init the K + tokens are exactly today's projection of each step) is added AFTER the + projection. ``HPT.stem_process`` skips its input-side sinusoid for such + a stem: that sinusoid is a constant [0, 1, 0, 1, ...] added to the + NORMALIZED values, i.e. amplitude-1 noise on [-1, 1] data. + + ``history_dropout`` is the per-sample probability, in training only, of + replacing every past step by the current one: what the first frames of + an episode look like after the dataset's front padding (later short + windows repeat their oldest real frame instead), so the policy trains + on "history missing" (the mitigation for copycat / causal confusion). + Inert at K = 1 and in eval. + """ super().__init__(**kwargs) + if history_len > 1 and num_of_copy > 1: + raise ValueError( + f"MLPPolicyStem: history_len={history_len} and " + f"num_of_copy={num_of_copy} both consume axis 1; use one." + ) + self.history_len = int(history_len) + self.history_dropout = float(history_dropout) modules = [nn.Linear(input_dim, widths[0]), nn.SiLU()] for i in range(len(widths) - 1): @@ -562,6 +587,15 @@ def __init__( self.net = nn.ModuleList( [nn.Sequential(*modules) for _ in range(num_of_copy)] ) + if self.history_len > 1: + self.time_embed = nn.Parameter(torch.zeros(self.history_len, output_dim)) + + def _drop_history(self, x: torch.Tensor) -> torch.Tensor: + """Per sample, with probability ``history_dropout``, replace every past + step by the current one (the last slot of axis 1).""" + drop = torch.rand(x.shape[0], device=x.device) < self.history_dropout + current = x[:, -1:].expand_as(x) + return torch.where(drop.view(-1, *([1] * (x.ndim - 1))), current, x) def forward(self, x: torch.Tensor) -> torch.Tensor: """ @@ -572,6 +606,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Returns: Flatten tensor with shape [B, M, 512] """ + if self.history_len > 1: + if x.shape[1] != self.history_len: + raise ValueError( + f"MLPPolicyStem got {x.shape[1]} history steps on axis 1 but " + f"history_len={self.history_len} (input shape {tuple(x.shape)})" + ) + if self.training and self.history_dropout > 0: + x = self._drop_history(x) if self.num_of_copy > 1: out = [] iter_num = min(self.num_of_copy, x.shape[1]) @@ -582,6 +624,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: y = torch.stack(out, dim=1) else: y = self.net(x) + if self.history_len > 1: + # Learned time-step embedding on axis 1, last slot = current step. + y = y + self.time_embed.view( + 1, self.history_len, *([1] * (y.ndim - 3)), y.shape[-1] + ) return y diff --git a/egomimic/rldb/embodiment/human.py b/egomimic/rldb/embodiment/human.py index d8e69ef3e6..064a00a660 100644 --- a/egomimic/rldb/embodiment/human.py +++ b/egomimic/rldb/embodiment/human.py @@ -25,6 +25,7 @@ QuaternionPoseToYPR, Reshape, RotateLocalFrame, + SelectCurrentStep, SplitKeys, Transform, UnpadGripperZeros, @@ -232,6 +233,8 @@ def get_keymap( include_ee_pose: bool = False, norm_mode: bool = False, annotation_key: str = None, + proprio_history: int = 1, + history_stride: int = 1, ): """Build the keymap. Per-vendor knobs are explicit args from the data config: ``has_head_pose`` (False only for contributors that truly omit @@ -241,12 +244,27 @@ def get_keymap( the keypoint transform modes can build the cartesian ``observations.state.ee_pose`` proprio, the pi0.5 prompt state). ``norm_mode``/``annotation_key`` behave as in the base. + + ``proprio_history`` (K) makes the proprio keys the model consumes read + the last K frames instead of one: the dataset emits ``(K, D)`` with the + current frame last. K = 1 (the default) leaves the keymap exactly as it + was. ``history_stride`` (s) + spaces that window out -- the steps are ``idx - s*(K-1) ... idx`` -- + because at 30 fps consecutive frames are nearly a duplicate of the + current pose (median wrist translation 6.8 mm at lag 1 against 42.9 mm + at lag 5), so s = 1 spends tokens on sensor noise. """ + if norm_mode: + # Norm stats are per-channel current-step stats, shared by every K: + # each history step is normalized with them. + proprio_history = 1 key_map = cls._get_keymap( keymap_mode, has_head_pose=has_head_pose, include_aria_keypoints=include_aria_keypoints, include_ee_pose=include_ee_pose, + proprio_history=proprio_history, + history_stride=history_stride, ) if annotation_key is not None and not norm_mode: key_map[annotation_key] = { @@ -270,12 +288,15 @@ def _get_keymap( has_head_pose: bool = True, include_aria_keypoints: bool = False, include_ee_pose: bool = False, + proprio_history: int = 1, + history_stride: int = 1, ): """Canonical MANO keymap. ``include_aria_keypoints`` additionally exposes the raw Aria-layout proprio keypoints alongside the MANO ones; ``include_ee_pose`` adds the palm-origin ee_pose keys to the keypoints keymap. The front image is always ``VIZ_IMAGE_KEY``; Pi renames it - onto its own slot. + onto its own slot. ``proprio_history`` > 1 marks the proprio entries + the model consumes with ``"history": K`` (see ``get_keymap``). """ base_mode = _strip_pi_keymap_mode(cls, keymap_mode) front_key = cls.VIZ_IMAGE_KEY @@ -386,6 +407,28 @@ def _get_keymap( "key_type": "proprio_keys", "zarr_key": "obs_head_pose", } + + if proprio_history > 1: + if history_stride < 1: + raise ValueError(f"history_stride must be >= 1, got {history_stride}") + # Backward window on the proprio the MODEL consumes. Not on + # ``obs_head_pose``: every history step is expressed in the CURRENT + # step's head frame (a per-step head frame would make the history + # unreadable), and not on the horizoned ``action_*`` entries, which + # are forward chunks. + if base_mode == "cartesian": + history_keys = ["left.obs_ee_pose", "right.obs_ee_pose"] + else: + history_keys = [ + f"{side}.{name}" + for side in ("left", "right") + for name in ("obs_keypoints", "obs_wrist_pose") + ] + if include_ee_pose: + history_keys += ["left.obs_ee_pose", "right.obs_ee_pose"] + for key in history_keys: + key_map[key]["history"] = int(proprio_history) + key_map[key]["history_stride"] = int(history_stride) return key_map @classmethod @@ -574,6 +617,9 @@ def _build_human_keypoints_revert_eef_frame_transform_list( else: pose_shape = 6 transform_list = [ + # proprio_history: the proprio defines the frame, so keep the CURRENT + # step (no-op at K = 1, idempotent after the 6D revert). + SelectCurrentStep(keys=[obs_key]), SplitKeys( input_key=obs_key, output_key_list=[ @@ -712,17 +758,22 @@ def _build_human_keypoints_eef_frame_transform_list( output_key=right_keypoints_action_headframe, shape=(chunk_length, 21, 3), ), + # target_history: the proprio wrist pose may carry a leading time + # axis (proprio_history); everything is expressed in the CURRENT + # step's wrist frame either way. ActionChunkCoordinateFrameTransform( target_world=left_wrist_obs_headframe, chunk_world=left_keypoints_action_headframe, transformed_key_name=left_keypoints_action_wristframe, mode="xyz", + target_history=True, ), ActionChunkCoordinateFrameTransform( target_world=right_wrist_obs_headframe, chunk_world=right_keypoints_action_headframe, transformed_key_name=right_keypoints_action_wristframe, mode="xyz", + target_history=True, ), Reshape( input_key=left_keypoints_action_wristframe, @@ -749,12 +800,14 @@ def _build_human_keypoints_eef_frame_transform_list( pose_world=left_keypoints_obs_headframe, transformed_key_name=left_keypoints_obs_wristframe, mode="xyz", + target_history=True, ), PoseCoordinateFrameTransform( target_world=right_wrist_obs_headframe, pose_world=right_keypoints_obs_headframe, transformed_key_name=right_keypoints_obs_wristframe, mode="xyz", + target_history=True, ), Reshape( input_key=left_keypoints_obs_wristframe, @@ -771,12 +824,14 @@ def _build_human_keypoints_eef_frame_transform_list( chunk_world=left_wrist_action_headframe, transformed_key_name=left_wrist_action_wristframe, mode="xyzwxyz", + target_history=True, ), ActionChunkCoordinateFrameTransform( target_world=right_wrist_obs_headframe, chunk_world=right_wrist_action_headframe, transformed_key_name=right_wrist_action_wristframe, mode="xyzwxyz", + target_history=True, ), ] ) @@ -1071,6 +1126,10 @@ def _build_human_cartesian_revert_eef_frame_transform_list( pose_shape = 7 if is_quat else 6 mode = "xyzwxyz" if is_quat else "xyzypr" transform_list = [ + # The proprio defines the frame the action chunk is projected out of; + # with proprio_history it arrives (K, D), so take the CURRENT step + # (no-op at K = 1, and idempotent when a 6D revert already did it). + SelectCurrentStep(keys=[obs_key]), SplitKeys( input_key=obs_key, output_key_list=[ @@ -1123,6 +1182,7 @@ def _build_human_cartesian_revert_6d_transform_list( grip-padded 18 -> 20) is reverted the same way. """ return [ + SelectCurrentStep(keys=[obs_key]), # proprio_history: keep the current step CartesianRot6DToYPR(action_key=action_key), CartesianRot6DToYPR(action_key=obs_key), UnpadGripperZeros(action_key=obs_key), @@ -1144,6 +1204,7 @@ def _build_human_cartesian_revert_6d_wristframe_transform_list( actions back into head frame using that ypr proprio to define the frame. """ return [ + SelectCurrentStep(keys=[obs_key]), # proprio_history: keep the current step CartesianRot6DToYPR(action_key=action_key), CartesianRot6DToYPR(action_key=obs_key), UnpadGripperZeros(action_key=obs_key), @@ -1162,6 +1223,7 @@ def _build_human_keypoints_revert_6d_transform_list( changes). The proprio is reverted the same way. """ return [ + SelectCurrentStep(keys=[obs_key]), # proprio_history: keep the current step KeypointsRot6DToYPR(action_key=action_key), KeypointsRot6DToYPR(action_key=obs_key), ] @@ -1180,6 +1242,7 @@ def _build_human_keypoints_revert_6d_wristframe_transform_list( ``actions_keypoints`` (the wrist-pose slices are consumed by the revert). """ return [ + SelectCurrentStep(keys=[obs_key]), # proprio_history: keep the current step KeypointsRot6DToYPR(action_key=action_key), KeypointsRot6DToYPR(action_key=obs_key), *_build_human_keypoints_revert_eef_frame_transform_list( @@ -1272,17 +1335,22 @@ def _build_human_cartesian_eef_frame_transform_list( mode="xyzwxyz", ), InterpolatePadMask(new_chunk_length=chunk_length, stride=stride), + # target_history: the proprio obs pose may carry a leading time axis + # (proprio_history); the action chunk is expressed in the CURRENT + # step's wrist frame either way. ActionChunkCoordinateFrameTransform( target_world=left_obs_headframe, chunk_world=left_action_headframe, transformed_key_name=left_action_wristframe, mode="xyzwxyz", + target_history=True, ), ActionChunkCoordinateFrameTransform( target_world=right_obs_headframe, chunk_world=right_action_headframe, transformed_key_name=right_action_wristframe, mode="xyzwxyz", + target_history=True, ), XYZWXYZ_to_XYZYPR( keys=[ diff --git a/egomimic/rldb/zarr/action_chunk_transforms.py b/egomimic/rldb/zarr/action_chunk_transforms.py index eba43193b7..dade43a4d8 100644 --- a/egomimic/rldb/zarr/action_chunk_transforms.py +++ b/egomimic/rldb/zarr/action_chunk_transforms.py @@ -201,6 +201,7 @@ def __init__( extra_batch_key: dict = None, mode: Literal["xyz", "xyzwxyz", "xyzypr"] = "xyzwxyz", inverse: bool = True, + target_history: bool = False, ): """ args: @@ -208,6 +209,10 @@ def __init__( chunk_world: transformed_key_name: is_quat: if True, inputs are xyz + quat(wxyz); otherwise xyz + ypr. + target_history: the target frame is a PROPRIO key that may be read + with history, ``(K, D)`` with the current frame last. Take that + current frame; without this flag a 2-D target is an error (a + target frame is a single pose). """ self.target_world = target_world self.chunk_world = chunk_world @@ -215,6 +220,7 @@ def __init__( self.extra_batch_key = extra_batch_key self.mode = mode self.inverse = inverse + self.target_history = target_history def transform(self, batch): """ @@ -237,6 +243,11 @@ def transform(self, batch): chunk_world = np.asarray(batch[self.chunk_world]) chunk_world_shape = None + if self.target_history and target_world.ndim > 1: + # Proprio read with history: (K, D), current frame last. The chunk + # is always expressed in the CURRENT step's frame. + target_world = target_world[-1] + if chunk_world.ndim > 2: chunk_world_shape = chunk_world.shape chunk_world = chunk_world.reshape(-1, chunk_world_shape[-1]) @@ -288,7 +299,11 @@ def transform(self, batch): class QuaternionPoseToYPR(Transform): - """Convert a single pose from xyz + quat(x,y,z,w) to xyz + ypr.""" + """Convert a single pose from xyz + quat(x,y,z,w) to xyz + ypr. + + A leading time axis (``(K, 7)``, a proprio key read with history) is + treated as a batch axis and comes back as ``(K, 6)``. + """ def __init__(self, pose_key: str, output_key: str): self.pose_key = pose_key @@ -296,15 +311,17 @@ def __init__(self, pose_key: str, output_key: str): def transform(self, batch: dict) -> dict: pose = np.asarray(batch[self.pose_key]) - if pose.shape != (7,): + if pose.ndim > 2 or pose.shape[-1] != 7: raise ValueError( - f"QuaternionPoseToYPR expects shape (7,), got {pose.shape} for key " - f"'{self.pose_key}'" + f"QuaternionPoseToYPR expects shape (7,) or (K, 7), got " + f"{pose.shape} for key '{self.pose_key}'" ) - xyz = pose[:3] - xyzw = wxyz_to_xyzw(pose[3:7]) + flat = pose.reshape(-1, 7) + xyz = flat[:, :3] + xyzw = wxyz_to_xyzw(flat[:, 3:7]) ypr = R.from_quat(xyzw).as_euler("ZYX", degrees=False) - batch[self.output_key] = np.concatenate([xyz, ypr], axis=0) + out = np.concatenate([xyz, ypr], axis=-1) + batch[self.output_key] = out.reshape(*pose.shape[:-1], 6) return batch @@ -380,6 +397,7 @@ def __init__( pose_world: str, transformed_key_name: str, mode: Literal["xyzwxyz", "xyzypr", "xyz"] = "xyzwxyz", + target_history: bool = False, ): self.target_world = target_world self.pose_world = pose_world @@ -390,6 +408,7 @@ def __init__( chunk_world=pose_world, transformed_key_name=transformed_key_name, mode=mode, + target_history=target_history, ) def transform(self, batch: dict) -> dict: @@ -770,6 +789,27 @@ def transform(self, batch: dict) -> dict: return batch +class SelectCurrentStep(Transform): + """Drop a proprio key's observation-history axis, keeping the CURRENT step. + + A proprio key read with ``proprio_history`` K > 1 is ``(K, D)`` per sample + with the current frame last. The revert / viz pipelines want the single + current pose the forward pipeline used as the frame, so they start by + collapsing it. A ``(D,)`` value is left alone, so the K = 1 path is + bit-for-bit unchanged and the transform is idempotent. + """ + + def __init__(self, keys: list[str]): + self.keys = list(keys) + + def transform(self, batch: dict) -> dict: + for key in self.keys: + value = batch.get(key) + if value is not None and value.ndim == 2: + batch[key] = value[-1] + return batch + + class ConcatKeys(Transform): def __init__(self, key_list, new_key_name, delete_old_keys=False): self.key_list = list(key_list) @@ -851,13 +891,22 @@ def transform(self, batch: dict) -> dict: class Reshape(Transform): + """Reshape a key to ``shape``. A value holding a whole multiple of + ``shape`` (a proprio key read with history, ``(K, ...)``) keeps its leading + time axis and is reshaped per step. + """ + def __init__(self, input_key: str, output_key: str, shape: tuple): self.input_key = input_key self.output_key = output_key self.shape = shape def transform(self, batch: dict) -> dict: - batch[self.output_key] = batch[self.input_key].reshape(*self.shape) + value = batch[self.input_key] + n = value.numel() if isinstance(value, torch.Tensor) else np.asarray(value).size + target = int(np.prod(self.shape)) + shape = self.shape if n == target else (-1, *self.shape) + batch[self.output_key] = value.reshape(*shape) return batch diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index e9f585715c..42a758bf7d 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -2118,6 +2118,24 @@ def _action_pad_mask(self, idx: int) -> np.ndarray | None: mask[: max(n_real, 0)] = 1.0 return mask + @staticmethod + def _history_window(array: np.ndarray, history: int, stride: int) -> np.ndarray: + """``history`` steps from a contiguous backward window: the current + frame and every ``stride``-th frame before it, in order, front-padded + by repeating the oldest real frame so the current frame is always + ``[-1]``. The mirror of ``_pad_sequences``, which repeat-pads the TAIL + of a forward chunk. + + Only a window whose sole real frame is the current one (``idx < + stride``) pads to exactly a history-dropout sample; a longer short + window repeats its oldest real frame, which dropout never produces. + """ + array = array[::-1][::stride][::-1] + if array.shape[0] >= history: + return array + padding = np.repeat(array[:1], history - array.shape[0], axis=0) + return np.concatenate([padding, array], axis=0) + def _pad_sequences(self, data, horizon: int | None) -> dict: if horizon is None: return data @@ -2174,6 +2192,8 @@ def _next(reason: str, key: str = "") -> int: zarr_key = self.key_map[k]["zarr_key"] key_type = self.key_map[k].get("key_type", None) horizon = self.key_map[k].get("horizon", None) + history = int(self.key_map[k].get("history") or 1) + stride = int(self.key_map[k].get("history_stride") or 1) if key_type == "annotation_keys": data[k] = self._annotation_text_for_frame(idx) @@ -2182,12 +2202,22 @@ def _next(reason: str, key: str = "") -> int: if horizon is not None: end_idx = self._chunk_end_idx(idx, horizon, key_type) read_interval = (idx, end_idx) + elif history > 1: + # Backward window ending at (and including) the current + # frame, every `history_stride`-th step; clamped at the + # episode start and front-padded below, so the key is + # always (K, D) with `[..., -1]` the single frame a K = 1 + # keymap would have read. Read contiguously and thin after: + # the zarr chunks are contiguous anyway. + read_interval = (max(0, idx - stride * (history - 1)), idx + 1) else: read_interval = (idx, None) read_dict = {zarr_key: read_interval} raw_data = self.episode_reader.read(read_dict) self._pad_sequences(raw_data, horizon) # should be able to pad images data[k] = raw_data[zarr_key] + if horizon is None and history > 1: + data[k] = self._history_window(data[k], history, stride) if zarr_key in self._image_keys: jpeg_bytes = data[k] diff --git a/tests/unit/test_proprio_history.py b/tests/unit/test_proprio_history.py new file mode 100644 index 0000000000..59e62bf399 --- /dev/null +++ b/tests/unit/test_proprio_history.py @@ -0,0 +1,847 @@ +"""Proprio observation history: the last K proprio vectors, one token per step. + +K = 1 must stay bit-for-bit today's pipeline -- same sample keys, same shapes, +same stem tokens, same sinusoidal PE -- so every check below has a K = 1 twin. +For K > 1 the dataset reads a backward window (front-padded with the first real +frame so the current frame is always ``[-1]``), the cartesian transforms treat +the leading axis as batch, the baseline's per-channel proprio stats broadcast +over it, and the stem replaces the input-side sinusoid with a learned +time-step embedding added after the projection. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +import torch +from omegaconf import OmegaConf + +import egomimic +from egomimic.algo.hpt import HPT, HPTModel +from egomimic.models.hpt_nets import MLPPolicyStem +from egomimic.rldb.embodiment.embodiment import get_embodiment_id +from egomimic.rldb.embodiment.human import Human +from egomimic.rldb.zarr.zarr_dataset_multi import ( + LocalEpisodeResolver, + MultiDataset, +) +from egomimic.rldb.zarr.zarr_writer import ZarrWriter + +T_FRAMES = 24 +HORIZON = Human.ACTION_HORIZON +CHUNK = 100 +PROPRIO_KEYS = ("left.obs_ee_pose", "right.obs_ee_pose") + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _pose_series(T: int, offset: float) -> np.ndarray: + """(T, 7) xyz+quat(wxyz) whose translation is a distinct ramp per frame. + + The shared synthetic fixture holds every pose CONSTANT over time (so the + per-cell quantile bounds stay usable); history needs frames that differ, + so this writes its own episode with an identity rotation and a ramp. + """ + t = np.arange(T, dtype=np.float64)[:, None] + xyz = offset + t * np.array([[0.001, 0.002, 0.003]]) + quat = np.tile(np.array([[1.0, 0.0, 0.0, 0.0]]), (T, 1)) + return np.concatenate([xyz, quat], axis=1) + + +def _write_varying_episode(root, *, T: int = T_FRAMES, seed: int = 0): + rng = np.random.default_rng(seed) + H = W = 64 + K = np.array( + [[100.0, 0.0, W / 2, 0.0], [0.0, 100.0, H / 2, 0.0], [0.0, 0.0, 1.0, 0.0]] + ) + numeric = { + "left.obs_ee_pose": _pose_series(T, 0.10 + seed), + "right.obs_ee_pose": _pose_series(T, 0.20 + seed), + "obs_head_pose": _pose_series(T, 0.30 + seed), + } + images = {"images.front_1": rng.integers(0, 255, (T, H, W, 3), dtype=np.uint8)} + return ZarrWriter.create_and_write( + root / f"varying_{seed:02d}.zarr", + numeric_data=numeric, + image_data=images, + embodiment="human_bimanual", + fps=30, + task_name="synthetic", + intrinsics={"front_1": K}, + operator="op_00", + lab="synthetic", + scene="test_scene", + objects=["cube"], + vendor_recording_id="synthetic_session", + ) + + +def _leaf(ds: MultiDataset): + return next(iter(ds.datasets.values())) + + +def _dataset( + tmp_path, + *, + proprio_history: int = 1, + history_stride: int = 1, + transforms=None, + n: int = 1, +): + for i in range(n): + _write_varying_episode(tmp_path, seed=i) + key_map = Human.get_keymap( + keymap_mode="cartesian", + proprio_history=proprio_history, + history_stride=history_stride, + ) + resolver = LocalEpisodeResolver( + tmp_path, key_map=key_map, transform_list=transforms or [] + ) + return MultiDataset._from_resolver(resolver, mode="total") + + +def _raw_poses(tmp_path, key: str) -> np.ndarray: + """The written (T, 7) array for ``key``, straight off disk.""" + import zarr + + store = zarr.open_group(str(next(tmp_path.glob("*.zarr"))), mode="r") + return np.asarray(store[key][:]) + + +# --------------------------------------------------------------------------- +# 1. dataset: backward window, front padding +# --------------------------------------------------------------------------- + + +def test_history_window_is_the_last_k_frames(tmp_path): + leaf = _leaf(_dataset(tmp_path, proprio_history=4)) + for key in PROPRIO_KEYS: + raw = _raw_poses(tmp_path, key) + got = np.asarray(leaf[10][key]) + assert got.shape == (4, 7) + np.testing.assert_allclose(got, raw[7:11], rtol=0, atol=1e-6) + # The current frame is always last. + np.testing.assert_allclose(got[-1], raw[10], rtol=0, atol=1e-6) + + +def test_history_stride_spaces_the_window_out(tmp_path): + """K = 3, s = 5 reads frames idx-10, idx-5, idx -- 0.33 s of past at 30 fps + instead of 0.07 s, with the current frame still last.""" + leaf = _leaf(_dataset(tmp_path, proprio_history=3, history_stride=5)) + for key in PROPRIO_KEYS: + raw = _raw_poses(tmp_path, key) + got = np.asarray(leaf[12][key]) + assert got.shape == (3, 7) + np.testing.assert_allclose(got, raw[[2, 7, 12]], rtol=0, atol=1e-6) + + +def test_strided_history_front_pads_at_the_episode_start(tmp_path): + leaf = _leaf(_dataset(tmp_path, proprio_history=3, history_stride=5)) + raw = _raw_poses(tmp_path, PROPRIO_KEYS[0]) + # idx 7: frames 7 and 2 are real, idx-10 would be negative -> the oldest + # real frame is repeated + got = np.asarray(leaf[7][PROPRIO_KEYS[0]]) + np.testing.assert_allclose(got, raw[[2, 2, 7]], rtol=0, atol=1e-6) + # idx 3: only the current frame is real -> exactly a history-dropout sample + got = np.asarray(leaf[3][PROPRIO_KEYS[0]]) + np.testing.assert_allclose(got, raw[[3, 3, 3]], rtol=0, atol=1e-6) + + +def test_stride_one_is_the_contiguous_window(tmp_path): + a = _leaf(_dataset(tmp_path, proprio_history=4)) + b = _leaf(_dataset(tmp_path, proprio_history=4, history_stride=1)) + for key in PROPRIO_KEYS: + np.testing.assert_array_equal(np.asarray(a[10][key]), np.asarray(b[10][key])) + + +def test_history_front_pads_by_repeating_the_first_real_frame(tmp_path): + leaf = _leaf(_dataset(tmp_path, proprio_history=4)) + for key in PROPRIO_KEYS: + raw = _raw_poses(tmp_path, key) + got = np.asarray(leaf[1][key]) + assert got.shape == (4, 7) + expected = np.stack([raw[0], raw[0], raw[0], raw[1]]) + np.testing.assert_allclose(got, expected, rtol=0, atol=1e-6) + + +def test_k1_sample_is_bit_for_bit_todays_sample(tmp_path): + """The K = 1 keymap must produce exactly what a keymap with no + ``proprio_history`` argument at all produces.""" + _write_varying_episode(tmp_path, seed=0) + baseline = Human.get_keymap(keymap_mode="cartesian") + with_arg = Human.get_keymap(keymap_mode="cartesian", proprio_history=1) + assert with_arg == baseline + assert not any("history" in spec for spec in with_arg.values()) + + def _sample(key_map): + resolver = LocalEpisodeResolver(tmp_path, key_map=key_map, transform_list=[]) + ds = MultiDataset._from_resolver(resolver, mode="total") + return _leaf(ds)[5] + + a, b = _sample(baseline), _sample(with_arg) + assert set(a) == set(b) + for key in PROPRIO_KEYS: + assert tuple(a[key].shape) == (7,) + for key, value in a.items(): + if isinstance(value, torch.Tensor): + torch.testing.assert_close(value, b[key]) + + +def test_keys_without_history_keep_the_single_frame_read(tmp_path): + """Only the proprio keys the model consumes get a time axis: the head pose + (the frame everything is expressed in) and the action chunks do not.""" + leaf = _leaf(_dataset(tmp_path, proprio_history=4)) + sample = leaf[10] + assert tuple(sample["obs_head_pose"].shape) == (7,) + assert tuple(sample["left.action_ee_pose"].shape) == (HORIZON, 7) + + +# --------------------------------------------------------------------------- +# 2. transforms: a leading time axis is a batch axis +# --------------------------------------------------------------------------- + + +def _wristframe_batch(poses: dict[str, np.ndarray]) -> dict: + rng = np.random.default_rng(0) + chunk = _pose_series(HORIZON, 0.4) + batch = { + "obs_head_pose": _pose_series(1, 0.3)[0], + "left.action_ee_pose": chunk.copy(), + "right.action_ee_pose": chunk.copy() + 0.01, + "action_pad_mask": np.ones(HORIZON, dtype=np.float32), + } + del rng + batch.update({k: v.copy() for k, v in poses.items()}) + return batch + + +def _run_wristframe(batch: dict) -> dict: + for t in Human.get_transform_list( + mode="cartesian_wristframe_6d", + stride=1, + fix_left_wrist_convention=True, + pad_proprio_gripper=True, + ): + batch = t.transform(batch) + return batch + + +def test_wristframe_chain_maps_k_by_7_proprio_to_k_by_20(): + K = 4 + hist = {key: _pose_series(K, 0.1 * (i + 1)) for i, key in enumerate(PROPRIO_KEYS)} + out = _run_wristframe(_wristframe_batch(hist)) + state = np.asarray(out["observations.state.ee_pose"]) + assert state.shape == (K, 20) + + current = {key: value[-1] for key, value in hist.items()} + out1 = _run_wristframe(_wristframe_batch(current)) + state1 = np.asarray(out1["observations.state.ee_pose"]) + assert state1.shape == (20,) + np.testing.assert_allclose(state[-1], state1, rtol=1e-6, atol=1e-8) + + +def test_wristframe_actions_are_unchanged_by_the_proprio_time_axis(): + """The action chunk is expressed in the CURRENT step's wrist frame, so it + must not move when history is added next to it.""" + K = 4 + hist = {key: _pose_series(K, 0.1 * (i + 1)) for i, key in enumerate(PROPRIO_KEYS)} + current = {key: value[-1] for key, value in hist.items()} + a = _run_wristframe(_wristframe_batch(hist))["actions_cartesian"] + b = _run_wristframe(_wristframe_batch(current))["actions_cartesian"] + np.testing.assert_allclose(np.asarray(a), np.asarray(b), rtol=1e-6, atol=1e-8) + + +def test_dataset_and_transforms_compose_to_a_k_by_20_state(tmp_path): + transforms = Human.get_transform_list( + mode="cartesian_wristframe_6d", + stride=1, + fix_left_wrist_convention=True, + pad_proprio_gripper=True, + ) + hist = _leaf(_dataset(tmp_path, proprio_history=4, transforms=transforms)) + base = _leaf(_dataset(tmp_path, proprio_history=1, transforms=transforms)) + for idx in (0, 3, 11): + h = hist[idx]["observations.state.ee_pose"] + b = base[idx]["observations.state.ee_pose"] + assert tuple(h.shape) == (4, 20) + assert tuple(b.shape) == (20,) + torch.testing.assert_close(h[-1], b, rtol=1e-5, atol=1e-6) + + +# --------------------------------------------------------------------------- +# 3. normalization / bounds over the time axis +# --------------------------------------------------------------------------- + + +def _norm_shell(key: str, width: int, *, lo=-1.0, hi=1.0): + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + md.norm_stats = { + 0: { + key: { + "quantile_1": np.full(width, lo, dtype=np.float32), + "quantile_99": np.full(width, hi, dtype=np.float32), + } + } + } + md.zarr_keys = {0: {key: key}} + md.key_types = {0: {key: "proprio_keys"}} + md._warned_violations = set() + return md + + +def test_per_channel_stats_broadcast_over_the_time_axis(): + key = "observations.state.ee_pose" + md = _norm_shell(key, 20, lo=-2.0, hi=4.0) + stats = md.norm_stats[0][key] + g = torch.Generator().manual_seed(0) + arr = torch.randn(4, 20, generator=g) + out = md._apply_norm_one(arr, stats) + assert out.shape == (4, 20) + for k in range(4): + torch.testing.assert_close(out[k], md._apply_norm_one(arr[k], stats)) + + +def test_bounds_check_catches_a_violation_in_a_past_step(): + key = "observations.state.ee_pose" + md = _norm_shell(key, 20) + arr = np.zeros((4, 20), dtype=np.float32) + assert md._check_bounds({"embodiment": 0, key: arr.copy()}, None, 0, "ep") is None + bad = arr.copy() + bad[0, 0] = 50.0 # oldest history step, translation channel + assert md._check_bounds({"embodiment": 0, key: bad}, None, 0, "ep") is not None + + +# --------------------------------------------------------------------------- +# 4. the stem +# --------------------------------------------------------------------------- + + +def _stems(D: int, out: int, **kwargs): + torch.manual_seed(0) + a = MLPPolicyStem(input_dim=D, output_dim=out, widths=[16], **kwargs) + torch.manual_seed(0) + b = MLPPolicyStem(input_dim=D, output_dim=out, widths=[16]) + return a, b + + +def test_history_stem_shape_and_zero_init_matches_the_baseline_stem(): + D, out, B, K = 5, 7, 3, 3 + hist, base = _stems(D, out, history_len=K) + x = torch.randn(B, K, 1, D) + y = hist(x) + assert y.shape == (B, K, 1, out) + # time_embed is zero at init, so every step is today's projection. + torch.testing.assert_close(y, base(x)) + for k in range(K): + torch.testing.assert_close(y[:, k], base(x[:, k])) + + +def test_time_embed_is_learned_zero_init_and_per_step(): + D, out, K = 5, 7, 3 + hist, _ = _stems(D, out, history_len=K) + assert hist.time_embed.shape == (K, out) + assert hist.time_embed.requires_grad + torch.testing.assert_close(hist.time_embed, torch.zeros(K, out)) + + with torch.no_grad(): + hist.time_embed.copy_( + torch.arange(K, dtype=torch.float32)[:, None].expand(K, out) + ) + x = torch.randn(2, K, 1, D) + y = hist(x) + assert not torch.allclose(y[:, 0], y[:, 1]) + # ... and the shift is exactly the per-step embedding. + hist.time_embed.data.zero_() + y0 = hist(x) + hist.time_embed.data.copy_( + torch.arange(K, dtype=torch.float32)[:, None].expand(K, out) + ) + torch.testing.assert_close( + hist(x) - y0, + torch.arange(K, dtype=torch.float32)[None, :, None, None].expand(2, K, 1, out), + ) + + +def test_history_len_one_stem_has_no_time_embed_and_is_todays_stem(): + D, out = 5, 7 + one, base = _stems(D, out, history_len=1) + assert not hasattr(one, "time_embed") + assert dict(one.named_parameters()).keys() == dict(base.named_parameters()).keys() + x = torch.randn(4, 1, 1, D) + torch.testing.assert_close(one(x), base(x)) + + +def test_history_dropout_replaces_the_past_with_the_current_step_in_train_mode(): + D, out, B, K = 5, 7, 6, 4 + hist, _ = _stems(D, out, history_len=K, history_dropout=1.0) + x = torch.randn(B, K, 1, D) + hist.train() + y = hist(x) + # Every step equals the current one: exactly an episode-start sample after + # front padding. + for k in range(K): + torch.testing.assert_close( + y[:, k], y[:, -1] - hist.time_embed[-1] + hist.time_embed[k] + ) + hist.eval() + torch.testing.assert_close(y, hist(_current_only(x))) + + +def _current_only(x: torch.Tensor) -> torch.Tensor: + out = x.clone() + out[:, :-1] = out[:, -1:] + return out + + +def test_history_dropout_is_a_noop_in_eval_mode(): + D, out, B, K = 5, 7, 6, 4 + drop, _ = _stems(D, out, history_len=K, history_dropout=1.0) + torch.manual_seed(0) + plain = MLPPolicyStem(input_dim=D, output_dim=out, widths=[16], history_len=K) + drop.eval() + plain.eval() + x = torch.randn(B, K, 1, D) + torch.testing.assert_close(drop(x), plain(x)) + assert not torch.allclose(drop(x), drop(_current_only(x))) + + +def test_history_dropout_zero_is_a_noop_in_train_mode(): + D, out, B, K = 5, 7, 6, 4 + a, _ = _stems(D, out, history_len=K, history_dropout=0.0) + torch.manual_seed(0) + b = MLPPolicyStem(input_dim=D, output_dim=out, widths=[16], history_len=K) + a.train() + b.eval() + x = torch.randn(B, K, 1, D) + torch.testing.assert_close(a(x), b(x)) + + +def test_history_stem_rejects_a_wrong_length_history(): + hist, _ = _stems(5, 7, history_len=3) + with pytest.raises(ValueError, match="3"): + hist(torch.randn(2, 4, 1, 5)) + + +# --------------------------------------------------------------------------- +# 5. stem_process: the input-side sinusoid goes once history_len > 1 +# --------------------------------------------------------------------------- + + +class _EchoStem: + """Stem stub whose ``compute_latent`` hands back what stem_process fed it.""" + + def __init__(self, history_len: int = 1): + self.history_len = history_len + self.seen = None + + def compute_latent(self, x): + self.seen = x.clone() + return x + + +def _stem_process(history_len: int, data_shape): + model = HPTModel.__new__(HPTModel) + model.modalities = {"human_bimanual": ["state_ee_pose"]} + model.shared_keys = [] + model.encoders = {} + stem = _EchoStem(history_len) + model.stems = {"human_bimanual_state_ee_pose": stem} + model.stem_spec = { + "human_bimanual": { + "state_ee_pose": OmegaConf.create( + {"specs": {"random_horizon_masking": False}} + ) + } + } + model.train_mode = False + torch.manual_seed(0) + x = torch.randn(*data_shape) + model.stem_process("human_bimanual", {"state_ee_pose": x.clone()}) + return x, stem.seen + + +def test_stem_process_skips_the_sinusoid_for_a_history_stem(): + x, seen = _stem_process(4, (2, 4, 20)) + torch.testing.assert_close(seen, x) + + +def test_stem_process_keeps_the_sinusoid_for_a_history_len_one_stem(): + x, seen = _stem_process(1, (2, 1, 20)) + from egomimic.utils.tensor_utils import get_sinusoid_encoding_table + + pe = get_sinusoid_encoding_table(0, 1, 20).to(x) + torch.testing.assert_close(seen, x + pe.view(1, 1, 20)) + assert not torch.allclose(seen, x) + + +# --------------------------------------------------------------------------- +# 6. HPT batch plumbing +# --------------------------------------------------------------------------- + + +class _NoRenameNormStats: + def zarr_key_to_keyname(self, zarr_key, embodiment_id): + return None + + +def _hpt_with_stem(history_len: int): + algo = HPT.__new__(HPT) + policy = HPTModel.__new__(HPTModel) + policy.stems = {"human_bimanual_state_ee_pose": _EchoStem(history_len)} + algo.nets = {"policy": policy} + algo.is_6dof = True + algo.shared_ac_key = None + algo.annotation_modality = "annotation" + algo.eval_image_augs = None + return algo + + +def _to_hpt(algo, proprio: torch.Tensor): + batch = { + "observations.state.ee_pose": proprio, + "pad_mask": torch.ones(proprio.shape[0], 5, 1), + "embodiment": torch.tensor([0]), + "actions_cartesian": torch.zeros(proprio.shape[0], 5, 18), + } + return HPT._robomimic_to_hpt_data( + algo, + batch, + cam_keys=[], + proprio_keys=["observations.state.ee_pose"], + lang_keys=[], + ac_key="actions_cartesian", + domain="human_bimanual", + ) + + +def test_two_dim_proprio_is_still_unsqueezed(): + algo = _hpt_with_stem(1) + data = _to_hpt(algo, torch.randn(3, 20)) + assert data["state_ee_pose"].shape == (3, 1, 20) + + +def test_three_dim_proprio_passes_through_without_unsqueeze(): + algo = _hpt_with_stem(4) + x = torch.randn(3, 4, 20) + data = _to_hpt(algo, x) + assert data["state_ee_pose"].shape == (3, 4, 20) + torch.testing.assert_close(data["state_ee_pose"], x) + + +def test_history_len_is_read_from_the_domain_not_by_suffix(): + """A cotrain run has one state_ee_pose stem per domain; a suffix match + returns whichever comes first in the dict, so the human side would be + checked against eva's history_len.""" + algo = _hpt_with_stem(3) + algo.nets["policy"].stems = { + "eva_bimanual_state_ee_pose": _EchoStem(1), + "human_bimanual_state_ee_pose": _EchoStem(3), + } + data = _to_hpt(algo, torch.randn(3, 3, 20)) + assert data["state_ee_pose"].shape == (3, 3, 20) + + +def test_history_length_mismatch_raises_naming_both(): + algo = _hpt_with_stem(2) + with pytest.raises(ValueError, match="state_ee_pose"): + _to_hpt(algo, torch.randn(3, 4, 20)) + + +def _process(action_pad_mask=None, S: int = 6): + algo = HPT.__new__(HPT) + algo.norm_stats = _NoRenameNormStats() + algo.device = "cpu" + algo.annotation_key = None + algo.use_pad_mask = False + emb = "human_bimanual" + emb_id = get_embodiment_id(emb) + algo.ac_keys = {emb_id: "actions_cartesian"} + inner = { + "actions_cartesian": torch.zeros(2, S, 18), + "observations.state.ee_pose": torch.zeros(2, 4, 20), + } + if action_pad_mask is not None: + inner["action_pad_mask"] = action_pad_mask + return HPT.process_batch_for_training(algo, {emb: inner})[emb_id] + + +def test_process_batch_keeps_the_history_axis(): + out = _process() + assert out["observations.state.ee_pose"].shape == (2, 4, 20) + torch.testing.assert_close(out["pad_mask"], torch.ones(2, 6, 1)) + + +def test_norm_mode_keymap_reads_the_current_step_only(): + """Norm stats are current-step stats for every K; get_keymap owns that.""" + km = Human.get_keymap(keymap_mode="cartesian", proprio_history=3, norm_mode=True) + assert km == Human.get_keymap(keymap_mode="cartesian", norm_mode=True) + + +# --------------------------------------------------------------------------- +# 7. configs +# --------------------------------------------------------------------------- + +RECIPE = "train_zarr_mecka_flagship_6d_hpt" +DATA = "data=mecka_fold_flagship_opsplit_hpt_6d" +KEYMAP = "data.train_datasets.human_bimanual.resolver.key_map" +STEM = "model.robomimic_model.stem_specs.human_bimanual.state_ee_pose" + + +def test_flagship_ships_k_three_stride_five(compose_resolve): + cfg = compose_resolve(RECIPE, [DATA]) + km = cfg.data.train_datasets.human_bimanual.resolver.key_map + assert km.proprio_history == 3 + assert km.history_stride == 5 + stem = cfg.model.robomimic_model.stem_specs.human_bimanual.state_ee_pose + assert stem.history_len == km.proprio_history, "data K and stem K must agree" + assert stem.history_dropout == 0.2 + + +def test_k_one_baseline_is_still_reachable(compose_resolve): + cfg = compose_resolve( + RECIPE, + [ + DATA, + f"{KEYMAP}.proprio_history=1", + f"{KEYMAP}.history_stride=1", + f"{STEM}.history_len=1", + ], + ) + km = cfg.data.train_datasets.human_bimanual.resolver.key_map + assert km.proprio_history == 1 and km.history_stride == 1 + assert ( + cfg.model.robomimic_model.stem_specs.human_bimanual.state_ee_pose.history_len + == 1 + ) + + +def test_k_two_override_reaches_every_loader(compose_resolve): + cfg = compose_resolve( + RECIPE, [DATA, f"{KEYMAP}.proprio_history=2", f"{STEM}.history_len=2"] + ) + data = cfg.data + for group in ("train_datasets", "valid_datasets", "unseen_op_valid_datasets"): + km = data[group]["human_bimanual"]["resolver"]["key_map"] + assert km["proprio_history"] == 2, group + assert ( + cfg.model.robomimic_model.stem_specs.human_bimanual.state_ee_pose.history_len + == 2 + ) + + +def test_k_two_keymap_marks_only_the_model_proprio_keys(): + km = Human.get_keymap(keymap_mode="cartesian", proprio_history=2) + assert {k for k, v in km.items() if "history" in v} == set(PROPRIO_KEYS) + assert all(km[k]["history"] == 2 for k in PROPRIO_KEYS) + + +# --------------------------------------------------------------------------- +# 8. the eval path: the revert transforms and the evaluator take the current +# step, so val metrics and the overlay video work at K > 1 +# --------------------------------------------------------------------------- + +REVERT_BUILDERS = [ + "_build_human_cartesian_revert_6d_wristframe_transform_list", + "_build_human_cartesian_revert_6d_transform_list", +] + + +def _revert_inputs(B: int = 2, K: int = 3): + """``(actions, current_proprio, history_proprio)`` whose last history step + is exactly the 2-D proprio.""" + g = torch.Generator().manual_seed(7) + actions = torch.randn(B, 100, 18, generator=g) + current = torch.randn(B, 20, generator=g) + history = torch.randn(B, K, 20, generator=g) + history[:, -1] = current + return actions, current, history + + +@pytest.mark.parametrize("builder_name", REVERT_BUILDERS) +def test_revert_transform_list_is_blind_to_proprio_history(builder_name): + from egomimic.rldb.embodiment import human as human_mod + from egomimic.rldb.embodiment.embodiment import Embodiment + + transform_list = getattr(human_mod, builder_name)() + actions, current, history = _revert_inputs() + + def _run(proprio): + return Embodiment.apply_transform( + { + "actions_cartesian": actions.clone(), + "observations.state.ee_pose": proprio.clone(), + }, + transform_list, + ) + + flat, hist = _run(current), _run(history) + assert flat["observations.state.ee_pose"].shape == (2, 12) + assert hist["observations.state.ee_pose"].shape == (2, 12) + for key in ("actions_cartesian", "observations.state.ee_pose"): + torch.testing.assert_close(hist[key], flat[key], rtol=1e-5, atol=1e-6) + + +def test_select_current_step_leaves_a_flat_value_alone_and_is_idempotent(): + from egomimic.rldb.zarr.action_chunk_transforms import SelectCurrentStep + + t = SelectCurrentStep(keys=["observations.state.ee_pose", "missing"]) + flat = np.arange(20, dtype=np.float32) + out = t.transform({"observations.state.ee_pose": flat}) + assert out["observations.state.ee_pose"] is flat # untouched, not a copy + + hist = np.stack([flat - 1, flat, flat + 1]) + once = t.transform({"observations.state.ee_pose": hist.copy()}) + np.testing.assert_array_equal(once["observations.state.ee_pose"], flat + 1) + twice = t.transform(dict(once)) + np.testing.assert_array_equal(twice["observations.state.ee_pose"], flat + 1) + + +def test_every_human_revert_list_selects_the_current_proprio_step(): + from egomimic.rldb.embodiment import human as human_mod + from egomimic.rldb.zarr.action_chunk_transforms import SelectCurrentStep + + builders = { + "_build_human_cartesian_revert_6d_transform_list": "observations.state.ee_pose", + "_build_human_cartesian_revert_6d_wristframe_transform_list": "observations.state.ee_pose", + "_build_human_cartesian_revert_eef_frame_transform_list": "observations.state.ee_pose", + "_build_human_keypoints_revert_6d_transform_list": "observations.state.keypoints", + "_build_human_keypoints_revert_6d_wristframe_transform_list": "observations.state.keypoints", + "_build_human_keypoints_revert_eef_frame_transform_list": "observations.state.keypoints", + } + for name, obs_key in builders.items(): + tl = getattr(human_mod, name)() + assert isinstance(tl[0], SelectCurrentStep), name + assert tl[0].keys == [obs_key], name + + +# --- the evaluator itself ----------------------------------------------- + + +class _IdentityNormStats: + def unnormalize(self, batch, embodiment_id): + return dict(batch) + + +class _StubAlgo: + """Just enough of HPT for HPTEvalVideo.compute_metrics_and_viz.""" + + def __init__(self, emb_id, preds): + self.norm_stats = _IdentityNormStats() + self.ac_keys = {emb_id: "actions_cartesian"} + self.shared_ac_key = None + self.auxiliary_ac_keys = {} + self._preds = preds + + def forward_eval(self, batch): + return self._preds + + +def _eval_once(proprio): + """Run the real HPTEvalVideo over one batch, returning its metrics and the + batch the viz function was handed.""" + from egomimic.eval.eval_hpt import HPTEvalVideo + from egomimic.rldb.embodiment.human import ( + _build_human_cartesian_revert_6d_wristframe_transform_list, + ) + + emb_id = get_embodiment_id("human_bimanual") + emb_name = "human_bimanual" + # Action horizon 1: the evaluator's paired/final MSE does + # ``pred[:, -1].cpu()``, which is a real (contiguous) copy for the CUDA + # tensors a run produces but a non-contiguous view for the CPU tensors a + # test hands it, and torchmetrics' ``view(-1)`` rejects that. S = 1 keeps + # the slice contiguous; the proprio path under test is unaffected. + g = torch.Generator().manual_seed(11) + actions = torch.randn(2, 1, 18, generator=g) + preds = { + f"{emb_name}_actions_cartesian": torch.randn(2, 1, 18, generator=g), + f"{emb_name}_loss": torch.tensor(0.5), + } + batch = { + emb_id: { + "actions_cartesian": actions.clone(), + "observations.state.ee_pose": proprio.clone(), + "embodiment": torch.tensor([emb_id, emb_id]), + } + } + + seen = {} + + def _viz(predictions, viz_batch): + seen["batch"] = viz_batch + seen["preds"] = predictions + return np.zeros((2, 4, 4, 3), dtype=np.uint8) + + ev = HPTEvalVideo.__new__(HPTEvalVideo) + ev.model = _StubAlgo(emb_id, preds) + ev.viz_func = {emb_name: _viz} + ev.transform_lists = { + emb_name: _build_human_cartesian_revert_6d_wristframe_transform_list() + } + metrics, images = ev.compute_metrics_and_viz(batch, do_viz=True) + return metrics, images, seen + + +def test_evaluator_metrics_and_viz_are_blind_to_proprio_history(): + _, current, history = _revert_inputs() + m_flat, im_flat, seen_flat = _eval_once(current) + m_hist, im_hist, seen_hist = _eval_once(history) + + assert set(m_flat) == set(m_hist) + assert m_flat, "the evaluator must produce metrics" + for key in m_flat: + torch.testing.assert_close( + torch.as_tensor(m_flat[key]).float(), + torch.as_tensor(m_hist[key]).float(), + rtol=1e-5, + atol=1e-6, + ) + # The viz function sees the reverted, current-step proprio either way. + assert seen_flat["batch"]["observations.state.ee_pose"].shape == (2, 12) + torch.testing.assert_close( + seen_hist["batch"]["observations.state.ee_pose"], + seen_flat["batch"]["observations.state.ee_pose"], + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + seen_hist["preds"]["human_bimanual_actions_cartesian"], + seen_flat["preds"]["human_bimanual_actions_cartesian"], + rtol=1e-5, + atol=1e-6, + ) + assert set(im_flat) == set(im_hist) + + +TRAIN_CONFIGS = sorted( + p.stem + for p in (Path(egomimic.__file__).parent / "hydra_configs").glob("train_*.yaml") +) + + +@pytest.mark.parametrize("top", TRAIN_CONFIGS) +def test_every_recipe_agrees_with_its_stem_on_K(top, compose_resolve): + """The data's `proprio_history` and the stem's `history_len` must match, or + process_batch_for_training raises at the first training step. A recipe that + inherits a data config from one model and overrides the model is exactly + how they drift apart, and nothing else composes both halves.""" + cfg = compose_resolve(top, []) + stem_specs = (cfg.model.robomimic_model.get("stem_specs") or {}).get( + "human_bimanual" + ) or {} + for split in ("train_datasets", "valid_datasets"): + node = (cfg.data.get(split) or {}).get("human_bimanual") + if node is None or node.get("resolver") is None: + continue + k = int(node.resolver.key_map.get("proprio_history", 1) or 1) + for name, stem in stem_specs.items(): + if stem is None or "history_len" not in stem: + continue + assert int(stem.history_len) == k, (top, split, name, k, stem.history_len)