diff --git a/cosmos_framework/data/generator/action/datasets/__init__.py b/cosmos_framework/data/generator/action/datasets/__init__.py index 62437797..7157bdcf 100644 --- a/cosmos_framework/data/generator/action/datasets/__init__.py +++ b/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", diff --git a/cosmos_framework/data/generator/action/datasets/human_hand_pose_lerobot_dataset.py b/cosmos_framework/data/generator/action/datasets/human_hand_pose_lerobot_dataset.py new file mode 100644 index 00000000..20e157ab --- /dev/null +++ b/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"] diff --git a/cosmos_framework/data/generator/action/domain_utils.py b/cosmos_framework/data/generator/action/domain_utils.py index 66cc51cb..6317a096 100644 --- a/cosmos_framework/data/generator/action/domain_utils.py +++ b/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. } diff --git a/cosmos_framework/data/generator/action/normalizer_stats/human_hand_pose_lerobot_stats.json b/cosmos_framework/data/generator/action/normalizer_stats/human_hand_pose_lerobot_stats.json new file mode 100644 index 00000000..40c57890 --- /dev/null +++ b/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 + ] +}