6d: PI goes 6D-only; Eva and Human gain the continuous-rotation modes

CONTRIBUTING_DATA.md

@@ -644,7 +644,7 @@ If you are contributing egocentric human data, you use `human_bimanual` (or the
 There is a **single concrete `Human` embodiment class** ([`egomimic/rldb/embodiment/human.py`](egomimic/rldb/embodiment/human.py)) shared by all human data. You do **not** write a per-vendor subclass, and there is **no per-vendor embodiment identifier** — every human contributor uses `human_*` and records their source in the `lab` field. Camera intrinsics travel with the data (`zarr.attrs`, §6.3 / §6.4); per-vendor structural choices are passed as explicit arguments from the data config:
 
 - `Human.get_keymap(keymap_mode="cartesian"|"keypoints", has_head_pose=<bool>, include_aria_keypoints=<bool>)`
-- `Human.get_transform_list(mode="cartesian"|"keypoints_headframe_ypr"|..., stride=<int>)`
+- `Human.get_transform_list(mode="cartesian_6d"|"keypoints_headframe_ypr"|..., stride=<int>)`
 
 Onboarding human data is just two steps:
 
@@ -947,7 +947,7 @@ import torch
 HAS_HEAD_POSE = True
 key_map = Human.get_keymap(keymap_mode="cartesian", has_head_pose=HAS_HEAD_POSE)
 transform_list = (
-    Human.get_transform_list(mode="cartesian", stride=3) if HAS_HEAD_POSE else None
+    Human.get_transform_list(mode="cartesian_6d", stride=3) if HAS_HEAD_POSE else None
 )
 
 resolver = LocalEpisodeResolver(
@@ -987,7 +987,7 @@ from egomimic.rldb.filters import DatasetFilter
 resolver = LocalEpisodeResolver(
     folder_path    = "/local/processed",
     key_map        = Human.get_keymap(keymap_mode="cartesian"),
-    transform_list = Human.get_transform_list(mode="cartesian", stride=3),
+    transform_list = Human.get_transform_list(mode="cartesian_6d", stride=3),
 )
 filters = DatasetFilter(filter_lambdas=[
     "lambda row: row['episode_hash'] == '2026-03-15-14-22-10-000000'"

egomimic/algo/pi.py

@@ -421,13 +421,7 @@ def forward_training(self, batch):
             )
 
             losses = self.nets["policy"].forward(processed_obs, action)
-
-            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)
-
-            loss = losses.mean()
+            loss = self._reduce_loss(losses, _batch[ac_key], embodiment_id, ac_key)
 
             predictions[f"{embodiment_name}_{ac_key}"] = _batch[ac_key]
             predictions[f"{embodiment_name}_loss"] = loss
@@ -465,13 +459,9 @@ def forward_eval(self, batch):
 
                 # Flow-matching val loss — same call as forward_training.
                 losses = self.nets["policy"].forward(processed_obs, action)
-                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
-                    )
-                unnorm_preds[f"{embodiment_name}_loss"] = losses.mean()
+                unnorm_preds[f"{embodiment_name}_loss"] = self._reduce_loss(
+                    losses, _batch[ac_key], embodiment_id, ac_key
+                )
 
                 pred_actions = self.nets["policy"].sample_actions(
                     device=self.device,
@@ -480,22 +470,36 @@ def forward_eval(self, batch):
                     num_steps=self.num_steps,
                 )
 
-                predictions = OrderedDict()
-                ref = _batch[ac_key]
-                B, T, D = ref.shape
-
-                converter = self.action_registry.get(embodiment_id, ac_key)
-                pred_actions_orig = converter.from32(pred_actions)
-
-                pred = pred_actions_orig[:, :T, :D]
-                predictions[ac_key] = pred
-
-                unnorm_actions = self.norm_stats.unnormalize(predictions, embodiment_id)
+                unnorm_actions = self._postprocess_sampled_actions(
+                    pred_actions, _batch, embodiment_id, ac_key
+                )
                 for key in unnorm_actions:
                     unnorm_preds[f"{embodiment_name}_{key}"] = unnorm_actions[key]
 
         return unnorm_preds
 
+    def _reduce_loss(self, losses, action, embodiment_id, ac_key) -> torch.Tensor:
+        """Mean flow-matching loss for the embodiment."""
+        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)
+        return losses.mean()
+
+    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)."""
+        ref = _batch[ac_key]
+        _, T, D = ref.shape
+        converter = self.action_registry.get(embodiment_id, ac_key)
+        predictions = OrderedDict()
+        pred_native = converter.from32_norm_6d(pred_actions)
+        predictions[ac_key] = pred_native[:, :T, :D]
+        return self.norm_stats.unnormalize(predictions, embodiment_id)
+
     @override
     def compute_losses(self, predictions, batch):
         """
@@ -555,7 +559,10 @@ def _robomimic_to_pi_data(self, batch, proprio_keys, lang_keys, ac_key, embodime
 
         emb_id = get_embodiment_id(embodiment)  # embodiment is a name string
         converter = self.action_registry.get(emb_id, ac_key)
-        action32 = converter.to32(action)
+        # 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)
 
         # 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
@@ -630,40 +637,3 @@ def _clone_batch(self, batch):
             return batch.clone()
         else:
             return batch  # Return as is for non-tensor types
-
-    def _extract_xyz(self, x):
-        """
-        Extract xyz (3D position) and rotation from 6DoF or 6DoF+gripper actions.
-
-        Supports:
-        - 6: 6DoF (single arm)
-        - 7: 6DoF + gripper (single arm)
-        - 12: 2 arms × 6DoF
-        - 14: 2 arms × (6DoF + gripper)
-
-        Returns:
-            xyz: Tensor with only xyz per arm (shape: ..., 3) or (..., 6) for dual-arm.
-            rot: Tensor with only rotation per arm (shape: ..., 3) or (..., 6) for dual-arm.
-        """
-        if x.shape[-1] == 6:
-            return x[..., :3], x[..., 3:6]
-        elif x.shape[-1] == 7:
-            return x[..., :3], x[..., 3:6]
-        elif x.shape[-1] == 12:
-            xyz_right = x[..., :3]
-            rot_right = x[..., 3:6]
-            xyz_left = x[..., 6:9]
-            rot_left = x[..., 9:12]
-            return torch.cat([xyz_right, xyz_left], dim=-1), torch.cat(
-                [rot_right, rot_left], dim=-1
-            )
-        elif x.shape[-1] == 14:
-            xyz_right = x[..., :3]
-            rot_right = x[..., 3:6]
-            xyz_left = x[..., 7:10]
-            rot_left = x[..., 10:13]
-            return torch.cat([xyz_right, xyz_left], dim=-1), torch.cat(
-                [rot_right, rot_left], dim=-1
-            )
-        else:
-            raise ValueError(f"Unexpected shape for 6DoF input: {x.shape}")

egomimic/eval/eval_latent.py

@@ -221,7 +221,7 @@ def on_validation_start(self):
                 exist_ok=True,
             )
 
-    def compute_metrics_and_viz(self, batch):
+    def compute_metrics_and_viz(self, batch, do_viz=True):
         algo = self.model
         metrics = {}
         images_dict = {}
@@ -260,7 +260,7 @@ def compute_metrics_and_viz(self, batch):
                 ref = _batch[ac_key]
                 B, T, D = ref.shape
                 converter = algo.action_registry.get(embodiment_id, ac_key)
-                pred_actions_orig = converter.from32(pred_actions)
+                pred_actions_orig = converter.from32_norm_6d(pred_actions)
                 pred = pred_actions_orig[:, :T, :D]
 
                 predictions = OrderedDict()
@@ -280,7 +280,7 @@ def compute_metrics_and_viz(self, batch):
                         unnorm_batch[ac_key][:, -1].cpu(),
                     )
 
-                if self.viz_func is not None:
+                if do_viz and self.viz_func is not None:
                     images_dict[embodiment_id] = self._visualize_preds(
                         unnorm_preds, unnorm_batch
                     )

egomimic/hydra_configs/data/aria.yaml

@@ -11,7 +11,10 @@ train_datasets:
         keymap_mode: cartesian
       transform_list:
         _target_: egomimic.rldb.embodiment.human.Human.get_transform_list
-        mode: cartesian
+        mode: cartesian_6d
+        # 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) so the pi0.5
+        # State: prompt bins align with the robot 20-D layout on every config.
+        pad_proprio_gripper: true
         stride: 3
     filters:
       _target_: egomimic.rldb.filters.DatasetFilter

egomimic/hydra_configs/data/aria_qwen.yaml

@@ -16,7 +16,10 @@ train_datasets:
         annotation_key: annotations
       transform_list:
         _target_: egomimic.rldb.embodiment.human.Human.get_transform_list
-        mode: cartesian
+        mode: cartesian_6d
+        # 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) so the pi0.5
+        # State: prompt bins align with the robot 20-D layout on every config.
+        pad_proprio_gripper: true
         stride: 3
     filters: null
     mode: train

egomimic/hydra_configs/data/cotrain_pi_base.yaml

@@ -12,7 +12,7 @@ train_datasets:
         annotation_key: annotations
       transform_list:
         _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list
-        mode: cartesian
+        mode: cartesian_wristframe_6d
     filters: null
     mode: train
     valid_ratio: 0.05
@@ -27,7 +27,10 @@ train_datasets:
         annotation_key: annotations
       transform_list:
         _target_: egomimic.rldb.embodiment.human.Human.get_transform_list
-        mode: cartesian
+        mode: cartesian_6d
+        # 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) so the pi0.5
+        # State: prompt bins align with the robot 20-D layout on every config.
+        pad_proprio_gripper: true
         stride: 3
     filters: null
     mode: train

egomimic/hydra_configs/data/cotrain_pi_lang_wrist.yaml

@@ -2,16 +2,11 @@ defaults:
   - cotrain_pi_lang
   - _self_
 
+# Kept for configs saved by earlier runs: cotrain_pi_lang with eva on the
+# wrist-frame 6D action, which is now the cotrain_pi_base default, so this
+# composes the same as cotrain_pi_lang. Pair with the default evaluator=eval_pi.
 train_datasets:
   eva_bimanual:
     resolver:
-      key_map:
-        keymap_mode: cartesian_wristframe_ypr
       transform_list:
-        mode: cartesian_wristframe_ypr
-  human_bimanual:
-    resolver:
-      key_map:
-        keymap_mode: cartesian
-      transform_list:
-        mode: cartesian_wristframe_ypr
+        mode: cartesian_wristframe_6d