obs: proprio history (K proprio steps, one token each, learned time embedding)

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:

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

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:

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