feat(action): add human hand pose dataset

cosmos_framework/data/generator/action/datasets/__init__.py

@@ -8,11 +8,14 @@
 statistics without instantiating the dataset.
 """
 
-from cosmos_framework.data.generator.action.datasets.agibotworld_beta_lerobot_dataset import AgiBotWorldBetaLeRobotDataset
+from cosmos_framework.data.generator.action.datasets.agibotworld_beta_lerobot_dataset import (
+    AgiBotWorldBetaLeRobotDataset,
+)
 from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset
 from cosmos_framework.data.generator.action.datasets.bridge_orig_lerobot_dataset import BridgeOrigLeRobotDataset
 from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset
 from cosmos_framework.data.generator.action.datasets.fractal_lerobot_dataset import FractalLeRobotDataset
+from cosmos_framework.data.generator.action.datasets.human_hand_pose_lerobot_dataset import HumanHandPoseLeRobotDataset
 from cosmos_framework.data.generator.action.datasets.libero_lerobot_dataset import LIBEROLeRobotDataset
 from cosmos_framework.data.generator.action.datasets.robomind_franka_dataset import RoboMINDFrankaDataset
 from cosmos_framework.data.generator.action.datasets.robomind_ur_dataset import RoboMINDURDataset
@@ -24,6 +27,7 @@
     "BridgeOrigLeRobotDataset",
     "DROIDLeRobotDataset",
     "FractalLeRobotDataset",
+    "HumanHandPoseLeRobotDataset",
     "LIBEROLeRobotDataset",
     "RoboMINDFrankaDataset",
     "RoboMINDURDataset",

cosmos_framework/data/generator/action/datasets/human_hand_pose_lerobot_dataset.py

@@ -0,0 +1,219 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Bimanual human hand-pose dataset in LeRobot v3 format."""
+
+from __future__ import annotations
+
+import random
+from pathlib import Path
+from typing import Any, Literal
+
+import numpy as np
+import pyarrow.parquet as pq
+import torch
+from lerobot.datasets.video_utils import decode_video_frames
+
+from cosmos_framework.data.generator.action.action_spec import ActionSpec, Pos, Rot, build_action_spec
+from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset
+from cosmos_framework.data.generator.action.pose_utils import build_abs_pose_from_components, pose_abs_to_rel
+
+PoseConvention = Literal["backward_framewise"]
+Viewpoint = Literal["ego_view"]
+
+_HAND_RIGHT_POSITION_KEY = "observation.state.hand_right_cam"
+_HAND_RIGHT_ROTATION_KEY = "observation.state.hand_right_cam_rotation"
+_HAND_LEFT_POSITION_KEY = "observation.state.hand_left_cam"
+_HAND_LEFT_ROTATION_KEY = "observation.state.hand_left_cam_rotation"
+_CAM_POSITION_KEY = "observation.state.camera_position"
+_CAM_ROTATION_KEY = "observation.state.camera_rotation"
+_IMAGE_FEATURE = "observation.images.main"
+
+_NUM_JOINTS = 21
+_WRIST_JOINT_IDX = 0
+_FINGERTIP_JOINT_IDXS = (4, 8, 12, 16, 20)
+_RAW_ACTION_DIM = 57
+_NORMALIZER_PATH = Path(__file__).parent.parent / "normalizer_stats/human_hand_pose_lerobot_stats.json"
+
+# Rotate the source wrist frames into the unified convention:
+# X = thumb-to-pinky, Y = outward palm normal, Z = wrist-to-fingertips.
+_WRIST_FRAME_ALIGNMENT = np.array(
+    [[0, 1, 0, 0], [-1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]],
+    dtype=np.float32,
+)
+
+
+class HumanHandPoseLeRobotDataset(ActionBaseDataset):
+    """Bimanual human hand-pose forward-dynamics data.
+
+    The 57D action layout is
+    ``[camera(9), right_wrist(9), right_fingertips(15),
+    left_wrist(9), left_fingertips(15)]``. Camera and wrist poses use
+    framewise 3D translation plus rot6d deltas. Each hand's five fingertip
+    positions are expressed in that frame's aligned wrist coordinate system.
+
+    Source video and pose annotations are sampled at 30 FPS by default and
+    decoded at 15 FPS for Cosmos3-Nano, yielding 17 video frames and 16 action
+    transitions for the default chunk.
+    """
+
+    def __init__(
+        self,
+        root: str,
+        fps: float = 15.0,
+        chunk_length: int = 16,
+        mode: str = "forward_dynamics",
+        pose_convention: PoseConvention = "backward_framewise",
+        tolerance_s: float = 2e-4,
+        viewpoint: Viewpoint = "ego_view",
+        action_normalization: str | None = "quantile",
+        sample_stride: int = 1,
+        image_key: str = _IMAGE_FEATURE,
+    ) -> None:
+        if viewpoint != "ego_view":
+            raise NotImplementedError("Human hand-pose data only supports ego_view.")
+        super().__init__(
+            root=root,
+            domain_name="hand_pose",
+            fps=fps,
+            chunk_length=chunk_length,
+            mode=mode,
+            pose_convention=pose_convention,
+            tolerance_s=tolerance_s,
+            viewpoint=viewpoint,
+            action_normalization=action_normalization,
+            sample_stride=sample_stride,
+        )
+        source_fps = float(self._info["fps"])
+        source_stride = source_fps / self._fps
+        if not source_stride.is_integer():
+            raise ValueError(f"Source FPS {source_fps} must be an integer multiple of target FPS {self._fps}.")
+        self._source_stride = int(source_stride)
+        self._image_key = image_key
+        required_source_steps = self._source_stride * self._chunk_length
+        self._valid_starts: list[int] = []
+        episode_start = 0
+        while episode_start < len(self._rows):
+            episode_index = int(self._rows[episode_start]["episode_index"])
+            episode_end = episode_start + 1
+            while episode_end < len(self._rows) and int(self._rows[episode_end]["episode_index"]) == episode_index:
+                episode_end += 1
+            self._valid_starts.extend(
+                range(episode_start, max(episode_start, episode_end - required_source_steps), self._sample_stride)
+            )
+            episode_start = episode_end
+        subtasks_path = self._root / "meta" / "subtasks.parquet"
+        self._subtasks = (
+            {int(row["subtask_index"]): str(row["subtask"]) for row in pq.read_table(subtasks_path).to_pylist()}
+            if subtasks_path.exists()
+            else {}
+        )
+
+    @property
+    def action_dim(self) -> int:
+        return _RAW_ACTION_DIM
+
+    def _action_spec(self) -> ActionSpec:
+        return build_action_spec(
+            Pos(prefix="camera"),
+            Rot("rot6d", prefix="camera"),
+            Pos(prefix="right_wrist"),
+            Rot("rot6d", prefix="right_wrist"),
+            Pos(dim=15, prefix="right_fingertip"),
+            Pos(prefix="left_wrist"),
+            Rot("rot6d", prefix="left_wrist"),
+            Pos(dim=15, prefix="left_fingertip"),
+        )
+
+    @classmethod
+    def _stats_path(cls) -> Path:
+        return _NORMALIZER_PATH
+
+    def __len__(self) -> int:
+        return len(self._valid_starts)
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        mode = self._choose_mode()
+        start = self._valid_starts[int(idx)]
+        stop = start + self._source_stride * self._chunk_length + 1
+        rows = self._rows[start : stop : self._source_stride]
+        if len(rows) != self._chunk_length + 1:
+            raise IndexError(f"Incomplete hand-pose window at index {idx}.")
+        episode_index = int(rows[0]["episode_index"])
+        if any(int(row["episode_index"]) != episode_index for row in rows):
+            raise IndexError(f"Hand-pose window at index {idx} crosses an episode boundary.")
+
+        episode = self._episodes[episode_index]
+        video = self._load_video(episode, rows)
+        raw_action = self._build_raw_action(rows)
+        subtask_index = int(rows[0].get("subtask_index", -1))
+        task = self._tasks[int(rows[0]["task_index"])]
+        caption = self._subtasks.get(subtask_index, task)
+        ai_caption = random.choice([part.strip() for part in caption.split(" | ") if part.strip()] or [caption])
+
+        result = self._build_result(mode=mode, video=video, action=raw_action, ai_caption=ai_caption)
+        if self.action_normalization is not None:
+            result["action"] = result["action"].clamp(-1.0, 1.0)
+        return result
+
+    def _load_video(self, episode: dict[str, Any], rows: list[dict[str, Any]]) -> torch.Tensor:
+        timestamps = [float(row["timestamp"]) for row in rows]
+        from_timestamp = float(episode.get(f"videos/{self._image_key}/from_timestamp", 0.0))
+        return decode_video_frames(
+            self._video_path(episode, self._image_key),
+            [from_timestamp + timestamp for timestamp in timestamps],
+            self._tolerance_s,
+        )
+
+    @staticmethod
+    def _finger_positions_in_wrist_frame(position_data: np.ndarray, wrist_poses: np.ndarray) -> np.ndarray:
+        future_positions = position_data[1:].reshape(-1, _NUM_JOINTS, 3)
+        fingertips = future_positions[:, _FINGERTIP_JOINT_IDXS, :]
+        fingertips_h = np.concatenate(
+            [fingertips, np.ones((*fingertips.shape[:-1], 1), dtype=np.float32)],
+            axis=-1,
+        )
+        wrist_inv = np.linalg.inv(wrist_poses[1:])
+        fingertips_wrist = np.einsum("tij,tnj->tni", wrist_inv, fingertips_h)[..., :3]
+        return fingertips_wrist.reshape(len(future_positions), -1)
+
+    def _build_raw_action(self, rows: list[dict[str, Any]]) -> torch.Tensor:
+        def values(key: str) -> np.ndarray:
+            return np.asarray([row[key] for row in rows], dtype=np.float32)
+
+        camera_pose = build_abs_pose_from_components(values(_CAM_POSITION_KEY), values(_CAM_ROTATION_KEY), "quat_xyzw")
+
+        right_positions = values(_HAND_RIGHT_POSITION_KEY)
+        right_rotations = values(_HAND_RIGHT_ROTATION_KEY).reshape(-1, _NUM_JOINTS, 4)
+        left_positions = values(_HAND_LEFT_POSITION_KEY)
+        left_rotations = values(_HAND_LEFT_ROTATION_KEY).reshape(-1, _NUM_JOINTS, 4)
+
+        right_wrist_camera = (
+            build_abs_pose_from_components(right_positions[:, :3], right_rotations[:, _WRIST_JOINT_IDX], "quat_xyzw")
+            @ _WRIST_FRAME_ALIGNMENT
+        )
+        left_wrist_camera = (
+            build_abs_pose_from_components(left_positions[:, :3], left_rotations[:, _WRIST_JOINT_IDX], "quat_xyzw")
+            @ _WRIST_FRAME_ALIGNMENT
+        )
+
+        right_wrist_world = camera_pose @ right_wrist_camera
+        left_wrist_world = camera_pose @ left_wrist_camera
+        action = np.concatenate(
+            [
+                pose_abs_to_rel(camera_pose, rotation_format="rot6d", pose_convention=self._pose_convention),
+                pose_abs_to_rel(right_wrist_world, rotation_format="rot6d", pose_convention=self._pose_convention),
+                self._finger_positions_in_wrist_frame(right_positions, right_wrist_camera),
+                pose_abs_to_rel(left_wrist_world, rotation_format="rot6d", pose_convention=self._pose_convention),
+                self._finger_positions_in_wrist_frame(left_positions, left_wrist_camera),
+            ],
+            axis=-1,
+        )
+        if action.shape != (self._chunk_length, _RAW_ACTION_DIM):
+            raise ValueError(
+                f"Expected hand-pose action shape {(self._chunk_length, _RAW_ACTION_DIM)}, got {action.shape}."
+            )
+        return torch.from_numpy(action).float()
+
+
+__all__ = ["HumanHandPoseLeRobotDataset"]

cosmos_framework/data/generator/action/domain_utils.py

@@ -29,6 +29,7 @@
 EMBODIMENT_TO_RAW_ACTION_DIM: dict[str, int] = {
     "av": 9,
     "camera_pose": 9,
+    "hand_pose": 57,
     "pusht": 2,
     "umi": 10,
     "bridge_orig_lerobot": 10,
@@ -43,11 +44,9 @@
     "xdof_yam": 20,
     "molmoact2_yam": 20,
     "fractal": 10,
-    # NOTE: ``libero`` (7/10/13 depending on ``rotation_space``) and ``hand_pose``
-    # (variable with ``keypoint_option`` and ``rotation_format``) are absent
-    # because their raw width is set per-dataset at construction time. Inference
-    # in inverse_dynamics/policy modes is not supported for these domains until
-    # canonical widths are added here.
+    # NOTE: ``libero`` (7/10/13 depending on ``rotation_space``) is absent because
+    # its raw width is selected per dataset at construction time. ``hand_pose``
+    # uses the released 57D wrist-plus-fingertips rot6d representation.
 }
 
 

cosmos_framework/data/generator/action/normalizer_stats/human_hand_pose_lerobot_stats.json

@@ -0,0 +1,120 @@
+{
+  "q01": [
+    -0.026688,
+    -0.024154,
+    -0.013269,
+    0.996827,
+    -0.056513,
+    -0.02372,
+    -0.055905,
+    0.996923,
+    -0.026957,
+    -0.048362,
+    -0.040973,
+    -0.040898,
+    0.869093,
+    -0.316019,
+    -0.240748,
+    -0.32718,
+    0.879926,
+    -0.21302,
+    -0.080012,
+    0.029713,
+    0.059227,
+    -0.038851,
+    0.010826,
+    0.062254,
+    -0.007452,
+    0.010061,
+    0.047149,
+    0.005736,
+    0.006631,
+    0.036123,
+    0.018989,
+    0.001386,
+    0.033526,
+    -0.037843,
+    -0.037665,
+    -0.036427,
+    0.90169,
+    -0.278892,
+    -0.200027,
+    -0.267251,
+    0.906997,
+    -0.187048,
+    -0.013628,
+    0.024588,
+    0.053585,
+    -0.012045,
+    0.000813,
+    0.065182,
+    -0.037355,
+    -0.001436,
+    0.052279,
+    -0.058561,
+    -0.004175,
+    0.042882,
+    -0.078014,
+    -0.007496,
+    0.039113
+  ],
+  "q99": [
+    0.026573,
+    0.015575,
+    0.015514,
+    1.0,
+    0.055889,
+    0.023992,
+    0.056601,
+    1.0,
+    0.027849,
+    0.044917,
+    0.040128,
+    0.043628,
+    0.999997,
+    0.322927,
+    0.238815,
+    0.311405,
+    0.999997,
+    0.222693,
+    0.016559,
+    0.105274,
+    0.129951,
+    0.014097,
+    0.089857,
+    0.188346,
+    0.037904,
+    0.087388,
+    0.19212,
+    0.059029,
+    0.08399,
+    0.17789,
+    0.074478,
+    0.069913,
+    0.145329,
+    0.041177,
+    0.0368,
+    0.039655,
+    0.999998,
+    0.27139,
+    0.203172,
+    0.282521,
+    0.999998,
+    0.193682,
+    0.087733,
+    0.10751,
+    0.128318,
+    0.040024,
+    0.08864,
+    0.191875,
+    0.008532,
+    0.08699,
+    0.195138,
+    -0.005823,
+    0.083213,
+    0.180562,
+    -0.019596,
+    0.069257,
+    0.148124
+  ]
+}