From 38944bf152f5920b92c21778341e46504952c660 Mon Sep 17 00:00:00 2001 From: Hans Yang Date: Tue, 30 Jun 2026 22:17:54 -0700 Subject: [PATCH 1/2] feat(action): add human hand pose dataset Signed-off-by: Hans Yang --- .../data/vfm/action/datasets/__init__.py | 2 + .../human_hand_pose_lerobot_dataset.py | 219 ++++++++++++++++++ .../human_hand_pose_lerobot_dataset_test.py | 45 ++++ .../stats/human_hand_pose_lerobot_stats.json | 120 ++++++++++ .../data/vfm/action/domain_utils.py | 9 +- 5 files changed, 390 insertions(+), 5 deletions(-) create mode 100644 cosmos_framework/data/vfm/action/datasets/human_hand_pose_lerobot_dataset.py create mode 100644 cosmos_framework/data/vfm/action/datasets/human_hand_pose_lerobot_dataset_test.py create mode 100644 cosmos_framework/data/vfm/action/datasets/stats/human_hand_pose_lerobot_stats.json diff --git a/cosmos_framework/data/vfm/action/datasets/__init__.py b/cosmos_framework/data/vfm/action/datasets/__init__.py index 64b1b278..0f800377 100644 --- a/cosmos_framework/data/vfm/action/datasets/__init__.py +++ b/cosmos_framework/data/vfm/action/datasets/__init__.py @@ -12,6 +12,7 @@ from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset from cosmos_framework.data.vfm.action.datasets.bridge_orig_lerobot_dataset import BridgeOrigLeRobotDataset from cosmos_framework.data.vfm.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset +from cosmos_framework.data.vfm.action.datasets.human_hand_pose_lerobot_dataset import HumanHandPoseLeRobotDataset from cosmos_framework.data.vfm.action.datasets.robomind_franka_dataset import RoboMINDFrankaDataset from cosmos_framework.data.vfm.action.datasets.umi_lerobot_dataset import UMILeRobotDataset @@ -20,6 +21,7 @@ "AgiBotWorldBetaLeRobotDataset", "BridgeOrigLeRobotDataset", "DROIDLeRobotDataset", + "HumanHandPoseLeRobotDataset", "RoboMINDFrankaDataset", "UMILeRobotDataset", ] diff --git a/cosmos_framework/data/vfm/action/datasets/human_hand_pose_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/human_hand_pose_lerobot_dataset.py new file mode 100644 index 00000000..f6188451 --- /dev/null +++ b/cosmos_framework/data/vfm/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.vfm.action.action_spec import ActionSpec, Pos, Rot, build_action_spec +from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.vfm.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 / "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"] diff --git a/cosmos_framework/data/vfm/action/datasets/human_hand_pose_lerobot_dataset_test.py b/cosmos_framework/data/vfm/action/datasets/human_hand_pose_lerobot_dataset_test.py new file mode 100644 index 00000000..e91af315 --- /dev/null +++ b/cosmos_framework/data/vfm/action/datasets/human_hand_pose_lerobot_dataset_test.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +import numpy as np + +from cosmos_framework.data.vfm.action.datasets.human_hand_pose_lerobot_dataset import ( + HumanHandPoseLeRobotDataset, +) + + +def _row(frame_index: int) -> dict: + positions = np.zeros((21, 3), dtype=np.float32) + positions[[4, 8, 12, 16, 20], 2] = np.arange(1, 6, dtype=np.float32) * 0.01 + rotations = np.zeros((21, 4), dtype=np.float32) + rotations[:, 3] = 1.0 + return { + "observation.state.camera_position": np.zeros(3, dtype=np.float32), + "observation.state.camera_rotation": np.array([0, 0, 0, 1], dtype=np.float32), + "observation.state.hand_right_cam": positions.reshape(-1), + "observation.state.hand_right_cam_rotation": rotations.reshape(-1), + "observation.state.hand_left_cam": positions.reshape(-1), + "observation.state.hand_left_cam_rotation": rotations.reshape(-1), + "frame_index": frame_index, + } + + +def test_static_hand_pose_builds_57d_action() -> None: + dataset = object.__new__(HumanHandPoseLeRobotDataset) + dataset._chunk_length = 2 + dataset._pose_convention = "backward_framewise" + + action = dataset._build_raw_action([_row(0), _row(1), _row(2)]).numpy() + + assert action.shape == (2, 57) + assert np.isfinite(action).all() + np.testing.assert_allclose(action[:, :3], 0.0) + np.testing.assert_allclose(action[:, 9:12], 0.0) + np.testing.assert_allclose(action[:, 33:36], 0.0) + + +def test_action_spec_matches_released_width() -> None: + dataset = object.__new__(HumanHandPoseLeRobotDataset) + + assert dataset.action_dim == 57 + assert dataset._action_spec().dim == dataset.action_dim diff --git a/cosmos_framework/data/vfm/action/datasets/stats/human_hand_pose_lerobot_stats.json b/cosmos_framework/data/vfm/action/datasets/stats/human_hand_pose_lerobot_stats.json new file mode 100644 index 00000000..40c57890 --- /dev/null +++ b/cosmos_framework/data/vfm/action/datasets/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 + ] +} diff --git a/cosmos_framework/data/vfm/action/domain_utils.py b/cosmos_framework/data/vfm/action/domain_utils.py index 6f433f73..38ee2a45 100644 --- a/cosmos_framework/data/vfm/action/domain_utils.py +++ b/cosmos_framework/data/vfm/action/domain_utils.py @@ -27,6 +27,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, @@ -39,11 +40,9 @@ "embodiment_c_gripper": 29, "embodiment_c_gripper_ext": 29, "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. } From 04ee4aeed61bc3da083d365ddd89ec0a62ee1d71 Mon Sep 17 00:00:00 2001 From: Hans Yang Date: Wed, 15 Jul 2026 17:19:09 -0700 Subject: [PATCH 2/2] test(action): remove hand pose dataset test --- .../generator/action/datasets/__init__.py | 4 +- .../human_hand_pose_lerobot_dataset_test.py | 45 ------------------- 2 files changed, 3 insertions(+), 46 deletions(-) delete mode 100644 cosmos_framework/data/generator/action/datasets/human_hand_pose_lerobot_dataset_test.py diff --git a/cosmos_framework/data/generator/action/datasets/__init__.py b/cosmos_framework/data/generator/action/datasets/__init__.py index 8d01f62b..7157bdcf 100644 --- a/cosmos_framework/data/generator/action/datasets/__init__.py +++ b/cosmos_framework/data/generator/action/datasets/__init__.py @@ -8,7 +8,9 @@ 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 diff --git a/cosmos_framework/data/generator/action/datasets/human_hand_pose_lerobot_dataset_test.py b/cosmos_framework/data/generator/action/datasets/human_hand_pose_lerobot_dataset_test.py deleted file mode 100644 index 6fd8f278..00000000 --- a/cosmos_framework/data/generator/action/datasets/human_hand_pose_lerobot_dataset_test.py +++ /dev/null @@ -1,45 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: OpenMDW-1.1 - -import numpy as np - -from cosmos_framework.data.generator.action.datasets.human_hand_pose_lerobot_dataset import ( - HumanHandPoseLeRobotDataset, -) - - -def _row(frame_index: int) -> dict: - positions = np.zeros((21, 3), dtype=np.float32) - positions[[4, 8, 12, 16, 20], 2] = np.arange(1, 6, dtype=np.float32) * 0.01 - rotations = np.zeros((21, 4), dtype=np.float32) - rotations[:, 3] = 1.0 - return { - "observation.state.camera_position": np.zeros(3, dtype=np.float32), - "observation.state.camera_rotation": np.array([0, 0, 0, 1], dtype=np.float32), - "observation.state.hand_right_cam": positions.reshape(-1), - "observation.state.hand_right_cam_rotation": rotations.reshape(-1), - "observation.state.hand_left_cam": positions.reshape(-1), - "observation.state.hand_left_cam_rotation": rotations.reshape(-1), - "frame_index": frame_index, - } - - -def test_static_hand_pose_builds_57d_action() -> None: - dataset = object.__new__(HumanHandPoseLeRobotDataset) - dataset._chunk_length = 2 - dataset._pose_convention = "backward_framewise" - - action = dataset._build_raw_action([_row(0), _row(1), _row(2)]).numpy() - - assert action.shape == (2, 57) - assert np.isfinite(action).all() - np.testing.assert_allclose(action[:, :3], 0.0) - np.testing.assert_allclose(action[:, 9:12], 0.0) - np.testing.assert_allclose(action[:, 33:36], 0.0) - - -def test_action_spec_matches_released_width() -> None: - dataset = object.__new__(HumanHandPoseLeRobotDataset) - - assert dataset.action_dim == 57 - assert dataset._action_spec().dim == dataset.action_dim