diff --git a/GR00T-WholeBodyControl/gear_sonic/__pycache__/__init__.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..72af732fc5b570aba105c9b7bbf7bd61c142f565 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/__pycache__/__init__.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/__pycache__/__init__.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3ac7032e30a4494ea80df510968138a062a5ac9b Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/__pycache__/__init__.cpython-311.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/__pycache__/__init__.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4cacd6c553f89c81b01b49a074c7ac2849ff3bbc Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/__pycache__/__init__.cpython-312.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/__pycache__/eval_agent_trl.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/__pycache__/eval_agent_trl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2ddca28e1d1c639aee80cdc8ad160ed8a4f6347a Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/__pycache__/eval_agent_trl.cpython-312.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/__pycache__/train_agent_trl.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/__pycache__/train_agent_trl.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0729bde080594e3e3e146dcd11f034ae80d55de0 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/__pycache__/train_agent_trl.cpython-311.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/__pycache__/train_agent_trl.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/__pycache__/train_agent_trl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..793396742b3f5eda4a4c37a8a5634be90b539ed5 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/__pycache__/train_agent_trl.cpython-312.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/__pycache__/version.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/__pycache__/version.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2affc490cf8d499f8132a22c404f1b6ebcb44435 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/__pycache__/version.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/__pycache__/version.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/__pycache__/version.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fcaa38c4c22f81fc23944d5ca167a77d22bc05b8 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/__pycache__/version.cpython-312.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/camera/sensor.py b/GR00T-WholeBodyControl/gear_sonic/camera/sensor.py new file mode 100644 index 0000000000000000000000000000000000000000..5b6f22120cc796d974432d6d79dbc94d6c8f3bc6 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/camera/sensor.py @@ -0,0 +1,33 @@ +"""Base sensor abstract class. + +The ``gymnasium`` dependency is lazy-imported so the camera server can +run without it. +""" + +from abc import abstractmethod +from typing import Any + + +class Sensor: + """Base class for camera / sensor implementations. + + Concrete drivers (OAK, RealSense, ZED, USB, …) inherit from this and + implement at least :meth:`read` and :meth:`serialize`. + """ + + def read(self, **kwargs) -> Any: + """Read the current sensor value (e.g. a dict of images).""" + + def observation_space(self): + """Return a ``gymnasium.Space`` describing the observation. + + Only used during init to report camera capabilities to the + composed-camera orchestrator; not required for data collection. + """ + + @abstractmethod + def serialize(self, data: dict[str, Any]) -> dict[str, Any]: + """Serialize the sensor reading for ZMQ transmission.""" + + def close(self): + """Release hardware resources.""" diff --git a/GR00T-WholeBodyControl/gear_sonic/camera/sensor_server.py b/GR00T-WholeBodyControl/gear_sonic/camera/sensor_server.py new file mode 100644 index 0000000000000000000000000000000000000000..7af89f68b841757d6759ae18c0b98d3859be57d1 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/camera/sensor_server.py @@ -0,0 +1,243 @@ +"""ZMQ PUB/SUB transport and image serialisation for the camera server.""" + +import base64 +from dataclasses import dataclass, field +from enum import Enum +from typing import Any + +import cv2 +import msgpack +import msgpack_numpy as m +import numpy as np +import zmq + + +# ============================================================================= +# Pose Message Schema +# ============================================================================= +@dataclass +class PoseData: + """Single pose data point with quaternion orientation and translation.""" + + qx: float = 0.0 + qy: float = 0.0 + qz: float = 0.0 + qw: float = 1.0 + tx: float = 0.0 + ty: float = 0.0 + tz: float = 0.0 + + def to_dict(self) -> dict[str, float]: + return { + "qx": self.qx, + "qy": self.qy, + "qz": self.qz, + "qw": self.qw, + "tx": self.tx, + "ty": self.ty, + "tz": self.tz, + } + + @staticmethod + def from_dict(data: dict[str, float]) -> "PoseData": + return PoseData( + qx=data.get("qx", 0.0), + qy=data.get("qy", 0.0), + qz=data.get("qz", 0.0), + qw=data.get("qw", 1.0), + tx=data.get("tx", 0.0), + ty=data.get("ty", 0.0), + tz=data.get("tz", 0.0), + ) + + def to_array(self) -> np.ndarray: + return np.array([self.qx, self.qy, self.qz, self.qw, self.tx, self.ty, self.tz]) + + @staticmethod + def from_array(arr: np.ndarray) -> "PoseData": + return PoseData( + qx=float(arr[0]), + qy=float(arr[1]), + qz=float(arr[2]), + qw=float(arr[3]), + tx=float(arr[4]), + ty=float(arr[5]), + tz=float(arr[6]), + ) + + +@dataclass +class PoseMessageSchema: + """Standardized message schema for pose / positional data.""" + + timestamp: float = 0.0 + device_id: str = "iphone" + pose: PoseData = field(default_factory=PoseData) + + def serialize(self) -> bytes: + data = { + "timestamp": self.timestamp, + "device_id": self.device_id, + "pose": self.pose.to_dict(), + } + return msgpack.packb(data, use_bin_type=True) + + @staticmethod + def deserialize(packed_data: bytes) -> "PoseMessageSchema": + data = msgpack.unpackb(packed_data, object_hook=m.decode) + return PoseMessageSchema( + timestamp=data.get("timestamp", 0.0), + device_id=data.get("device_id", "iphone"), + pose=PoseData.from_dict(data.get("pose", {})), + ) + + def asdict(self) -> dict[str, Any]: + return { + "timestamp": self.timestamp, + "device_id": self.device_id, + "pose": self.pose.to_dict(), + } + + +# ============================================================================= +# Image Message Schema +# ============================================================================= +@dataclass +class ImageMessageSchema: + """Standardized message schema for camera images. + + Handles two encodings on the wire: + + * **str** – legacy base64-encoded JPEG. + * **bytes** – raw JPEG from on-device MJPEG encoder (e.g. OAK). + """ + + timestamps: dict[str, float] + images: dict[str, np.ndarray] + + def serialize(self) -> dict[str, Any]: + serialized_msg: dict[str, Any] = {"timestamps": self.timestamps, "images": {}} + for key, image in self.images.items(): + if isinstance(image, bytes | bytearray): + serialized_msg["images"][key] = image + else: + serialized_msg["images"][key] = ImageUtils.encode_image(image) + return serialized_msg + + @staticmethod + def deserialize(data: dict[str, Any]) -> "ImageMessageSchema": + timestamps = data.get("timestamps", {}) + images = {} + for key, value in data.get("images", {}).items(): + if isinstance(value, bytes | bytearray): + mat = cv2.imdecode(np.frombuffer(value, dtype=np.uint8), cv2.IMREAD_COLOR) + images[key] = mat[..., ::-1] # BGR -> RGB + elif isinstance(value, str): + images[key] = ImageUtils.decode_image(value) + elif isinstance(value, np.ndarray): + images[key] = value + elif isinstance(value, dict) and b"nd" in value: + images[key] = m.decode(value) + else: + images[key] = value + return ImageMessageSchema(timestamps=timestamps, images=images) + + def asdict(self) -> dict[str, Any]: + return {"timestamps": self.timestamps, "images": self.images} + + +# ============================================================================= +# ZMQ Server / Client +# ============================================================================= +class SensorServer: + """ZMQ PUB server that streams msgpack-encoded sensor payloads.""" + + def start_server(self, port: int): + self.context = zmq.Context() + self.socket = self.context.socket(zmq.PUB) + self.socket.setsockopt(zmq.SNDHWM, 20) + self.socket.setsockopt(zmq.LINGER, 0) + self.socket.bind(f"tcp://*:{port}") + print(f"Sensor server running at tcp://*:{port}") + + self.message_sent = 0 + self.message_dropped = 0 + + def stop_server(self): + self.socket.close() + self.context.term() + + def send_message(self, data: dict[str, Any]): + try: + packed = msgpack.packb(data, use_bin_type=True) + self.socket.send(packed, flags=zmq.NOBLOCK) + except zmq.Again: + self.message_dropped += 1 + print(f"[Warning] message dropped: {self.message_dropped}") + self.message_sent += 1 + + if self.message_sent % 100 == 0: + print( + f"[Sensor server] Message sent: {self.message_sent}, " + f"message dropped: {self.message_dropped}" + ) + + +class SensorClient: + """ZMQ SUB client that receives msgpack-encoded sensor payloads.""" + + def start_client(self, server_ip: str, port: int): + self.context = zmq.Context() + self.socket = self.context.socket(zmq.SUB) + self.socket.setsockopt_string(zmq.SUBSCRIBE, "") + self.socket.setsockopt(zmq.CONFLATE, True) + self.socket.setsockopt(zmq.RCVHWM, 3) + self.socket.connect(f"tcp://{server_ip}:{port}") + + def stop_client(self): + self.socket.close() + self.context.term() + + def receive_message(self): + packed = self.socket.recv() + return msgpack.unpackb(packed, object_hook=m.decode) + + def receive_message_nonblocking(self, timeout_ms: int = 0): + if self.socket.poll(timeout_ms): + packed = self.socket.recv() + return msgpack.unpackb(packed, object_hook=m.decode) + return None + + +# ============================================================================= +# Helpers +# ============================================================================= +class CameraMountPosition(Enum): + EGO_VIEW = "ego_view" + HEAD = "head" + LEFT_WRIST = "left_wrist" + RIGHT_WRIST = "right_wrist" + + +class ImageUtils: + @staticmethod + def encode_image(image: np.ndarray) -> str: + _, color_buffer = cv2.imencode(".jpg", image, [int(cv2.IMWRITE_JPEG_QUALITY), 80]) + return base64.b64encode(color_buffer).decode("utf-8") + + @staticmethod + def encode_depth_image(image: np.ndarray) -> str: + depth_compressed = cv2.imencode(".png", image)[1].tobytes() + return base64.b64encode(depth_compressed).decode("utf-8") + + @staticmethod + def decode_image(image: str) -> np.ndarray: + color_data = base64.b64decode(image) + color_array = np.frombuffer(color_data, dtype=np.uint8) + return cv2.imdecode(color_array, cv2.IMREAD_COLOR) + + @staticmethod + def decode_depth_image(image: str) -> np.ndarray: + depth_data = base64.b64decode(image) + depth_array = np.frombuffer(depth_data, dtype=np.uint8) + return cv2.imdecode(depth_array, cv2.IMREAD_UNCHANGED) diff --git a/GR00T-WholeBodyControl/gear_sonic/data/exporter.py b/GR00T-WholeBodyControl/gear_sonic/data/exporter.py new file mode 100644 index 0000000000000000000000000000000000000000..2aa21ca9abf5cbcdaebcc57d222cd4c9bc5417c6 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/data/exporter.py @@ -0,0 +1,472 @@ +""" +Gr00t data exporter for LeRobot-format datasets. +""" + +import copy +from dataclasses import asdict, dataclass +from functools import partial +import json +import os +from pathlib import Path +import shutil +from typing import Any, Optional + +import datasets +from datasets import load_dataset +from datasets.utils import disable_progress_bars +from huggingface_hub.errors import RepositoryNotFoundError +from lerobot.common.datasets.lerobot_dataset import ( + LeRobotDataset, + LeRobotDatasetMetadata, + compute_episode_stats, +) +from lerobot.common.datasets.utils import ( + check_timestamps_sync, + get_episode_data_index, + validate_episode_buffer, + validate_frame, +) +import numpy as np +from PIL import Image as PILImage +import torch +from torchvision import transforms + +from gear_sonic.data.video_writer import VideoWriter + +disable_progress_bars() + + +# --------------------------------------------------------------------------- +# ArgsConfig (inlined from decoupled_wbc.control.main.config_template) +# --------------------------------------------------------------------------- + + +@dataclass +class ArgsConfig: + """Minimal config dataclass for script_config serialization.""" + + def update( + self, + config_dict: dict, + strict: bool = False, + skip_keys: list[str] = [], + allowed_keys: list[str] | None = None, + ): + for k, v in config_dict.items(): + if k in skip_keys: + continue + if allowed_keys is not None and k not in allowed_keys: + continue + if strict and not hasattr(self, k): + raise ValueError(f"Config {k} not found in {self.__class__.__name__}") + if not strict and not hasattr(self, k): + continue + setattr(self, k, v) + + @classmethod + def from_dict( + cls, + config_dict: dict, + strict: bool = False, + skip_keys: list[str] = [], + allowed_keys: list[str] | None = None, + ): + instance = cls() + instance.update( + config_dict=config_dict, strict=strict, skip_keys=skip_keys, allowed_keys=allowed_keys + ) + return instance + + def to_dict(self): + return asdict(self) + + def get(self, key: str, default: Any = None): + return getattr(self, key) if hasattr(self, key) else default + + +# --------------------------------------------------------------------------- +# Gr00tDatasetMetadata +# --------------------------------------------------------------------------- + + +class Gr00tDatasetMetadata(LeRobotDatasetMetadata): + """Additional metadata on top of LeRobotDatasetMetadata: + - modality_config: Written to ``meta/modality.json`` + - discarded_episode_indices: Written to ``meta/info.json`` + """ + + MODALITY_CONFIG_REL_PATH = Path("meta/modality.json") + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + with open(self.root / self.MODALITY_CONFIG_REL_PATH, "rb") as f: + self.modality_config = json.load(f) + + @classmethod + def create( + cls, + modality_config: dict, + script_config: dict, + *args, + **kwargs, + ): + cls.validate_modality_config(modality_config) + + obj = super().create(*args, **kwargs) + + obj.info["script_config"] = script_config + obj.info["discarded_episode_indices"] = [] + with open(obj.root / "meta" / "info.json", "w") as f: + json.dump(obj.info, f, indent=4) + + obj.__class__ = cls + with open(obj.root / cls.MODALITY_CONFIG_REL_PATH, "w") as f: + json.dump(modality_config, f, indent=4) + obj.modality_config = modality_config + return obj + + @staticmethod + def validate_modality_config(modality_config: dict) -> None: + valid_keys = ["state", "action", "video", "annotation"] + if not all(key in modality_config for key in valid_keys): + raise ValueError( + f"Modality config must contain all of the following keys: {valid_keys}" + ) + for key in valid_keys: + if key not in modality_config: + raise ValueError(f"Modality config must contain a '{key}' key") + + +# --------------------------------------------------------------------------- +# Gr00tDataExporter +# --------------------------------------------------------------------------- + + +class Gr00tDataExporter(LeRobotDataset): + """Exports data collected for a single session to LeRobot Dataset. + + Lifecycle: + 1. Create a Gr00tDataExporter object + 2. Add frames using add_frame() + 3. Save the episode using save_episode() + - Flushes the episode buffer to disk + - Closes the video writers + - Creates new video writer and ep buffer for the next episode + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.video_writers = self.create_video_writer() + + @property + def repo_id(self): + return self.meta.repo_id + + @property + def root(self): + return self.meta.root + + @property + def local_files_only(self): + return self.meta.local_files_only + + @property + def video_keys(self): + return self.meta.video_keys + + @classmethod + def create( + cls, + save_root: str | Path, + fps: int, + features: dict, + modality_config: dict, + task: str, + script_config: ArgsConfig | dict = None, + robot_type: str | None = None, + tolerance_s: float = 1e-4, + vcodec: str = "h264", + overwrite_existing: bool = False, + ) -> "Gr00tDataExporter": + if script_config is None: + script_config = {} + + obj = cls.__new__(cls) + repo_id = "tmp/tmp_dataset" + + if overwrite_existing and (Path(save_root)).exists(): + print( + f"Found existing dataset at {save_root}", + "Cleaning up this directory since overwrite_existing is True.", + ) + shutil.rmtree(save_root) + + if (Path(save_root)).exists(): + try: + obj.meta = Gr00tDatasetMetadata( + repo_id=repo_id, + root=save_root, + ) + except RepositoryNotFoundError as e: + raise ValueError( + f"Failed to resume from corrupted dataset. " + f"Please manually check the dataset at {save_root}" + ) from e + else: + if not isinstance(script_config, dict): + script_config = script_config.to_dict() + obj.meta = Gr00tDatasetMetadata.create( + repo_id=repo_id, + fps=fps, + root=save_root, + robot=None, + robot_type=robot_type, + features=features, + modality_config=modality_config, + script_config=script_config, + use_videos=True, + ) + + obj.tolerance_s = tolerance_s + obj.video_backend = "pyav" + obj.vcodec = vcodec + obj.task = task + obj.image_writer = None + + obj.episode_buffer = obj.create_episode_buffer() + + obj.episodes = None + obj.hf_dataset = obj.create_hf_dataset() + obj.image_transforms = None + obj.delta_timestamps = None + obj.delta_indices = None + obj.episode_data_index = None + obj.video_writers = obj.create_video_writer() + return obj + + def create_video_writer(self) -> dict[str, VideoWriter]: + video_writers = {} + for key in self.meta.video_keys: + video_writers[key] = VideoWriter( + self.root + / self.meta.get_video_file_path(self.episode_buffer["episode_index"], key), + self.meta.shapes[key][1], + self.meta.shapes[key][0], + self.fps, + self.vcodec, + ) + return video_writers + + def add_frame(self, frame: dict) -> None: + """Add a frame to the episode buffer. Videos are handled by the video_writer.""" + frame = copy.deepcopy(frame) + frame["task"] = frame.get("task", self.task) + + for name in frame: + if isinstance(frame[name], torch.Tensor): + frame[name] = frame[name].numpy() + + validate_frame(frame, self.features) + + if self.episode_buffer is None: + self.episode_buffer = self.create_episode_buffer() + + frame_index = self.episode_buffer["size"] + timestamp = frame.pop("timestamp") if "timestamp" in frame else frame_index / self.fps + self.episode_buffer["frame_index"].append(frame_index) + self.episode_buffer["timestamp"].append(timestamp) + + for key in frame: + if key == "task": + self.episode_buffer["task"].append(frame["task"]) + continue + + if key not in self.features: + raise ValueError( + f"An element of the frame is not in the features. " + f"'{key}' not in '{self.features.keys()}'." + ) + + if self.features[key]["dtype"] in ["image", "video"]: + img_path = self._get_image_file_path( + episode_index=self.episode_buffer["episode_index"], + image_key=key, + frame_index=frame_index, + ) + if frame_index == 0: + img_path.parent.mkdir(parents=True, exist_ok=True) + + self.video_writers[key].add_frame(frame[key]) + self.episode_buffer[key].append(str(img_path)) + else: + self.episode_buffer[key].append(frame[key]) + + self.episode_buffer["size"] += 1 + + def stop_video_writers(self): + if not hasattr(self, "video_writers"): + raise RuntimeError( + "Can't stop video writers because they haven't been initialized. Call create() first." + ) + for key in self.video_writers: + self.video_writers[key].stop() + + def skip_and_start_new_episode(self) -> None: + """Skip the current episode and start a new one.""" + self.stop_video_writers() + self.episode_buffer = self.create_episode_buffer() + self.video_writers = self.create_video_writer() + + def save_episode(self, episode_data: dict | None = None) -> None: + if not episode_data: + episode_buffer = self.episode_buffer + + validate_episode_buffer(episode_buffer, self.meta.total_episodes, self.features) + + episode_length = episode_buffer.pop("size") + tasks = episode_buffer.pop("task") + episode_tasks = list(set(tasks)) + episode_index = episode_buffer["episode_index"] + + episode_buffer["index"] = np.arange( + self.meta.total_frames, self.meta.total_frames + episode_length + ) + episode_buffer["episode_index"] = np.full((episode_length,), episode_index) + + for task in episode_tasks: + task_index = self.meta.get_task_index(task) + if task_index is None: + self.meta.add_task(task) + + episode_buffer["task_index"] = np.array([self.meta.get_task_index(task) for task in tasks]) + + for key, ft in self.features.items(): + if key in ["index", "episode_index", "task_index"] or ft["dtype"] in ["image", "video"]: + continue + episode_buffer[key] = np.stack(episode_buffer[key]) + + self._wait_image_writer() + self._save_episode_table(episode_buffer, episode_index) + + non_video_features = {k: v for k, v in self.features.items() if v["dtype"] not in ["video"]} + non_vid_ep_buffer = { + k: v for k, v in episode_buffer.items() if k in non_video_features.keys() + } + ep_stats = compute_episode_stats(non_vid_ep_buffer, non_video_features) + + if len(self.meta.video_keys) > 0: + video_paths = self.encode_episode_videos(episode_index) + for key in self.meta.video_keys: + episode_buffer[key] = video_paths[key] + + self.meta.save_episode(episode_index, episode_length, episode_tasks, ep_stats) + + ep_data_index = get_episode_data_index(self.meta.episodes, [episode_index]) + ep_data_index_np = {k: t.numpy() for k, t in ep_data_index.items()} + check_timestamps_sync( + episode_buffer["timestamp"], + episode_buffer["episode_index"], + ep_data_index_np, + self.fps, + self.tolerance_s, + ) + + video_files = list(self.root.rglob("*.mp4")) + assert len(video_files) == self.num_episodes * len(self.meta.video_keys) + + parquet_files = list(self.root.rglob("*.parquet")) + assert len(parquet_files) == self.num_episodes + + img_dir = self.root / "images" + if img_dir.is_dir(): + shutil.rmtree(self.root / "images") + + if not episode_data: + self.episode_buffer = self.create_episode_buffer() + self.video_writers = self.create_video_writer() + + for key in self.meta.video_keys: + video_path = os.path.join(self.root, self.meta.get_video_file_path(episode_index, key)) + if not os.path.exists(video_path): + raise FileNotFoundError( + f"Video path: {video_path} does not exist for episode {episode_index}" + ) + + parquet_path = os.path.join(self.root, self.meta.get_data_file_path(episode_index)) + if not os.path.exists(parquet_path): + raise FileNotFoundError( + f"Parquet path: {parquet_path} does not exist for episode {episode_index}" + ) + + def encode_episode_videos(self, episode_index: int) -> dict: + video_paths = {} + for key in self.meta.video_keys: + video_paths[key] = self.video_writers[key].stop() + return video_paths + + def save_episode_as_discarded(self) -> None: + """Flag ongoing episode as discarded and save it to disk.""" + self.meta.info["discarded_episode_indices"] = self.meta.info.get( + "discarded_episode_indices", [] + ) + [self.episode_buffer["episode_index"]] + self.save_episode() + + +# --------------------------------------------------------------------------- +# HF dataset helpers (for loading saved datasets) +# --------------------------------------------------------------------------- + + +def hf_transform_to_torch_by_features( + features: datasets.Sequence, items_dict: dict[torch.Tensor | None] +): + for key in items_dict: + first_item = items_dict[key][0] + if isinstance(first_item, PILImage.Image): + to_tensor = transforms.ToTensor() + items_dict[key] = [to_tensor(img) for img in items_dict[key]] + elif first_item is None: + pass + else: + if isinstance(features[key], datasets.Value): + dtype_str = features[key].dtype + elif isinstance(features[key], datasets.Sequence): + assert isinstance(features[key].feature, datasets.Value) + dtype_str = features[key].feature.dtype + else: + raise ValueError(f"Unsupported feature type for key '{key}': {features[key]}") + dtype_mapping = { + "float32": torch.float32, + "float64": torch.float64, + "int32": torch.int32, + "int64": torch.int64, + } + items_dict[key] = [ + torch.tensor(x, dtype=dtype_mapping[dtype_str]) for x in items_dict[key] + ] + return items_dict + + +class TypedLeRobotDataset(LeRobotDataset): + def __init__(self, load_video=True, *args, **kwargs): + super().__init__(*args, **kwargs) + if not load_video: + video_keys = [] + for key in self.meta.features.keys(): + if self.meta.features[key]["dtype"] == "video": + video_keys.append(key) + for key in video_keys: + self.meta.features.pop(key) + + def load_hf_dataset(self) -> datasets.Dataset: + if self.episodes is None: + path = str(self.root / "data") + hf_dataset = load_dataset("parquet", data_dir=path, split="train") + else: + files = [ + str(self.root / self.meta.get_data_file_path(ep_idx)) for ep_idx in self.episodes + ] + hf_dataset = load_dataset("parquet", data_files=files, split="train") + + hf_dataset.set_transform(partial(hf_transform_to_torch_by_features, hf_dataset.features)) + return hf_dataset diff --git a/GR00T-WholeBodyControl/gear_sonic/data/features_sonic_vla.py b/GR00T-WholeBodyControl/gear_sonic/data/features_sonic_vla.py new file mode 100644 index 0000000000000000000000000000000000000000..b59fd0f358677ab9253dc36510c024f9a79240bb --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/data/features_sonic_vla.py @@ -0,0 +1,410 @@ +""" +Dataset configuration for the Sonic VLA pipeline. + +Provides feature/modality config dicts and a convenience function to +instantiate the G1 RobotModel needed for FK and joint configuration +assembly during data collection. + +Joint names, counts, and group indices are derived at runtime from the +``RobotModel`` (via its ``supplemental_info``). +""" + +from __future__ import annotations + +from typing import Literal + +from gear_sonic.data.robot_model import RobotModel + +EGO_VIEW_HEIGHT: int = 480 +EGO_VIEW_WIDTH: int = 640 +WRIST_VIEW_HEIGHT: int = 480 +WRIST_VIEW_WIDTH: int = 640 +FPS: int = 50 + + +_JOINT_GROUPS_FOR_STATE: list[str] = [ + "left_leg", + "right_leg", + "waist", + "left_arm", + "left_hand", + "right_arm", + "right_hand", +] + + +def _get_joint_group_slices(robot_model: RobotModel) -> dict[str, dict[str, int]]: + """Derive ``{group_name: {"start": ..., "end": ...}}`` from the robot model.""" + slices: dict[str, dict[str, int]] = {} + for group in _JOINT_GROUPS_FOR_STATE: + indices = sorted(robot_model.get_joint_group_indices(group)) + slices[group] = {"start": indices[0], "end": indices[-1] + 1} + return slices + + +def get_modality_config_sonic_vla(robot_model: RobotModel) -> dict: + """Return the modality config for the Sonic VLA dataset. + + Produces the exact content of meta/modality.json. + """ + group_slices = _get_joint_group_slices(robot_model) + + return { + "state": { + **group_slices, + "left_wrist_pos": { + "start": 0, + "end": 3, + "original_key": "observation.eef_state", + }, + "left_wrist_abs_quat": { + "start": 3, + "end": 7, + "original_key": "observation.eef_state", + "rotation_type": "quaternion", + }, + "right_wrist_pos": { + "start": 7, + "end": 10, + "original_key": "observation.eef_state", + }, + "right_wrist_abs_quat": { + "start": 10, + "end": 14, + "original_key": "observation.eef_state", + "rotation_type": "quaternion", + }, + "root_orientation": { + "start": 0, + "end": 4, + "original_key": "observation.root_orientation", + "rotation_type": "quaternion", + }, + "projected_gravity": { + "start": 0, + "end": 3, + "original_key": "observation.projected_gravity", + }, + "cpp_rotation_offset": { + "start": 0, + "end": 4, + "original_key": "observation.cpp_rotation_offset", + "rotation_type": "quaternion", + }, + "init_base_quat": { + "start": 0, + "end": 4, + "original_key": "observation.init_base_quat", + "rotation_type": "quaternion", + }, + }, + "action": { + "delta_heading": { + "start": 0, + "end": 1, + "original_key": "teleop.delta_heading", + }, + "motion_token": { + "start": 0, + "end": 64, + "original_key": "action.motion_token", + }, + "smpl_joints": { + "start": 0, + "end": 72, + "original_key": "teleop.smpl_joints", + }, + "smpl_pose": { + "start": 0, + "end": 63, + "original_key": "teleop.smpl_pose", + }, + "body_quat_w": { + "start": 0, + "end": 4, + "original_key": "teleop.body_quat_w", + "rotation_type": "quaternion", + }, + "target_body_orientation": { + "start": 0, + "end": 6, + "original_key": "teleop.target_body_orientation", + "rotation_type": "rotation_6d", + }, + "left_hand_joints": { + "start": 0, + "end": 7, + "original_key": "teleop.left_hand_joints", + }, + "right_hand_joints": { + "start": 0, + "end": 7, + "original_key": "teleop.right_hand_joints", + }, + "left_wrist_joints": { + "start": 0, + "end": 3, + "original_key": "teleop.left_wrist_joints", + }, + "right_wrist_joints": { + "start": 0, + "end": 3, + "original_key": "teleop.right_wrist_joints", + }, + "stream_mode": { + "start": 0, + "end": 1, + "original_key": "teleop.stream_mode", + }, + "planner_mode": { + "start": 0, + "end": 1, + "original_key": "teleop.planner_mode", + }, + "planner_movement": { + "start": 0, + "end": 3, + "original_key": "teleop.planner_movement", + }, + "planner_facing": { + "start": 0, + "end": 3, + "original_key": "teleop.planner_facing", + }, + "planner_speed": { + "start": 0, + "end": 1, + "original_key": "teleop.planner_speed", + }, + "planner_height": { + "start": 0, + "end": 1, + "original_key": "teleop.planner_height", + }, + "vr_3pt_position": { + "start": 0, + "end": 9, + "original_key": "teleop.vr_3pt_position", + }, + "vr_3pt_orientation": { + "start": 0, + "end": 18, + "original_key": "teleop.vr_3pt_orientation", + "rotation_type": "rotation_6d", + }, + }, + "video": { + "ego_view": {"original_key": "observation.images.ego_view"}, + }, + "annotation": { + "human.task_description": {"original_key": "task_index"}, + }, + } + + +def get_features_sonic_vla(robot_model: RobotModel) -> dict: + """Return the dataset features for the Sonic VLA dataset. + + The returned dict populates the "features" key of meta/info.json. + """ + joint_names = robot_model.joint_names + num_joints = robot_model.num_joints + + return { + "observation.images.ego_view": { + "dtype": "video", + "shape": [EGO_VIEW_HEIGHT, EGO_VIEW_WIDTH, 3], + "names": ["height", "width", "channel"], + }, + "observation.state": { + "dtype": "float64", + "shape": (num_joints,), + "names": joint_names, + }, + "observation.eef_state": { + "dtype": "float64", + "shape": (14,), + "names": [ + "left_wrist_pos", + "left_wrist_abs_quat", + "right_wrist_pos", + "right_wrist_abs_quat", + ], + }, + "action.wbc": { + "dtype": "float64", + "shape": (num_joints,), + "names": joint_names, + }, + "observation.root_orientation": { + "dtype": "float64", + "shape": (4,), + "names": ["base_qw", "base_qx", "base_qy", "base_qz"], + }, + "observation.projected_gravity": { + "dtype": "float64", + "shape": (3,), + "names": ["gravity_x", "gravity_y", "gravity_z"], + }, + "observation.cpp_rotation_offset": { + "dtype": "float64", + "shape": (4,), + "names": ["rot_offset_qw", "rot_offset_qx", "rot_offset_qy", "rot_offset_qz"], + }, + "observation.init_base_quat": { + "dtype": "float64", + "shape": (4,), + "names": ["init_base_qw", "init_base_qx", "init_base_qy", "init_base_qz"], + }, + "teleop.delta_heading": { + "dtype": "float64", + "shape": (1,), + "names": ["delta_heading"], + }, + "action.motion_token": { + "dtype": "float64", + "shape": (64,), + "names": "motion_token", + }, + "teleop.smpl_joints": { + "dtype": "float32", + "shape": (72,), + "names": "smpl_joints", + }, + "teleop.smpl_pose": { + "dtype": "float32", + "shape": (63,), + "names": "smpl_pose", + }, + "teleop.body_quat_w": { + "dtype": "float32", + "shape": (4,), + "names": "body_quat_w", + }, + "teleop.target_body_orientation": { + "dtype": "float32", + "shape": (6,), + "names": [ + "target_body_r00", + "target_body_r10", + "target_body_r01", + "target_body_r11", + "target_body_r02", + "target_body_r12", + ], + }, + "teleop.left_hand_joints": { + "dtype": "float32", + "shape": (7,), + "names": "left_hand_joints", + }, + "teleop.right_hand_joints": { + "dtype": "float32", + "shape": (7,), + "names": "right_hand_joints", + }, + "teleop.smpl_frame_index": { + "dtype": "int64", + "shape": (1,), + "names": ["smpl_frame_index"], + }, + "teleop.left_wrist_joints": { + "dtype": "float32", + "shape": (3,), + "names": ["left_wrist_roll", "left_wrist_pitch", "left_wrist_yaw"], + }, + "teleop.right_wrist_joints": { + "dtype": "float32", + "shape": (3,), + "names": ["right_wrist_roll", "right_wrist_pitch", "right_wrist_yaw"], + }, + "teleop.stream_mode": { + "dtype": "int32", + "shape": (1,), + "names": ["stream_mode"], + }, + "teleop.planner_mode": { + "dtype": "int32", + "shape": (1,), + "names": ["locomotion_mode"], + }, + "teleop.planner_movement": { + "dtype": "float32", + "shape": (3,), + "names": ["movement_x", "movement_y", "movement_z"], + }, + "teleop.planner_facing": { + "dtype": "float32", + "shape": (3,), + "names": ["facing_x", "facing_y", "facing_z"], + }, + "teleop.planner_speed": { + "dtype": "float32", + "shape": (1,), + "names": ["speed"], + }, + "teleop.planner_height": { + "dtype": "float32", + "shape": (1,), + "names": ["height"], + }, + "teleop.vr_3pt_position": { + "dtype": "float32", + "shape": (9,), + "names": [ + "lwrist_x", "lwrist_y", "lwrist_z", + "rwrist_x", "rwrist_y", "rwrist_z", + "neck_x", "neck_y", "neck_z", + ], + }, + "teleop.vr_3pt_orientation": { + "dtype": "float32", + "shape": (18,), + "names": [ + "lwrist_r00", "lwrist_r10", "lwrist_r01", "lwrist_r11", "lwrist_r02", "lwrist_r12", + "rwrist_r00", "rwrist_r10", "rwrist_r01", "rwrist_r11", "rwrist_r02", "rwrist_r12", + "neck_r00", "neck_r10", "neck_r01", "neck_r11", "neck_r02", "neck_r12", + ], + }, + } + + +def get_wrist_camera_features() -> dict: + """Features for optional wrist cameras (added when ``record_wrist_cameras`` is enabled).""" + return { + "observation.images.left_wrist": { + "dtype": "video", + "shape": [WRIST_VIEW_HEIGHT, WRIST_VIEW_WIDTH, 3], + "names": ["height", "width", "channel"], + }, + "observation.images.right_wrist": { + "dtype": "video", + "shape": [WRIST_VIEW_HEIGHT, WRIST_VIEW_WIDTH, 3], + "names": ["height", "width", "channel"], + }, + } + + +def get_wrist_camera_modality_config() -> dict: + """Modality config entries for optional wrist cameras.""" + return { + "video": { + "left_wrist": {"original_key": "observation.images.left_wrist"}, + "right_wrist": {"original_key": "observation.images.right_wrist"}, + }, + } + + +def get_g1_robot_model( + waist_location: Literal[ + "lower_body", "upper_body", "lower_and_upper_body" + ] = "lower_and_upper_body", + high_elbow_pose: bool = False, +): + """Instantiate the G1 + ThreeFinger RobotModel for Sonic VLA.""" + from gear_sonic.data.robot_model.instantiation.g1 import instantiate_g1_robot_model + + return instantiate_g1_robot_model( + waist_location=waist_location, + high_elbow_pose=high_elbow_pose, + ) diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/left_ankle_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/left_ankle_pitch_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..69de8490184afc698f633e34bb74e65351bd4689 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/left_ankle_pitch_link.STL differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/left_elbow_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/left_elbow_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..1a96d99ba469960173129084ab6dd3bf8a732a71 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/left_elbow_link.STL differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/left_hand_thumb_0_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/left_hand_thumb_0_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..3028bb4d6e1ae3d30d2504259c08e4106bbacf63 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/left_hand_thumb_0_link.STL differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/right_ankle_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/right_ankle_pitch_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..e77d8a2fe1e5d56fac049833d254d6ffa4f6b350 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/right_ankle_pitch_link.STL differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/right_elbow_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/right_elbow_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..f259e3812efb9985d6463c04d3e8a4b53793a699 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/right_elbow_link.STL differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/right_hand_thumb_0_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/right_hand_thumb_0_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..1cae7f18e16605cb9d6b1d1a0cf6e5c5c360a344 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/right_hand_thumb_0_link.STL differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/right_wrist_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/right_wrist_pitch_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..da194543c40df9d492abb0553c06cc614a78caae Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/right_wrist_pitch_link.STL differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/torso_constraint_L_rod_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/torso_constraint_L_rod_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..6747f3f9341bd72b3803385c135c3ce750322e9d Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/torso_constraint_L_rod_link.STL differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/torso_constraint_R_rod_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/torso_constraint_R_rod_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..95cf415f72f1679a5867ca21a4af774f0b217cad Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/torso_constraint_R_rod_link.STL differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/waist_roll_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/waist_roll_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..65831abd2a64bc8c36e31016964413a3d2116725 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/waist_roll_link.STL differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/waist_roll_link_rev_1_0.STL b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/waist_roll_link_rev_1_0.STL new file mode 100644 index 0000000000000000000000000000000000000000..a64f330c592582dc31cdf38f4d08cfff06681c5f Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/waist_roll_link_rev_1_0.STL differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/video_writer.py b/GR00T-WholeBodyControl/gear_sonic/data/video_writer.py new file mode 100644 index 0000000000000000000000000000000000000000..777b6d28663417a01287b60f28e588cf75f6b383 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/data/video_writer.py @@ -0,0 +1,99 @@ +import os +import queue +import sys +import threading +import time + +import av +import numpy as np + + +class VideoWriter: + def __init__( + self, + output_path: str, + width: int, + height: int, + fps: float, + codec: str = "h264", + buffer_size: int = 50, + ): + self.output_path = output_path + self._first_frame = True + + output_dir = os.path.dirname(output_path) + if output_dir and not os.path.exists(output_dir): + os.makedirs(output_dir, exist_ok=True) + + self.queue = queue.Queue(maxsize=buffer_size) + self.container = av.open(output_path, mode="w") + self.stream = self.container.add_stream(codec, rate=fps) + self.stream.width = width + self.stream.height = height + thread = threading.Thread(target=self._writer_worker, daemon=True) + thread.start() + + def _assert_dimensions(self, frame: np.ndarray) -> None: + assert ( + frame.shape[1] == self.stream.width and frame.shape[0] == self.stream.height + ), ( + f"Incorrect frame dimensions. Input dimensions: {frame.shape[1]}x{frame.shape[0]}. " + f"Expected dimensions: {self.stream.width}x{self.stream.height}" + ) + + def add_frame(self, frame: np.ndarray) -> None: + self._assert_dimensions(frame) + self.queue.put(frame) + + def _writer_worker(self) -> None: + while True: + frame = self.queue.get() + if frame is None: + continue + self._assert_dimensions(frame) + frame = av.VideoFrame.from_ndarray(frame, format="rgb24") + + if self._first_frame: + stderr_fd = sys.stderr.fileno() + old_stderr = os.dup(stderr_fd) + devnull = os.open(os.devnull, os.O_WRONLY) + os.dup2(devnull, stderr_fd) + try: + packets = self.stream.encode(frame) + for packet in packets: + self.container.mux(packet) + finally: + os.dup2(old_stderr, stderr_fd) + os.close(old_stderr) + os.close(devnull) + self._first_frame = False + else: + packets = self.stream.encode(frame) + for packet in packets: + self.container.mux(packet) + + def _flush_stream(self) -> None: + packets = self.stream.encode() + for packet in packets: + self.container.mux(packet) + + def stop(self) -> str: + """Blocking call. Waits for queue to drain, flushes, and closes the container.""" + if not self.queue.empty(): + print("Waiting for video writer queue to empty...") + while not self.queue.empty(): + time.sleep(0.1) + + print("Video writer queue is empty, flushing stream...") + self._flush_stream() + self.container.close() + return self.output_path + + def cancel(self) -> None: + """Immediately stops writing and deletes the output file.""" + if os.path.exists(self.output_path): + os.remove(self.output_path) + self.container.close() + + def __del__(self) -> None: + self.container.close() diff --git a/GR00T-WholeBodyControl/gear_sonic/data_process/convert_soma_csv_to_motion_lib.py b/GR00T-WholeBodyControl/gear_sonic/data_process/convert_soma_csv_to_motion_lib.py new file mode 100644 index 0000000000000000000000000000000000000000..6c7888a175056945de75bf58c3bb6eb652219fb3 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/data_process/convert_soma_csv_to_motion_lib.py @@ -0,0 +1,584 @@ +#!/usr/bin/env python3 # noqa: EXE001 +# ruff: noqa: T201, DOC +"""Convert SOMA retargeter CSV/PKL data to motion_lib format for SONIC training. + +SOMA retargeter outputs G1 29-DOF motion data as CSV files (joint_pos.csv, +body_pos.csv, body_quat.csv) or as a joblib PKL with the same fields. This +script converts that data into the motion_lib PKL format expected by SONIC +training (root_trans_offset, pose_aa, dof, root_rot, fps). + +Supports five input modes: + 1. Single motion directory with CSVs (joint_pos.csv, body_pos.csv, body_quat.csv) + 2. Parent directory containing multiple motion subdirectories + 3. Deploy PKL file (joblib dict with joint_pos, body_pos_w, body_quat_w per sequence) + 4. Directory of flat Bones-SEED CSVs (single CSV per motion, degrees+cm) + 5. Parent directory of session dirs containing Bones-SEED CSVs + +Usage: + # Single CSV directory + python scripts/motion/convert_soma_csv_to_motion_lib.py \ + --input data/soma_retarget/tired_squat_003__A360 \ + --output data/soma_test.pkl --fps 50 + + # Batch: parent dir with multiple motion subdirs + python scripts/motion/convert_soma_csv_to_motion_lib.py \ + --input data/soma_retarget/all_demo_4seqs \ + --output data/soma_demo_4seqs.pkl --fps 50 + + # Deploy PKL file + python scripts/motion/convert_soma_csv_to_motion_lib.py \ + --input data/soma_retarget/bones_test.pkl \ + --output data/soma_bones_test.pkl --fps 50 + + # Bones-SEED: directory of flat CSVs (single session) + python scripts/motion/convert_soma_csv_to_motion_lib.py \ + --input /path/to/bones_SEED/g1/csv/210531 \ + --output data/bones_seed_210531.pkl --fps 50 + + # Bones-SEED: all sessions (parent dir) + python scripts/motion/convert_soma_csv_to_motion_lib.py \ + --input /path/to/bones_SEED/g1/csv \ + --output data/bones_seed_all.pkl --fps 50 +""" + +import argparse +import os +import sys + +import joblib +import numpy as np +from scipy.spatial import transform + +# IsaacLab ↔ MuJoCo joint reordering (29 DOFs for G1). +# MJ_TO_IL[mj] = il: for MuJoCo DOF index mj, gives the IsaacLab index il. +# Source: external_dependencies/SONIC_Web/demo_python.py +MJ_TO_IL = np.array( + [ + 0, + 3, + 6, + 9, + 13, + 17, + 1, + 4, + 7, + 10, + 14, + 18, + 2, + 5, + 8, + 11, + 15, + 19, + 21, + 23, + 25, + 27, + 12, + 16, + 20, + 22, + 24, + 26, + 28, + ], + dtype=np.int32, +) + +# G1 29-DOF axis definitions (from Humanoid_Batch / g1_29dof_rev_1_0.xml). +# Each DOF rotates around a single axis. Hardcoded to avoid torch dependency. +NUM_DOF = 29 +NUM_BODIES = 30 # pelvis + 29 actuated links +DOF_AXIS = np.array( + [ + [0, 1, 0], + [1, 0, 0], + [0, 0, 1], + [0, 1, 0], + [0, 1, 0], + [1, 0, 0], # left leg + [0, 1, 0], + [1, 0, 0], + [0, 0, 1], + [0, 1, 0], + [0, 1, 0], + [1, 0, 0], # right leg + [0, 0, 1], + [1, 0, 0], + [0, 1, 0], # waist + [0, 1, 0], + [1, 0, 0], + [0, 0, 1], + [0, 1, 0], + [1, 0, 0], + [0, 1, 0], + [0, 0, 1], # left arm + [0, 1, 0], + [1, 0, 0], + [0, 0, 1], + [0, 1, 0], + [1, 0, 0], + [0, 1, 0], + [0, 0, 1], # right arm + ], + dtype=np.float32, +) + + +# Joint names in Bones-SEED CSV column order (after Frame + 6 root columns). +# These are in MuJoCo/MJCF actuator order (same as g1_29dof_rev_1_0.xml motors). +BONES_CSV_JOINT_NAMES = [ + "left_hip_pitch_joint_dof", + "left_hip_roll_joint_dof", + "left_hip_yaw_joint_dof", + "left_knee_joint_dof", + "left_ankle_pitch_joint_dof", + "left_ankle_roll_joint_dof", + "right_hip_pitch_joint_dof", + "right_hip_roll_joint_dof", + "right_hip_yaw_joint_dof", + "right_knee_joint_dof", + "right_ankle_pitch_joint_dof", + "right_ankle_roll_joint_dof", + "waist_yaw_joint_dof", + "waist_roll_joint_dof", + "waist_pitch_joint_dof", + "left_shoulder_pitch_joint_dof", + "left_shoulder_roll_joint_dof", + "left_shoulder_yaw_joint_dof", + "left_elbow_joint_dof", + "left_wrist_roll_joint_dof", + "left_wrist_pitch_joint_dof", + "left_wrist_yaw_joint_dof", + "right_shoulder_pitch_joint_dof", + "right_shoulder_roll_joint_dof", + "right_shoulder_yaw_joint_dof", + "right_elbow_joint_dof", + "right_wrist_roll_joint_dof", + "right_wrist_pitch_joint_dof", + "right_wrist_yaw_joint_dof", +] + + +def load_bones_csv(csv_path: str) -> dict: + """Load a single Bones-SEED flat CSV motion file. + + Bones-SEED CSV format: Frame, root_translate{X,Y,Z}, root_rotate{X,Y,Z}, 29 joint DOFs. + All angles in degrees, positions in centimeters. + """ + import pandas as pd + + data = pd.read_csv(csv_path) + T = len(data) + + # Root position: cm → meters + root_pos = ( + np.stack( + [ + data["root_translateX"].values, # noqa: PD011 + data["root_translateY"].values, # noqa: PD011 + data["root_translateZ"].values, # noqa: PD011 + ], + axis=1, + ).astype(np.float32) + / 100.0 + ) # cm → m + + # Root rotation: Euler xyz (intrinsic) degrees → quaternion (xyzw scipy convention) + # Reference: gear_sonic/data_process/process_bones_to_motionlib.py uses "xyz" (intrinsic) + euler_deg = np.stack( + [ + data["root_rotateX"].values, # noqa: PD011 + data["root_rotateY"].values, # noqa: PD011 + data["root_rotateZ"].values, # noqa: PD011 + ], + axis=1, + ).astype(np.float64) + root_quat_xyzw = ( + transform.Rotation.from_euler("xyz", euler_deg, degrees=True).as_quat().astype(np.float32) + ) + # Convert xyzw → wxyz for body_quat_w format + root_quat_wxyz = root_quat_xyzw[:, [3, 0, 1, 2]] + + # Joint DOFs: degrees → radians, already in MuJoCo/MJCF actuator order + joint_cols = [c for c in data.columns if c.endswith("_dof")] + joint_pos_mj = np.deg2rad(data[joint_cols].values).astype(np.float32) # (T, 29) + + # Create dummy body_pos_w and body_quat_w (only root body populated, rest zeros) + # The converter only uses body_pos_w[:,0] for root_trans and body_quat_w[:,0] for root_rot + body_pos_w = np.zeros((T, 14, 3), dtype=np.float32) + body_pos_w[:, 0, :] = root_pos + body_quat_w = np.zeros((T, 14, 4), dtype=np.float32) + body_quat_w[:, :, 0] = 1.0 # identity quaternion wxyz + body_quat_w[:, 0, :] = root_quat_wxyz + + return { + "joint_pos": joint_pos_mj, # (T, 29) MuJoCo order, radians + "body_pos_w": body_pos_w, # (T, 14, 3) + "body_quat_w": body_quat_w, # (T, 14, 4) wxyz + "joint_order": "mj", # already in MuJoCo order, skip IL→MJ reorder + } + + +def load_csv_motion(motion_dir: str) -> dict: + """Load a single motion from a directory of CSV files.""" + joint_pos_f = os.path.join(motion_dir, "joint_pos.csv") + body_pos_f = os.path.join(motion_dir, "body_pos.csv") + body_quat_f = os.path.join(motion_dir, "body_quat.csv") + + if not os.path.exists(joint_pos_f): + return None + + joint_pos = np.loadtxt(joint_pos_f, delimiter=",", skiprows=1, dtype=np.float32) + body_pos = np.loadtxt(body_pos_f, delimiter=",", skiprows=1, dtype=np.float32) + body_quat = np.loadtxt(body_quat_f, delimiter=",", skiprows=1, dtype=np.float32) + + # Reshape body data: (T, 14*3) → (T, 14, 3), (T, 14*4) → (T, 14, 4) + T = joint_pos.shape[0] + body_pos = body_pos.reshape(T, -1, 3) + body_quat = body_quat.reshape(T, -1, 4) + + return { + "joint_pos": joint_pos, # (T, 29) IsaacLab order + "body_pos_w": body_pos, # (T, 14, 3) world frame + "body_quat_w": body_quat, # (T, 14, 4) wxyz format + } + + +def convert_sequence(seq_data: dict, fps: int, humanoid_fk=None) -> dict: # noqa: ARG001 + """Convert a single deploy-format sequence to motion_lib format. + + Args: + seq_data: dict with joint_pos (T, 29), body_pos_w (T, 14, 3), + body_quat_w (T, 14, 4 wxyz) + fps: frame rate of the input data + humanoid_fk: Optional Humanoid_Batch instance (unused, kept for compat) + + Returns: + motion_lib entry dict with root_trans_offset, pose_aa, dof, root_rot, fps + """ + joint_pos = seq_data["joint_pos"] # (T, 29) + body_pos_w = seq_data["body_pos_w"] # (T, 14, 3) + body_quat_w = seq_data["body_quat_w"] # (T, 14, 4) wxyz + joint_order = seq_data.get("joint_order", "il") # "il" or "mj" + + T = joint_pos.shape[0] + + # 1. Root position: body_0 (pelvis) position + root_trans_offset = body_pos_w[:, 0, :].copy() # (T, 3) + + # 2. Root quaternion: body_0 quaternion, convert wxyz → xyzw (scipy convention) + root_quat_wxyz = body_quat_w[:, 0, :] # (T, 4) [w, x, y, z] + root_quat_xyzw = root_quat_wxyz[:, [1, 2, 3, 0]] # (T, 4) [x, y, z, w] + + # 3. Reorder DOFs to MuJoCo order if needed + if joint_order == "il": + # Input is IsaacLab order → reorder to MuJoCo (MJCF actuator order) + dof_mj = joint_pos[:, MJ_TO_IL] # (T, 29) + else: + # Input is already in MuJoCo order (e.g., Bones-SEED CSVs) + dof_mj = joint_pos # (T, 29) + + # 4. Convert DOF → pose_aa using hardcoded G1 axis definitions + dof = dof_mj[:, :NUM_DOF] + + # pose_aa[body_idx] = dof_axis * dof_value (axis-angle representation) + # Body 0 = pelvis (root), bodies 1-29 = actuated joints + pose_aa = np.zeros((T, NUM_BODIES, 3), dtype=np.float32) + # Actuated joints: body idx = dof idx + 1 + pose_aa[:, 1:NUM_BODIES, :] = DOF_AXIS[None, :, :] * dof[:, :, None] + + # Set root rotation as axis-angle + pose_aa[:, 0, :] = transform.Rotation.from_quat(root_quat_xyzw).as_rotvec() + + return { + "root_trans_offset": root_trans_offset.astype(np.float32), + "pose_aa": pose_aa.astype(np.float32), + "dof": dof.astype(np.float32), + "root_rot": root_quat_xyzw.astype(np.float32), # xyzw (scipy convention) + "smpl_joints": np.zeros((T, 24, 3), dtype=np.float32), # placeholder + "fps": fps, + } + + +def downsample_sequence(entry: dict, fps_source: int, fps_target: int) -> dict: + """Downsample a motion_lib entry using stride-based frame skipping. + + Matches process_bones_to_motionlib.py: jump = int(fps_source / fps_target). + Best used when fps_source is an exact multiple of fps_target (e.g. 120→30). + The resulting PKL is stored at fps_target; fk_batch handles the final + resampling to target_fps at load time using the canonical interploate_pose formula. + """ + if fps_source == fps_target: + return entry + jump = int(fps_source / fps_target) + if jump <= 1: + return entry + return { + "root_trans_offset": entry["root_trans_offset"][::jump], + "pose_aa": entry["pose_aa"][::jump], + "dof": entry["dof"][::jump], + "root_rot": entry["root_rot"][::jump], + "smpl_joints": entry["smpl_joints"][::jump], + "fps": fps_target, + } + + +def init_humanoid_fk(): + """Initialize Humanoid_Batch from the G1 MJCF config. + + Only needed for non-Bones-SEED inputs (deploy PKL, SOMA CSV dirs). + Bones-SEED path uses hardcoded DOF_AXIS constants instead. + """ + import omegaconf + + motion_cfg = omegaconf.OmegaConf.create( + { + "asset": { + "assetRoot": "gear_sonic/data/assets/robot_description/mjcf/", + "assetFileName": "g1_29dof_rev_1_0.xml", + "urdfFileName": "", + }, + "extend_config": [], + } + ) + from gear_sonic.utils.motion_lib import torch_humanoid_batch + + return torch_humanoid_batch.Humanoid_Batch(motion_cfg) + + +def process_session_csvs(args_tuple): + """Process all CSVs in a single session directory. Used by multiprocessing.""" + session_dir, session_name, out_dir, fps, fps_source = args_tuple + import warnings + + warnings.filterwarnings("ignore") + + csv_files = sorted([f for f in os.listdir(session_dir) if f.endswith(".csv")]) + + session_out = os.path.join(out_dir, session_name) + os.makedirs(session_out, exist_ok=True) + + converted = 0 + failed = 0 + for csv_f in csv_files: + name = os.path.splitext(csv_f)[0] + out_path = os.path.join(session_out, name + ".pkl") + if os.path.exists(out_path): + converted += 1 # skip existing + continue + try: + seq = load_bones_csv(os.path.join(session_dir, csv_f)) + fps_for_convert = fps_source if fps_source else fps + entry = convert_sequence(seq, fps_for_convert) + if fps_source and fps_source != fps: + entry = downsample_sequence(entry, fps_source, fps) + joblib.dump({name: entry}, out_path, compress=True) + converted += 1 + except Exception: # noqa: BLE001 + failed += 1 + return session_name, converted, failed, len(csv_files) + + +def main(): + parser = argparse.ArgumentParser(description="Convert SOMA CSV/PKL to motion_lib format") + parser.add_argument( + "--input", required=True, help="CSV dir, parent dir of CSV dirs, or deploy PKL" + ) + parser.add_argument( + "--output", required=True, help="Output path (PKL file or directory for individual PKLs)" + ) + parser.add_argument( + "--fps", + type=int, + default=30, + help="Target output FPS (default: 30, matches process_bones_to_motionlib)", + ) + parser.add_argument( + "--fps_source", + type=int, + default=None, + help="Source data FPS. If set and != --fps, data is downsampled. " + "Bones-SEED CSVs are typically 120fps.", + ) + parser.add_argument( + "--individual", + action="store_true", + help="Write individual PKLs per motion (preserves session dir structure)", + ) + parser.add_argument( + "--num_workers", + type=int, + default=8, + help="Number of parallel workers for --individual mode", + ) + args = parser.parse_args() + + print(f"G1 {NUM_DOF} DOFs, {NUM_BODIES} bodies (hardcoded axes)") + + # Individual PKL mode: skip scanning, go straight to parallel per-session processing + if args.individual: + if not os.path.isdir(args.input): + print("ERROR: --individual requires a directory input") + sys.exit(1) + + # Detect: is input a single session dir (contains CSVs) or parent of sessions? + has_csvs = any(f.endswith(".csv") for f in os.listdir(args.input)) + subdirs = sorted( + [d for d in os.listdir(args.input) if os.path.isdir(os.path.join(args.input, d))] + ) + has_session_subdirs = ( + any( + any(f.endswith(".csv") for f in os.listdir(os.path.join(args.input, d))) + for d in subdirs[:3] + ) + if subdirs + else False + ) + + session_dirs = [] + if has_session_subdirs: + for d in subdirs: + subdir = os.path.join(args.input, d) + if any(f.endswith(".csv") for f in os.listdir(subdir)): + session_dirs.append((subdir, d, args.output, args.fps, args.fps_source)) + elif has_csvs: + session_name = os.path.basename(args.input.rstrip("/")) + session_dirs.append((args.input, session_name, args.output, args.fps, args.fps_source)) + + print(f"\nBatch converting {len(session_dirs)} sessions with {args.num_workers} workers") + print(f"Output: {args.output}") + os.makedirs(args.output, exist_ok=True) + + import multiprocessing + + total_converted = 0 + total_failed = 0 + total_csvs = 0 + with multiprocessing.Pool(processes=args.num_workers) as pool: + for session_name, converted, failed, n_csvs in pool.imap_unordered( + process_session_csvs, session_dirs + ): + total_converted += converted + total_failed += failed + total_csvs += n_csvs + print( + f" {session_name}: {converted}/{n_csvs} converted" + + (f" ({failed} failed)" if failed else "") + ) + + print( + f"\nDone: {total_converted} motions converted, {total_failed} failed, {total_csvs} total CSVs" + ) + return + + # Detect input mode (combined PKL output path) + sequences = {} + + if args.input.endswith(".pkl"): + # Mode 3: Deploy PKL file + print(f"Loading deploy PKL: {args.input}") + data = joblib.load(args.input) + for name, seq in data.items(): + sequences[name] = seq + print(f" Found {len(sequences)} sequences") + + elif os.path.isfile(os.path.join(args.input, "joint_pos.csv")): + # Mode 1: Single CSV directory + name = os.path.basename(args.input) + print(f"Loading single CSV motion: {name}") + seq = load_csv_motion(args.input) + if seq is None: + print("ERROR: joint_pos.csv not found") + sys.exit(1) + sequences[name] = seq + print(f" {seq['joint_pos'].shape[0]} frames") + + elif os.path.isdir(args.input): + # Check if directory contains flat CSVs (Bones-SEED format) + csv_files = sorted([f for f in os.listdir(args.input) if f.endswith(".csv")]) + subdirs = sorted( + [d for d in os.listdir(args.input) if os.path.isdir(os.path.join(args.input, d))] + ) + + if csv_files and not any( + os.path.exists(os.path.join(args.input, d, "joint_pos.csv")) + for d in subdirs[:5] # check first 5 subdirs + ): + # Mode 4: Directory of flat Bones-SEED CSVs + print(f"Scanning directory for Bones-SEED CSVs: {args.input}") + for csv_f in csv_files: + csv_path = os.path.join(args.input, csv_f) + name = os.path.splitext(csv_f)[0] + try: + seq = load_bones_csv(csv_path) + sequences[name] = seq + except Exception as e: # noqa: BLE001 + print(f" WARNING: Failed to load {csv_f}: {e}") + print(f" Found {len(sequences)} Bones-SEED CSV motions") + elif subdirs: + # Check if subdirs contain flat CSVs (batch of session dirs) + has_session_csvs = False + for dname in subdirs[:3]: + subdir = os.path.join(args.input, dname) + sub_csvs = [f for f in os.listdir(subdir) if f.endswith(".csv")] + if sub_csvs and not os.path.exists(os.path.join(subdir, "joint_pos.csv")): + has_session_csvs = True + break + + if has_session_csvs: + # Mode 5: Parent dir of session dirs containing Bones-SEED CSVs + print(f"Scanning session directories for Bones-SEED CSVs: {args.input}") + for dname in sorted(subdirs): + subdir = os.path.join(args.input, dname) + sub_csvs = sorted([f for f in os.listdir(subdir) if f.endswith(".csv")]) + for csv_f in sub_csvs: + csv_path = os.path.join(subdir, csv_f) + name = os.path.splitext(csv_f)[0] + try: + seq = load_bones_csv(csv_path) + sequences[name] = seq + except Exception as e: # noqa: BLE001 + print(f" WARNING: Failed to load {dname}/{csv_f}: {e}") + if sub_csvs: + print(f" Session {dname}: {len(sub_csvs)} CSVs") + print(f" Found {len(sequences)} total Bones-SEED CSV motions") + else: + # Mode 2: Parent directory with SOMA-style subdirectories + print(f"Scanning directory: {args.input}") + for dname in sorted(subdirs): + subdir = os.path.join(args.input, dname) + seq = load_csv_motion(subdir) + if seq is not None: + sequences[dname] = seq + print(f" Found {len(sequences)} motion directories with CSVs") + else: + print(f"ERROR: {args.input} is not a valid input") + sys.exit(1) + + if not sequences: + print("ERROR: No sequences found") + sys.exit(1) + + # Convert each sequence (combined PKL mode) + motion_lib_dict = {} + for name, seq_data in sequences.items(): + T = seq_data["joint_pos"].shape[0] + print(f" Converting {name}: {T} frames @ {args.fps} fps") + fps_for_convert = args.fps_source if args.fps_source else args.fps + entry = convert_sequence(seq_data, fps_for_convert) + if args.fps_source and args.fps_source != args.fps: + entry = downsample_sequence(entry, args.fps_source, args.fps) + motion_lib_dict[name] = entry + + # Save + os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True) + print(f"\nSaving motion_lib PKL: {args.output}") + joblib.dump(motion_lib_dict, args.output, compress=True) + print(f"Done: {len(motion_lib_dict)} sequences saved") + + +if __name__ == "__main__": + main() diff --git a/GR00T-WholeBodyControl/gear_sonic/data_process/extract_soma_joints_from_bvh.py b/GR00T-WholeBodyControl/gear_sonic/data_process/extract_soma_joints_from_bvh.py new file mode 100644 index 0000000000000000000000000000000000000000..c4f0f15ee09ef46eb7feb7472ed99b63bab39066 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/data_process/extract_soma_joints_from_bvh.py @@ -0,0 +1,394 @@ +#!/usr/bin/env python3 +# ruff: noqa: T201, BLE001, DOC +"""Extract SOMA skeleton FK joint positions from BVH files. + +Parses NOVA-skeleton BVH files and computes forward kinematics to extract +world-space 3D joint positions for a selected 26-joint subset. Outputs +per-motion PKL files in the same directory structure as the robot PKLs. + +The 26 selected joints cover the major body landmarks with symmetric arms +(including Thumb1 + Middle1 per hand for orientation): hips, spine chain, +shoulders, arms, hands+fingers, legs, feet. + +Input: BVH files from bones_update_240924/anims_uniform_novaskel_v1/BVH/ +Output: Per-motion PKL files with soma_joints (T, 26, 3) Z-up meters body-local, + soma_root_quat (T, 4) wxyz Y-up BVH world rotation + +Usage: + # Single session + python scripts/motion/extract_soma_joints_from_bvh.py \ + --input /path/to/novaskel_v1/BVH/210531 \ + --output /path/to/output/bones_soma_joints/210531 \ + --fps 30 + + # All sessions (parent dir) + python scripts/motion/extract_soma_joints_from_bvh.py \ + --input /path/to/novaskel_v1/BVH \ + --output /path/to/output/bones_soma_joints \ + --fps 30 --num_workers 8 +""" + +import argparse +import glob +import multiprocessing +import os +import os.path as osp +import re +import sys +import time + +import joblib +import numpy as np +from scipy.spatial import transform + +# 26-joint subset of the 78-joint NOVA skeleton (Root excluded). +# Covers major body landmarks, excluding most fingers, face details, end sites. +# Arms are fully symmetric with two finger joints per hand (Thumb1 + Middle1) +# to determine hand orientation. +SOMA_JOINTS = [ + "Hips", # 0 - pelvis + "Spine1", # 1 - lower spine + "Spine2", # 2 - mid spine + "Chest", # 3 - upper spine + "Neck1", # 4 - neck + "Head", # 5 - head + "LeftShoulder", # 6 - left clavicle + "LeftArm", # 7 - left upper arm + "LeftForeArm", # 8 - left elbow + "LeftHand", # 9 - left wrist + "LeftHandThumb1", # 10 - left thumb (hand orientation) + "LeftHandMiddle1", # 11 - left middle finger (hand orientation) + "RightShoulder", # 12 - right clavicle + "RightArm", # 13 - right upper arm + "RightForeArm", # 14 - right elbow + "RightHand", # 15 - right wrist + "RightHandThumb1", # 16 - right thumb (hand orientation) + "RightHandMiddle1", # 17 - right middle finger (hand orientation) + "LeftLeg", # 18 - left hip / upper leg + "LeftShin", # 19 - left knee + "LeftFoot", # 20 - left ankle + "LeftToeBase", # 21 - left toe + "RightLeg", # 22 - right hip / upper leg + "RightShin", # 23 - right knee + "RightFoot", # 24 - right ankle + "RightToeBase", # 25 - right toe +] + +NUM_SOMA_JOINTS = len(SOMA_JOINTS) + + +def parse_bvh(filepath): + """Parse BVH hierarchy and motion data. + + Returns: + joints: list of dicts with name, offset, channels, parent_idx + channel_order: list of (joint_idx, channel_name) tuples + motion_data: (n_frames, n_channels) numpy array + n_frames: int + frame_time: float (seconds per frame) + """ + with open(filepath) as f: + lines = f.readlines() + + joints = [] + joint_stack = [] + channel_order = [] + i = 0 + + while i < len(lines): + line = lines[i].strip() + if line == "MOTION": + i += 1 + break + + m = re.match(r"(ROOT|JOINT)\s+(\S+)", line) + if m: + name = m.group(2) + parent_idx = joint_stack[-1] if joint_stack else -1 + joints.append({"name": name, "offset": None, "channels": [], "parent_idx": parent_idx}) + joint_stack.append(len(joints) - 1) + elif line.startswith("OFFSET") and joint_stack: + vals = [float(x) for x in line.split()[1:]] + joints[joint_stack[-1]]["offset"] = np.array(vals) + elif line.startswith("CHANNELS") and joint_stack: + parts = line.split() + n_ch = int(parts[1]) + ch_names = parts[2 : 2 + n_ch] + joints[joint_stack[-1]]["channels"] = ch_names + for ch in ch_names: + channel_order.append((joint_stack[-1], ch)) + elif line == "}": + if joint_stack: + joint_stack.pop() + i += 1 + + # Parse MOTION section + frames_line = lines[i].strip() + n_frames = int(frames_line.split(":")[1]) + i += 1 + frame_time = float(lines[i].strip().split(":")[1]) + i += 1 + + motion_data = np.empty((n_frames, len(channel_order))) + for f_idx in range(n_frames): + vals = lines[i].strip().split() + motion_data[f_idx] = [float(v) for v in vals] + i += 1 + + return joints, channel_order, motion_data, n_frames, frame_time + + +def compute_fk_selected(joints, channel_order, motion_data, selected_names): + """Compute FK world positions for selected joints only. + + Uses vectorized rotation computation per joint across all frames. + + Args: + joints: parsed joint hierarchy + channel_order: channel mapping + motion_data: (n_frames, n_channels) + selected_names: list of joint names to extract + + Returns: + selected_positions: (n_frames, len(selected_names), 3) in BVH units (cm) + root_quats: (n_frames, 4) root orientation quaternions (xyzw) + """ + n_frames = motion_data.shape[0] + n_joints = len(joints) + joint_names = [j["name"] for j in joints] + + # Build selected indices + selected_indices = set() + for name in selected_names: + if name in joint_names: + selected_indices.add(joint_names.index(name)) + + # Also include all ancestors needed for FK + ancestors = set() + for idx in selected_indices: + j = idx + while j >= 0: + ancestors.add(j) + j = joints[j]["parent_idx"] + compute_joints = sorted(ancestors | selected_indices) + + # Pre-compute per-joint channel indices + joint_channels = {j: [] for j in range(n_joints)} + for ch_idx, (j_idx, ch_name) in enumerate(channel_order): + joint_channels[j_idx].append((ch_idx, ch_name)) + + # Compute FK for all frames + world_rots = np.zeros((n_frames, n_joints, 3, 3)) + world_pos = np.zeros((n_frames, n_joints, 3)) + + for j_idx in compute_joints: + joint = joints[j_idx] + offset = joint["offset"] if joint["offset"] is not None else np.zeros(3) + + # Extract position and rotation channels + pos_channels = {} + rot_order = "" + rot_ch_indices = [] + for ch_idx, ch_name in joint_channels[j_idx]: + if ch_name.endswith("position"): + pos_channels[ch_name] = ch_idx + elif ch_name.endswith("rotation"): + rot_order += ch_name[0].lower() + rot_ch_indices.append(ch_idx) + + # Local position (all frames) + has_pos_channels = bool(pos_channels) + if has_pos_channels: + # Joints with position channels: use channels directly (not additive to offset) + local_pos = np.zeros((n_frames, 3)) + if "Xposition" in pos_channels: + local_pos[:, 0] = motion_data[:, pos_channels["Xposition"]] + if "Yposition" in pos_channels: + local_pos[:, 1] = motion_data[:, pos_channels["Yposition"]] + if "Zposition" in pos_channels: + local_pos[:, 2] = motion_data[:, pos_channels["Zposition"]] + else: + # Joints with only rotation channels: use static offset + local_pos = np.tile(offset, (n_frames, 1)) + + # Local rotation (all frames) + # BVH uses extrinsic rotations: uppercase in scipy convention + if rot_order: + rot_vals = motion_data[:, rot_ch_indices] # (n_frames, n_rot_channels) + local_rot = transform.Rotation.from_euler( + rot_order.upper(), rot_vals, degrees=True + ).as_matrix() + else: + local_rot = np.tile(np.eye(3), (n_frames, 1, 1)) + + if joint["parent_idx"] < 0: + # Root joint: no parent transform + world_rots[:, j_idx] = local_rot + world_pos[:, j_idx] = local_pos + else: + p = joint["parent_idx"] + parent_rot = world_rots[:, p] # (n_frames, 3, 3) + parent_pos = world_pos[:, p] # (n_frames, 3) + # world_pos = parent_pos + parent_rot @ local_pos + world_pos[:, j_idx] = parent_pos + np.einsum("fij,fj->fi", parent_rot, local_pos) + # world_rot = parent_rot @ local_rot + world_rots[:, j_idx] = np.einsum("fij,fjk->fik", parent_rot, local_rot) + + # Extract selected joints + sel_indices = [joint_names.index(name) for name in selected_names if name in joint_names] + selected_positions = world_pos[:, sel_indices, :] # (n_frames, len(selected_names), 3) + + # Extract root quaternion (Hips joint, index 1) + hips_idx = joint_names.index("Hips") if "Hips" in joint_names else 0 + root_quats_scipy = transform.Rotation.from_matrix(world_rots[:, hips_idx]) + root_quats = root_quats_scipy.as_quat() # (n_frames, 4) as xyzw + + return selected_positions, root_quats + + +def process_single_bvh(args): + """Process a single BVH file → PKL with soma_joints. + + Returns (motion_name, success, error_msg) + """ + bvh_path, output_dir, fps_target, skip_existing = args + motion_name = osp.splitext(osp.basename(bvh_path))[0] + output_path = osp.join(output_dir, f"{motion_name}.pkl") + + if skip_existing and osp.exists(output_path): + return motion_name, True, "skipped" + + try: + joints, channel_order, motion_data, n_frames, frame_time = parse_bvh(bvh_path) + fps_source = round(1.0 / frame_time) + + positions, root_quats = compute_fk_selected(joints, channel_order, motion_data, SOMA_JOINTS) + + # Convert cm → meters + positions_m = positions / 100.0 + + # Extract hips translation and subtract from all joints to get + # body-local positions (matching SMPL's compute_human_joints which + # produces joints without global translation). + hips_idx = 0 # Hips is joint index 0 in SOMA_JOINTS (Root removed) + transl = positions_m[:, hips_idx, :].copy() # (T, 3) Y-up + positions_m = positions_m - transl[:, None, :] # body-local + + # Convert Y-up → Z-up: (x, y, z) → (x, -z, y) + # Same as applying rot90x, matching SMPL's convert_smpl_bones which + # applies rot90x to global_orient before FK to produce Z-up joints. + positions_zup = positions_m.copy() + positions_zup[..., 1] = -positions_m[..., 2] + positions_zup[..., 2] = positions_m[..., 1] + + # Downsample to target fps using stride-based frame skipping. + # Matches convert_soma_csv_to_motion_lib.py: jump = int(fps_source / fps_target). + # For Bones-SEED (120fps BVH → 30fps), this is stride-4 (exact division). + # Both BVH and CSV have identical source frame counts at 120fps, so + # stride-based downsampling produces identical frame counts. + if fps_source != fps_target: + jump = max(1, int(fps_source / fps_target)) + positions_zup = positions_zup[::jump] + transl = transl[::jump] + root_quats = root_quats[::jump] + + # Convert xyzw → wxyz for compatibility with IsaacLab quat pipeline. + # Root quats stay Y-up — runtime converts via smpl_root_ytoz_up + + # remove_bvh_base_rot (same pattern as SMPL pose_aa). + root_quats = root_quats[:, [3, 0, 1, 2]] + + # Store as PKL + entry = { + motion_name: { + "soma_joints": positions_zup.astype( + np.float32 + ), # (T, 26, 3) Z-up meters, body-local + "soma_root_quat": root_quats.astype( + np.float32 + ), # (T, 4) wxyz, Y-up BVH world rotation + "soma_transl": transl.astype(np.float32), # (T, 3) Hips world position, Y-up + "fps": fps_target, + "joint_names": SOMA_JOINTS, + } + } + + os.makedirs(output_dir, exist_ok=True) + joblib.dump(entry, output_path) + return motion_name, True, None + + except Exception as e: + return motion_name, False, str(e) + + +def main(): + parser = argparse.ArgumentParser(description="Extract SOMA joints from BVH files") + parser.add_argument("--input", required=True, help="BVH dir (session or parent)") + parser.add_argument("--output", required=True, help="Output dir for PKL files") + parser.add_argument( + "--fps", + type=int, + default=30, + help="Target FPS (default: 30, matches process_bones_to_motionlib)", + ) + parser.add_argument("--num_workers", type=int, default=4, help="Parallel workers") + parser.add_argument("--skip_existing", action="store_true", help="Skip existing PKL files") + args = parser.parse_args() + + # Discover BVH files + bvh_files = sorted(glob.glob(osp.join(args.input, "*.bvh"))) + + if bvh_files: + # Single session directory + sessions = {osp.basename(args.input): bvh_files} + else: + # Parent directory with session subdirs + session_dirs = sorted([d for d in glob.glob(osp.join(args.input, "*")) if osp.isdir(d)]) + sessions = {} + for sd in session_dirs: + files = sorted(glob.glob(osp.join(sd, "*.bvh"))) + if files: + sessions[osp.basename(sd)] = files + + if not sessions: + print(f"No BVH files found in {args.input}") + sys.exit(1) + + total_bvh = sum(len(v) for v in sessions.values()) + print(f"Found {total_bvh} BVH files across {len(sessions)} sessions") + print(f"Output: {args.output}, FPS: {args.fps}, Workers: {args.num_workers}") + total_converted = 0 + total_failed = 0 + t0 = time.time() + + for session_name, bvh_list in sessions.items(): + session_output = osp.join(args.output, session_name) + tasks = [(bvh_path, session_output, args.fps, args.skip_existing) for bvh_path in bvh_list] + + with multiprocessing.Pool(args.num_workers) as pool: + results = pool.map(process_single_bvh, tasks) + + converted = sum(1 for _, s, e in results if s and e != "skipped") + skipped = sum(1 for _, s, e in results if e == "skipped") + failed = sum(1 for _, s, _ in results if not s) + + if failed > 0: + for name, success, err in results: + if not success: + print(f" FAILED: {name}: {err}") + + total_converted += converted + skipped + total_failed += failed + + elapsed = time.time() - t0 + rate = total_converted / elapsed if elapsed > 0 else 0 + print( + f" {session_name}: {converted} converted, {skipped} skipped, " + f"{failed} failed [{total_converted}/{total_bvh}, {rate:.0f}/s]" + ) + + elapsed = time.time() - t0 + print(f"\nDone: {total_converted} converted, {total_failed} failed, " f"{elapsed:.1f}s elapsed") + + +if __name__ == "__main__": + main() diff --git a/GR00T-WholeBodyControl/gear_sonic/data_process/filter_and_copy_bones_data.py b/GR00T-WholeBodyControl/gear_sonic/data_process/filter_and_copy_bones_data.py new file mode 100644 index 0000000000000000000000000000000000000000..f09061012a1f1d1da3ef35d614440629138659cd --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/data_process/filter_and_copy_bones_data.py @@ -0,0 +1,294 @@ +#!/usr/bin/env python3 +""" +Script to filter and copy bones data from bones_gmr to single_pkls directory. + +This script copies motion files from the bones_gmr directory structure while +filtering out unwanted sequences based on keywords in filenames. +It preserves the bones_xxx directory structure in the destination. +""" + +import argparse +from functools import partial +import glob +from multiprocessing import Pool, cpu_count +import os.path as osp +from pathlib import Path +import shutil + +from tqdm import tqdm + + +def should_filter_out(filename, filter_keywords, include_keywords=None): + """ + Check if a filename contains any of the filter keywords. + + Args: + filename (str): The filename to check + filter_keywords (list): List of keywords to filter out + + Returns: + bool: True if file should be filtered out, False otherwise + """ + filename_lower = filename.lower() + if include_keywords is None: + return any(keyword.lower() in filename_lower for keyword in filter_keywords) + else: + return any(keyword.lower() in filename_lower for keyword in filter_keywords) or ( + not any(keyword.lower() in filename_lower for keyword in include_keywords) + ) + + +def process_bones_directory( + bones_dir, dest_path, filter_keywords, dry_run, verbose, include_keywords=None +): + """ + Process a single bones directory - worker function for multiprocessing. + + Args: + bones_dir (Path): Source bones directory to process + dest_path (Path): Destination base directory + filter_keywords (list): Keywords to filter out + dry_run (bool): If True, don't actually copy files + verbose (bool): If True, show detailed output + + Returns: + tuple: (total_files, copied_files, filtered_files) + """ + dest_bones_dir = dest_path / bones_dir.name + + if not dry_run: + dest_bones_dir.mkdir(parents=True, exist_ok=True) + + # Find all pkl files in this bones directory + + pkl_files = list(glob.glob(osp.join(bones_dir, "**", "*.pkl"), recursive=True)) + total_files = len(pkl_files) + copied_files = 0 + filtered_files = 0 + + if verbose: + print(f"Processing {bones_dir.name}: {total_files} files") + + for pkl_file in pkl_files: + base = osp.basename(pkl_file) + parent = osp.basename(osp.dirname(pkl_file)) + name_to_check = f"{parent}/{base}" + + if ( + should_filter_out(name_to_check, filter_keywords, include_keywords) + and not base == "metadata.pkl" + ): + filtered_files += 1 + if verbose: + print(f" FILTERED: {osp.basename(pkl_file)}") + else: + copied_files += 1 + dest_file = osp.join(dest_bones_dir, osp.basename(pkl_file)) + + if not dry_run: + shutil.copy2(pkl_file, dest_file) + + return (total_files, copied_files, filtered_files) + + +def copy_filtered_bones_data( + source_dir, + dest_dir, + filter_keywords, + dry_run=False, + verbose=False, + workers=None, + filter_file=None, +): + """ + Copy bones data while filtering out unwanted sequences. + + Args: + source_dir (Path): Source directory containing bones_xxx subdirs + dest_dir (Path): Destination directory + filter_keywords (list): Keywords to filter out + dry_run (bool): If True, only show what would be copied + verbose (bool): If True, show detailed output + workers (int): Number of worker processes. If None, uses all CPU cores. + """ + source_path = Path(source_dir) + dest_path = Path(dest_dir) + include_keywords = None + if filter_file is not None: + with open(filter_file) as f: + include_keywords = f.read().splitlines() + + if not source_path.exists(): + print(f"Error: Source directory {source_path} does not exist!") + return False + + # Find all bones_xxx directories + bones_dirs = [d for d in source_path.iterdir() if d.is_dir()] + + if not bones_dirs: + print(f"No bones_xxx directories found in {source_path}") + return False + + # Determine number of workers + if workers is None: + workers = cpu_count() + workers = max(1, min(workers, len(bones_dirs))) # Don't use more workers than directories + + print(f"Found {len(bones_dirs)} bones directories to process") + print(f"Using {workers} worker processes") + + if verbose: + for d in bones_dirs: + print(f" {d.name}") + + # Process directories in parallel + worker_func = partial( + process_bones_directory, + dest_path=dest_path, + filter_keywords=filter_keywords, + dry_run=dry_run, + verbose=verbose, + include_keywords=include_keywords, + ) + + total_files = 0 + copied_files = 0 + filtered_files = 0 + if workers == 1: + # Single-threaded execution for easier debugging + results = [] + for bones_dir in tqdm(bones_dirs, desc="Processing bones directories"): + results.append(worker_func(bones_dir)) + else: + # Multi-process execution + with Pool(processes=workers) as pool: + results = list( + tqdm( + pool.imap(worker_func, bones_dirs), + total=len(bones_dirs), + desc="Processing bones directories", + ) + ) + + # Aggregate results + for result in results: + total_files += result[0] + copied_files += result[1] + filtered_files += result[2] + + # Summary + print(f"\n{'='*60}") + print("FILTERING SUMMARY") + print(f"{'='*60}") + print(f"Source directory: {source_path}") + print(f"Destination directory: {dest_path}") + print(f"Total files found: {total_files}") + print(f"Files copied: {copied_files}") + print(f"Files filtered out: {filtered_files}") + print(f"Filter keywords: {', '.join(filter_keywords)}") + + if dry_run: + print("\nDRY RUN - No files were actually copied") + else: + print(f"\nFiles successfully copied to: {dest_path}") + + return True + + +def main(): + parser = argparse.ArgumentParser(description="Filter and copy bones data") + parser.add_argument( + "--source", + default="data/bones_gmr/0903_all/", + help="Source directory containing bones_xxx subdirectories", + ) + parser.add_argument( + "--dest", default="data/single_pkls/", help="Destination directory for filtered bones data" + ) + parser.add_argument( + "--filter-keywords", + default=[ + "bed", + "bike", + "chair", + "climb", + "com_up_50cm", + "sitting", + "step_on", + "seat", + "table", + "_sit_", + "sit_", "ladder", + "crutch", + "_bed_", + "_ride_", + "scooter", + "stepdown", + "acrobatics_", + "box_HSPU", + "cartwheel", + "50cm_box_", + "on_box", "fall_from", + "handstand_ff_", + "on_1m", + "form_box", + "off_1m", + "230m", + "jump_over_obstacle_", + "lift_crate_come_up_", + "jump_to_shoulder_roll", + "kozak_dance", + "stair", + "handstand", + "box_jump", + "monkey_jump", + "safety_roll", + "box_dips", + "walking_on_edge", + "push_obstacle", + ], + nargs="+", + help="Keywords to filter out from filenames", + ) + parser.add_argument( + "--dry-run", action="store_true", help="Show what would be copied without actually copying" + ) + parser.add_argument("--verbose", "-v", action="store_true", help="Show detailed output") + parser.add_argument( + "--add-keywords", nargs="+", help="Additional keywords to add to the default filter list" + ) + parser.add_argument( + "--workers", + type=int, + default=None, + help="Number of worker processes (default: use all CPU cores)", + ) + + parser.add_argument("--filter_file", default=None, help="Filter file to use") + + args = parser.parse_args() + + # Combine default and additional keywords + filter_keywords = args.filter_keywords + if args.add_keywords: + filter_keywords.extend(args.add_keywords) + + print(f"Filtering out files containing: {', '.join(filter_keywords)}") + + success = copy_filtered_bones_data( + source_dir=args.source, + dest_dir=args.dest, + filter_keywords=filter_keywords, + dry_run=args.dry_run, + verbose=args.verbose, + workers=args.workers, + filter_file=args.filter_file, + ) + + return 0 if success else 1 + + +if __name__ == "__main__": + import sys + + sys.exit(main()) diff --git a/GR00T-WholeBodyControl/gear_sonic/data_process/split_pkl_files.py b/GR00T-WholeBodyControl/gear_sonic/data_process/split_pkl_files.py new file mode 100644 index 0000000000000000000000000000000000000000..420aad005dd0628272238315be06f5cdd718fbe8 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/data_process/split_pkl_files.py @@ -0,0 +1,193 @@ +#!/usr/bin/env python3 +""" +Script to break large pickle files into individual motion sequence files. + +This script reads large pickle files and breaks each motion sequence +within them into individual pickle files in subdirectories. +This enables motion_lib_base.py to use directory mode for efficient loading. +Supports any pkl file format, not just bone motion files. +""" + +import argparse +from pathlib import Path +import shutil +import sys +import time + +import joblib +from tqdm import tqdm + + +def create_output_structure(output_dir, clean=False): + """Create the output directory structure.""" + output_path = Path(output_dir) + + if clean and output_path.exists(): + print(f"Removing existing output directory: {output_path}") + shutil.rmtree(output_path) + + output_path.mkdir(parents=True, exist_ok=True) + print(f"Created output directory: {output_path}") + return output_path + + +def extract_motion_metadata(motion_data): + """Extract metadata (length, fps) from motion data.""" + metadata = {} + + # Check common fields that might indicate fps + fps = None + if hasattr(motion_data, "get"): + fps = motion_data.get("fps", motion_data.get("frame_rate", motion_data.get("framerate"))) + + # If no fps found, try to infer from common values or set default + if fps is None: + fps = 30.0 # Default fps + + # Get length - check for common motion data structures + length = 0 + length = motion_data["root_trans_offset"].shape[0] + + # If still no length found, try to get it from the data structure + if length == 0 and hasattr(motion_data, "__len__"): + length = len(motion_data) + + return {"length": length, "fps": fps, "duration": length / fps if fps > 0 else 0.0} + + +def process_motion_file(input_file, output_dir, verbose=False): + """Process a single large motion file and break it into individual files.""" + input_path = Path(input_file) + + # Create subdirectory named after the input file + file_subdir = input_path.stem + output_path = Path(output_dir) / file_subdir + output_path.mkdir(parents=True, exist_ok=True) + + try: + print(f"Loading {input_path.name}...") + motion_data = joblib.load(input_path) + print(f"Loaded {len(motion_data)} motion sequences -> {output_path}") + + # Collect metadata for all motion sequences + metadata = {} + + # Process each motion sequence + for motion_key, motion_sequence_data in tqdm( + motion_data.items(), desc=f"Processing {input_path.name}" + ): + individual_filepath = output_path / f"{motion_key}.pkl" + individual_dict = {motion_key: motion_sequence_data} + joblib.dump(individual_dict, individual_filepath) + + # Extract metadata for this motion sequence + motion_metadata = extract_motion_metadata(motion_sequence_data) + metadata[motion_key] = motion_metadata + + # Save metadata file + metadata_filepath = output_path / "metadata.pkl" + joblib.dump(metadata, metadata_filepath) + + if verbose: + print( + f"Successfully processed {input_path.name} -> {len(motion_data)} individual files + metadata" + ) + return True + + except Exception as e: + print(f"Error processing {input_path.name}: {e}") + return False + + +def main(): + parser = argparse.ArgumentParser(description="Break pickle files into individual sequences") + parser.add_argument("input", help="Input directory containing pkl files or single pkl file") + parser.add_argument( + "--output", + default="data/processed_pkl/", + help="Output directory for individual motion files", + ) + parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose output") + parser.add_argument( + "--file-pattern", + default="*.pkl", + help="Pattern to match input files (only used for directories)", + ) + parser.add_argument( + "--clean", action="store_true", help="Remove output directory if it already exists" + ) + + args = parser.parse_args() + + input_path = Path(args.input) + if not input_path.exists(): + print(f"Error: Input path {input_path} does not exist!") + return 1 + + # Create output directory + output_path = create_output_structure(args.output, args.clean) + + # Determine input files + if input_path.is_file(): + # Single file input + if not input_path.suffix == ".pkl": + print(f"Error: Input file must be a .pkl file, got {input_path.suffix}") + return 1 + input_files = [input_path] + print(f"Processing single file: {input_path}") + else: + # Directory input + input_files = sorted(input_path.glob(args.file_pattern)) + if not input_files: + print(f"No files found matching pattern {args.file_pattern} in {input_path}") + return 1 + print(f"Processing directory: {input_path}") + + print(f"Found {len(input_files)} files to process:") + for f in input_files[:5]: + print(f" {f.name}") + if len(input_files) > 5: + print(f" ... and {len(input_files) - 5} more files") + + # Process all files + successful = failed = total_individual_files = 0 + start_time = time.time() + + for input_file in input_files: + print(f"\n{'='*60}") + if process_motion_file(input_file, output_path, args.verbose): + successful += 1 + # Count files created (excluding metadata) + subdir_path = output_path / input_file.stem + individual_files = [ + f for f in subdir_path.glob("*.pkl") if not f.name.endswith("metadata.pkl") + ] + total_individual_files += len(individual_files) + print( + f"Created {len(individual_files)} individual files + metadata from {input_file.name}" + ) + else: + failed += 1 + + elapsed = time.time() - start_time + + # Summary + print(f"\n{'='*60}") + print("PROCESSING SUMMARY") + print(f"{'='*60}") + print(f"Input: {input_path}") + print(f"Output directory: {output_path}") + print(f"Files processed successfully: {successful}") + print(f"Files failed: {failed}") + print(f"Total individual motion files created: {total_individual_files}") + print(f"Processing time: {elapsed:.2f} seconds") + + if successful > 0: + print(f"\nSuccess! Individual motion files are available in: {output_path}") + print(f"To use with motion_lib_base.py: motion_file = '{output_path}'") + + return 0 if failed == 0 else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/__init__.py b/GR00T-WholeBodyControl/gear_sonic/envs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__init__.py b/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/GR00T-WholeBodyControl/gear_sonic/isaac_utils/maths.py b/GR00T-WholeBodyControl/gear_sonic/isaac_utils/maths.py new file mode 100644 index 0000000000000000000000000000000000000000..50efe6972e1787ace0401c0360f1f1848dab562f --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/isaac_utils/maths.py @@ -0,0 +1,52 @@ +"""Basic math primitives for Isaac Lab: normalization, random floats, copysign, seeding.""" + +import torch +import numpy as np +import random +import os + + +@torch.jit.script +def normalize(x, eps: float = 1e-9): + return x / x.norm(p=2, dim=-1).clamp(min=eps, max=None).unsqueeze(-1) + + +@torch.jit.script +def torch_rand_float(lower, upper, shape, device): + # type: (float, float, Tuple[int, int], str) -> Tensor + return (upper - lower) * torch.rand(*shape, device=device) + lower + + +@torch.jit.script +def copysign(a, b): + # type: (float, Tensor) -> Tensor + a = torch.tensor(a, device=b.device, dtype=torch.float).repeat(b.shape[0]) + return torch.abs(a) * torch.sign(b) + + +def set_seed(seed, torch_deterministic=False): + """set seed across modules""" + if seed == -1 and torch_deterministic: + seed = 42 + elif seed == -1: + seed = np.random.randint(0, 10000) + print("Setting seed: {}".format(seed)) + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + os.environ["PYTHONHASHSEED"] = str(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + if torch_deterministic: + # refer to https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + torch.use_deterministic_algorithms(True) + else: + torch.backends.cudnn.benchmark = True + torch.backends.cudnn.deterministic = False + + return seed diff --git a/GR00T-WholeBodyControl/gear_sonic/isaac_utils/rotations.py b/GR00T-WholeBodyControl/gear_sonic/isaac_utils/rotations.py new file mode 100644 index 0000000000000000000000000000000000000000..71ac8208dcd8c8384a8925ee0621171d46361ee8 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/isaac_utils/rotations.py @@ -0,0 +1,787 @@ +"""JIT-compiled quaternion and rotation utilities for Isaac environments. + +Provides quaternion arithmetic (multiply, inverse, conjugate, slerp), conversions +(axis-angle, rotation matrix, euler), and specialized helpers for SMPL root +orientation transforms (Y-up to Z-up, base rotation removal). +""" + +import torch +from torch import Tensor +import torch.nn.functional as F +from gear_sonic.isaac_utils.maths import ( + normalize, + copysign, +) +from gear_sonic.trl.utils.torch_transform import angle_axis_to_quaternion, quaternion_to_angle_axis +from typing import Tuple +import numpy as np +from typing import List, Optional + + +@torch.jit.script +def quat_unit(a): + """Normalize quaternion to unit length.""" + return normalize(a) + + +@torch.jit.script +def quat_apply(a: Tensor, b: Tensor, w_last: bool) -> Tensor: + shape = b.shape + a = a.reshape(-1, 4) + b = b.reshape(-1, 3) + if w_last: + xyz = a[:, :3] + w = a[:, 3:] + else: + xyz = a[:, 1:] + w = a[:, :1] + t = xyz.cross(b, dim=-1) * 2 + return (b + w * t + xyz.cross(t, dim=-1)).view(shape) + + +def get_yaw_quat_from_quat(quat_angle): + rpy = get_euler_xyz_in_tensor(quat_angle) + roll, pitch, yaw = rpy[:, 0], rpy[:, 1], rpy[:, 2] + roll = torch.zeros_like(roll) + pitch = torch.zeros_like(pitch) + return quat_from_euler_xyz(roll, pitch, yaw) + + +@torch.jit.script +def yaw_quat(quat: torch.Tensor) -> torch.Tensor: + """Extract the yaw component of a quaternion. + + Args: + quat: The orientation in (w, x, y, z). Shape is (..., 4) + + Returns: + A quaternion with only yaw component. + """ + shape = quat.shape + quat_yaw = quat.view(-1, 4) + qw = quat_yaw[:, 0] + qx = quat_yaw[:, 1] + qy = quat_yaw[:, 2] + qz = quat_yaw[:, 3] + yaw = torch.atan2(2 * (qw * qz + qx * qy), 1 - 2 * (qy * qy + qz * qz)) + quat_yaw = torch.zeros_like(quat_yaw) + quat_yaw[:, 3] = torch.sin(yaw / 2) + quat_yaw[:, 0] = torch.cos(yaw / 2) + quat_yaw = normalize(quat_yaw) + return quat_yaw.view(shape) + + +@torch.jit.script +def wrap_to_pi(angles): + angles %= 2 * np.pi + angles -= 2 * np.pi * (angles > np.pi) + return angles + + +@torch.jit.script +def quat_conjugate(a: Tensor, w_last: bool) -> Tensor: + shape = a.shape + a = a.reshape(-1, 4) + if w_last: + return torch.cat((-a[:, :3], a[:, -1:]), dim=-1).view(shape) + else: + return torch.cat((a[:, 0:1], -a[:, 1:]), dim=-1).view(shape) + + + + +@torch.jit.script +def quat_rotate(q: Tensor, v: Tensor, w_last: bool) -> Tensor: + shape = q.shape + if w_last: + q_w = q[:, -1] + q_vec = q[:, :3] + else: + q_w = q[:, 0] + q_vec = q[:, 1:] + a = v * (2.0 * q_w**2 - 1.0).unsqueeze(-1) + b = torch.cross(q_vec, v, dim=-1) * q_w.unsqueeze(-1) * 2.0 + c = q_vec * torch.bmm(q_vec.view(shape[0], 1, 3), v.view(shape[0], 3, 1)).squeeze(-1) * 2.0 + return a + b + c + + +@torch.jit.script +def quat_rotate_inverse(q: Tensor, v: Tensor, w_last: bool) -> Tensor: + # Same as quat_rotate but with the cross-product term (b) negated, + # which is equivalent to rotating by the conjugate quaternion (inverse rotation). + shape = q.shape + if w_last: + q_w = q[:, -1] + q_vec = q[:, :3] + else: + q_w = q[:, 0] + q_vec = q[:, 1:] + a = v * (2.0 * q_w**2 - 1.0).unsqueeze(-1) + b = torch.cross(q_vec, v, dim=-1) * q_w.unsqueeze(-1) * 2.0 + c = q_vec * torch.bmm(q_vec.view(shape[0], 1, 3), v.view(shape[0], 3, 1)).squeeze(-1) * 2.0 + return a - b + c + + +@torch.jit.script +def quat_angle_axis(x: Tensor, w_last: bool) -> Tuple[Tensor, Tensor]: + """ + The (angle, axis) representation of the rotation. The axis is normalized to unit length. + The angle is guaranteed to be between [0, pi]. + """ + if w_last: + w = x[..., -1] + axis = x[..., :3] + else: + w = x[..., 0] + axis = x[..., 1:] + # cos(theta) = 2*w^2 - 1, derived from w = cos(theta/2) and double-angle formula + s = 2 * (w**2) - 1 + angle = s.clamp(-1, 1).arccos() # just to be safe + axis /= axis.norm(p=2, dim=-1, keepdim=True).clamp(min=1e-9) + return angle, axis + + +@torch.jit.script +def quat_from_angle_axis(angle: Tensor, axis: Tensor, w_last: bool) -> Tensor: + theta = (angle / 2).unsqueeze(-1) + xyz = normalize(axis) * theta.sin() + w = theta.cos() + if w_last: + return quat_unit(torch.cat([xyz, w], dim=-1)) + else: + return quat_unit(torch.cat([w, xyz], dim=-1)) + + +@torch.jit.script +def vec_to_heading(h_vec): + h_theta = torch.atan2(h_vec[..., 1], h_vec[..., 0]) + return h_theta + + +@torch.jit.script +def heading_to_quat(h_theta, w_last: bool): + axis = torch.zeros( + h_theta.shape + + [ + 3, + ], + device=h_theta.device, + ) + axis[..., 2] = 1 + heading_q = quat_from_angle_axis(h_theta, axis, w_last=w_last) + return heading_q + + +@torch.jit.script +def quat_axis(q: Tensor, axis: int, w_last: bool) -> Tensor: + basis_vec = torch.zeros(q.shape[0], 3, device=q.device) + basis_vec[:, axis] = 1 + return quat_rotate(q, basis_vec, w_last) + + +@torch.jit.script +def normalize_angle(x): + return torch.atan2(torch.sin(x), torch.cos(x)) + + +@torch.jit.script +def get_basis_vector(q: Tensor, v: Tensor, w_last: bool) -> Tensor: + return quat_rotate(q, v, w_last) + + +@torch.jit.script +def quat_to_angle_axis(q, w_last: bool): + # type: (Tensor, bool) -> Tuple[Tensor, Tensor] + # computes axis-angle representation from quaternion q + # q must be normalized + # ZL: could have issues. + min_theta = 1e-5 + if w_last: + qx, qy, qz, qw = 0, 1, 2, 3 + else: + qw, qx, qy, qz = 0, 1, 2, 3 + + sin_theta = torch.sqrt(1 - q[..., qw] * q[..., qw]) + angle = 2 * torch.acos(q[..., qw]) + angle = normalize_angle(angle) + sin_theta_expand = sin_theta.unsqueeze(-1) + axis = q[..., qx:qw] / sin_theta_expand + + mask = torch.abs(sin_theta) > min_theta + default_axis = torch.zeros_like(axis) + default_axis[..., -1] = 1 + + angle = torch.where(mask, angle, torch.zeros_like(angle)) + mask_expand = mask.unsqueeze(-1) + axis = torch.where(mask_expand, axis, default_axis) + return angle, axis + + +@torch.jit.script +def slerp(q0, q1, t): + # type: (Tensor, Tensor, Tensor) -> Tensor + cos_half_theta = torch.sum(q0 * q1, dim=-1) + + neg_mask = cos_half_theta < 0 + q1 = q1.clone() + + # Replace: q1[neg_mask] = -q1[neg_mask] + # With: torch.where for safer broadcasting + neg_mask_expanded = neg_mask.unsqueeze(-1).expand_as(q1) + q1 = torch.where(neg_mask_expanded, -q1, q1) + + cos_half_theta = torch.abs(cos_half_theta) + cos_half_theta = torch.unsqueeze(cos_half_theta, dim=-1) + + half_theta = torch.acos(cos_half_theta) + sin_half_theta = torch.sqrt(1.0 - cos_half_theta * cos_half_theta) + + ratioA = torch.sin((1 - t) * half_theta) / sin_half_theta + ratioB = torch.sin(t * half_theta) / sin_half_theta + + new_q = ratioA * q0 + ratioB * q1 + + new_q = torch.where(torch.abs(sin_half_theta) < 0.001, 0.5 * q0 + 0.5 * q1, new_q) + new_q = torch.where(torch.abs(cos_half_theta) >= 1, q0, new_q) + + return new_q + + +@torch.jit.script +def angle_axis_to_exp_map(angle, axis): + # type: (Tensor, Tensor) -> Tensor + # compute exponential map from axis-angle + angle_expand = angle.unsqueeze(-1) + exp_map = angle_expand * axis + return exp_map + + +@torch.jit.script +def my_quat_rotate(q, v, w_last=True): + # type: (Tensor, Tensor, bool) -> Tensor + shape = q.shape + if w_last: + q_w = q[:, -1] + q_vec = q[:, :3] + else: + q_w = q[:, 0] + q_vec = q[:, 1:] + a = v * (2.0 * q_w**2 - 1.0).unsqueeze(-1) + b = torch.cross(q_vec, v, dim=-1) * q_w.unsqueeze(-1) * 2.0 + c = q_vec * torch.bmm(q_vec.view(shape[0], 1, 3), v.view(shape[0], 3, 1)).squeeze(-1) * 2.0 + return a + b + c + + +@torch.jit.script +def quat_to_tan_norm(q, w_last): + # type: (Tensor, bool) -> Tensor + # represents a rotation using the tangent and normal vectors + ref_tan = torch.zeros_like(q[..., 0:3]) + ref_tan[..., 0] = 1 + if w_last: + tan = my_quat_rotate(q, ref_tan) + else: + raise NotImplementedError + + ref_norm = torch.zeros_like(q[..., 0:3]) + ref_norm[..., -1] = 1 + if w_last: + norm = my_quat_rotate(q, ref_norm) + else: + raise NotImplementedError + + norm_tan = torch.cat([tan, norm], dim=len(tan.shape) - 1) + return norm_tan + + +@torch.jit.script +def calc_heading(q, w_last=True): + # type: (Tensor, bool) -> Tensor + # calculate heading direction from quaternion + # the heading is the direction on the xy plane + # q must be normalized + # this is the x axis heading + ref_dir = torch.zeros_like(q[..., 0:3]) + ref_dir[..., 0] = 1 + rot_dir = my_quat_rotate(q, ref_dir, w_last) + + heading = torch.atan2(rot_dir[..., 1], rot_dir[..., 0]) + return heading + + +@torch.jit.script +def quat_to_exp_map(q, w_last): + # type: (Tensor, bool) -> Tensor + # compute exponential map from quaternion + # q must be normalized + angle, axis = quat_to_angle_axis(q, w_last) + exp_map = angle_axis_to_exp_map(angle, axis) + return exp_map + + +@torch.jit.script +def calc_heading_quat(q, w_last): + # type: (Tensor, bool) -> Tensor + # calculate heading rotation from quaternion + # the heading is the direction on the xy plane + # q must be normalized + heading = calc_heading(q, w_last) + axis = torch.zeros_like(q[..., 0:3]) + axis[..., 2] = 1 + + heading_q = quat_from_angle_axis(heading, axis, w_last=w_last) + return heading_q + + +@torch.jit.script +def calc_heading_quat_inv(q, w_last): + # type: (Tensor, bool) -> Tensor + # calculate heading rotation from quaternion + # the heading is the direction on the xy plane + # q must be normalized + heading = calc_heading(q, w_last) + axis = torch.zeros_like(q[..., 0:3]) + axis[..., 2] = 1 + + heading_q = quat_from_angle_axis(-heading, axis, w_last=w_last) + return heading_q + + +@torch.jit.script +def quat_inverse(x, w_last): + # type: (Tensor, bool) -> Tensor + """ + The inverse of the rotation + """ + return quat_conjugate(x, w_last=w_last) + + +@torch.jit.script +def get_euler_xyz(q: Tensor, w_last: bool) -> Tuple[Tensor, Tensor, Tensor]: + if w_last: + qx, qy, qz, qw = 0, 1, 2, 3 + else: + qw, qx, qy, qz = 0, 1, 2, 3 + # roll (x-axis rotation) + sinr_cosp = 2.0 * (q[:, qw] * q[:, qx] + q[:, qy] * q[:, qz]) + cosr_cosp = ( + q[:, qw] * q[:, qw] - q[:, qx] * q[:, qx] - q[:, qy] * q[:, qy] + q[:, qz] * q[:, qz] + ) + roll = torch.atan2(sinr_cosp, cosr_cosp) + + # pitch (y-axis rotation) + sinp = 2.0 * (q[:, qw] * q[:, qy] - q[:, qz] * q[:, qx]) + pitch = torch.where(torch.abs(sinp) >= 1, copysign(np.pi / 2.0, sinp), torch.asin(sinp)) + + # yaw (z-axis rotation) + siny_cosp = 2.0 * (q[:, qw] * q[:, qz] + q[:, qx] * q[:, qy]) + cosy_cosp = ( + q[:, qw] * q[:, qw] + q[:, qx] * q[:, qx] - q[:, qy] * q[:, qy] - q[:, qz] * q[:, qz] + ) + yaw = torch.atan2(siny_cosp, cosy_cosp) + + return roll % (2 * np.pi), pitch % (2 * np.pi), yaw % (2 * np.pi) + + +# @torch.jit.script +def get_euler_xyz_in_tensor(q): + qx, qy, qz, qw = 0, 1, 2, 3 + # roll (x-axis rotation) + sinr_cosp = 2.0 * (q[:, qw] * q[:, qx] + q[:, qy] * q[:, qz]) + cosr_cosp = ( + q[:, qw] * q[:, qw] - q[:, qx] * q[:, qx] - q[:, qy] * q[:, qy] + q[:, qz] * q[:, qz] + ) + roll = torch.atan2(sinr_cosp, cosr_cosp) + + # pitch (y-axis rotation) + sinp = 2.0 * (q[:, qw] * q[:, qy] - q[:, qz] * q[:, qx]) + pitch = torch.where(torch.abs(sinp) >= 1, copysign(np.pi / 2.0, sinp), torch.asin(sinp)) + + # yaw (z-axis rotation) + siny_cosp = 2.0 * (q[:, qw] * q[:, qz] + q[:, qx] * q[:, qy]) + cosy_cosp = ( + q[:, qw] * q[:, qw] + q[:, qx] * q[:, qx] - q[:, qy] * q[:, qy] - q[:, qz] * q[:, qz] + ) + yaw = torch.atan2(siny_cosp, cosy_cosp) + + return torch.stack((roll, pitch, yaw), dim=-1) + + +@torch.jit.script +def quat_pos(x): + """ + make all the real part of the quaternion positive + """ + q = x + z = (q[..., 3:] < 0).float() + q = (1 - 2 * z) * q + return q + + +@torch.jit.script +def is_valid_quat(q): + x, y, z, w = q[..., 0], q[..., 1], q[..., 2], q[..., 3] + return (w * w + x * x + y * y + z * z).allclose(torch.ones_like(w)) + + +@torch.jit.script +def quat_normalize(q): + """ + Construct 3D rotation from quaternion (the quaternion needs not to be normalized). + """ + q = quat_unit(quat_pos(q)) # normalized to positive and unit quaternion + return q + + +@torch.jit.script +def quat_mul(a, b, w_last: bool): + assert a.shape == b.shape + shape = a.shape + a = a.reshape(-1, 4) + b = b.reshape(-1, 4) + + if w_last: + x1, y1, z1, w1 = a[..., 0], a[..., 1], a[..., 2], a[..., 3] + x2, y2, z2, w2 = b[..., 0], b[..., 1], b[..., 2], b[..., 3] + else: + w1, x1, y1, z1 = a[..., 0], a[..., 1], a[..., 2], a[..., 3] + w2, x2, y2, z2 = b[..., 0], b[..., 1], b[..., 2], b[..., 3] + ww = (z1 + x1) * (x2 + y2) + yy = (w1 - y1) * (w2 + z2) + zz = (w1 + y1) * (w2 - z2) + xx = ww + yy + zz + qq = 0.5 * (xx + (z1 - x1) * (x2 - y2)) + w = qq - ww + (z1 - y1) * (y2 - z2) + x = qq - xx + (x1 + w1) * (x2 + w2) + y = qq - yy + (w1 - x1) * (y2 + z2) + z = qq - zz + (z1 + y1) * (w2 - x2) + + if w_last: + quat = torch.stack([x, y, z, w], dim=-1).view(shape) + else: + quat = torch.stack([w, x, y, z], dim=-1).view(shape) + + return quat + + + +@torch.jit.script +def quat_mul_norm(x, y, w_last): + # type: (Tensor, Tensor, bool) -> Tensor + """ + Combine two set of 3D rotations together using \**\* operator. The shape needs to be + broadcastable + """ + return quat_unit(quat_mul(x, y, w_last)) + + +@torch.jit.script +def quat_identity(shape: List[int]): + """ + Construct 3D identity rotation given shape + """ + w = torch.ones(shape + [1]) + xyz = torch.zeros(shape + [3]) + q = torch.cat([xyz, w], dim=-1) + return quat_normalize(q) + + +@torch.jit.script +def quat_identity_like(x): + """ + Construct identity 3D rotation with the same shape + """ + return quat_identity(list(x.shape[:-1])) + + +@torch.jit.script +def transform_from_rotation_translation( + r: Optional[torch.Tensor] = None, t: Optional[torch.Tensor] = None +): + """ + Construct a transform from a quaternion and 3D translation. Only one of them can be None. + """ + assert r is not None or t is not None, "rotation and translation can't be all None" + if r is None: + assert t is not None + r = quat_identity(list(t.shape)) + if t is None: + t = torch.zeros(list(r.shape) + [3]) + return torch.cat([r, t], dim=-1) + + +@torch.jit.script +def transform_rotation(x): + """Get rotation from transform""" + return x[..., :4] + + +@torch.jit.script +def transform_translation(x): + """Get translation from transform""" + return x[..., 4:] + + +@torch.jit.script +def transform_mul(x, y): + """ + Combine two transformation together + """ + z = transform_from_rotation_translation( + r=quat_mul_norm(transform_rotation(x), transform_rotation(y), w_last=True), + t=quat_rotate(transform_rotation(x), transform_translation(y), w_last=True) + + transform_translation(x), + ) + return z + + +##################################### FROM PHC rotation_conversions.py ##################################### +@torch.jit.script +def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as quaternions to rotation matrices. + + Args: + quaternions: quaternions with real part first, + as tensor of shape (..., 4). + + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + r, i, j, k = torch.unbind(quaternions, -1) + two_s = 2.0 / (quaternions * quaternions).sum(-1) + + o = torch.stack( + ( + 1 - two_s * (j * j + k * k), + two_s * (i * j - k * r), + two_s * (i * k + j * r), + two_s * (i * j + k * r), + 1 - two_s * (i * i + k * k), + two_s * (j * k - i * r), + two_s * (i * k - j * r), + two_s * (j * k + i * r), + 1 - two_s * (i * i + j * j), + ), + -1, + ) + return o.reshape(quaternions.shape[:-1] + (3, 3)) + + +@torch.jit.script +def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as axis/angle to quaternions. + + Args: + axis_angle: Rotations given as a vector in axis angle form, + as a tensor of shape (..., 3), where the magnitude is + the angle turned anticlockwise in radians around the + vector's direction. + + Returns: + quaternions with real part first, as tensor of shape (..., 4). + """ + angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True) + half_angles = angles * 0.5 + eps = 1e-6 + small_angles = angles.abs() < eps + sin_half_angles_over_angles = torch.empty_like(angles) + sin_half_angles_over_angles[~small_angles] = ( + torch.sin(half_angles[~small_angles]) / angles[~small_angles] + ) + # for x small, sin(x/2) is about x/2 - (x/2)^3/6 + # so sin(x/2)/x is about 1/2 - (x*x)/48 + sin_half_angles_over_angles[small_angles] = ( + 0.5 - (angles[small_angles] * angles[small_angles]) / 48 + ) + quaternions = torch.cat( + [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1 + ) + return quaternions + + +# @torch.jit.script +def wxyz_to_xyzw(quat): + return quat[..., [1, 2, 3, 0]] + + +# @torch.jit.script +def xyzw_to_wxyz(quat): + return quat[..., [3, 0, 1, 2]] + + +def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor: + """ + w x y z + Convert rotations given as rotation matrices to quaternions. + + Args: + matrix: Rotation matrices as tensor of shape (..., 3, 3). + + Returns: + quaternions with real part first, as tensor of shape (..., 4). + """ + if matrix.size(-1) != 3 or matrix.size(-2) != 3: + raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.") + + batch_dim = matrix.shape[:-2] + m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind( + matrix.reshape(batch_dim + (9,)), dim=-1 + ) + + q_abs = _sqrt_positive_part( + torch.stack( + [ + 1.0 + m00 + m11 + m22, + 1.0 + m00 - m11 - m22, + 1.0 - m00 + m11 - m22, + 1.0 - m00 - m11 + m22, + ], + dim=-1, + ) + ) + + # we produce the desired quaternion multiplied by each of r, i, j, k + quat_by_rijk = torch.stack( + [ + torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1), + torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1), + torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1), + torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1), + ], + dim=-2, + ) + + # We floor here at 0.1 but the exact level is not important; if q_abs is small, + # the candidate won't be picked. + flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device) + quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr)) + + # if not for numerical problems, quat_candidates[i] should be same (up to a sign), + # forall i; we pick the best-conditioned one (with the largest denominator) + + return quat_candidates[ + F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, : # pyre-ignore[16] + ].reshape(batch_dim + (4,)) + + +def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor: + """ + Returns torch.sqrt(torch.max(0, x)) + but with a zero subgradient where x is 0. + """ + ret = torch.zeros_like(x) + positive_mask = x > 0 + ret[positive_mask] = torch.sqrt(x[positive_mask]) + return ret + + +def quat_w_first(rot): + rot = torch.cat([rot[..., [-1]], rot[..., :-1]], -1) + return rot + + +@torch.jit.script +def quat_from_euler_xyz(roll, pitch, yaw): + cy = torch.cos(yaw * 0.5) + sy = torch.sin(yaw * 0.5) + cr = torch.cos(roll * 0.5) + sr = torch.sin(roll * 0.5) + cp = torch.cos(pitch * 0.5) + sp = torch.sin(pitch * 0.5) + + qw = cy * cr * cp + sy * sr * sp + qx = cy * sr * cp - sy * cr * sp + qy = cy * cr * sp + sy * sr * cp + qz = sy * cr * cp - cy * sr * sp + + return torch.stack([qx, qy, qz, qw], dim=-1) + + + +@torch.jit.script +def remove_smpl_base_rot(quat, w_last: bool): + # [0.5,0.5,0.5,0.5] is a 120° rotation about the [1,1,1] axis — SMPL's default rest orientation. + # Conjugating it out aligns with a neutral standing pose. + base_rot = quat_conjugate(torch.tensor([[0.5, 0.5, 0.5, 0.5]]).to(quat), w_last=w_last) # SMPL + return quat_mul(quat, base_rot.repeat(quat.shape[0], 1), w_last=w_last) + + +@torch.jit.script +def smpl_root_ytoz_up(root_quat_y_up) -> torch.Tensor: + """Convert SMPL root quaternion from Y-up to Z-up coordinate system""" + # 90° rotation about X-axis maps Y-up (SMPL convention) to Z-up (robot convention) + base_rot = angle_axis_to_quaternion(torch.tensor([[np.pi / 2, 0.0, 0.0]]).to(root_quat_y_up)) + root_quat_z_up = quat_mul( + base_rot.repeat(root_quat_y_up.shape[0], 1), root_quat_y_up, w_last=False + ) + return root_quat_z_up + + +@torch.jit.script +def rotate_vectors_by_quaternion(quat: torch.Tensor, vec: torch.Tensor) -> torch.Tensor: + """ + Rotate `vec` by `quat`, elementwise. + + Args: + quat (torch.Tensor): Tensor of shape (..., 4), quaternions in [x, y, z, w] format. + vec (torch.Tensor): Tensor of shape (..., 3), vectors to rotate. + + Returns: + torch.Tensor: Rotated vectors, same shape as `vec`. + """ + q_xyz = quat[..., :3] # (..., 3) + q_w = quat[..., 3:].unsqueeze(-1) # (..., 1, 1) -> we'll squeeze to (...,1) + + # Compute intermediate cross products + # t = 2 * q_xyz × v + t = 2.0 * torch.cross(q_xyz, vec, dim=-1) # (..., 3) + + # v' = v + w * t + q_xyz × t + rotated = vec + q_w.squeeze(-1) * t + torch.cross(q_xyz, t, dim=-1) + return rotated + + +def rot6d_to_quat_first_two_cols(rot_6d: torch.Tensor) -> torch.Tensor: + """ + Convert 6D rotation representation (first 2 columns of rotation matrix) to quaternion. + + This function handles the 6D representation where the first 6 elements represent + the first 2 columns of a 3x3 rotation matrix (flattened). The third column is + reconstructed via cross product of the first two columns. + + Args: + rot_6d (torch.Tensor): Tensor of shape (..., 6) representing the first 2 columns + of rotation matrix flattened. + + Returns: + torch.Tensor: Quaternion in (w, x, y, z) format, shape (..., 4). + """ + # Reshape to get first 2 columns: (..., 3, 2) + rot_2cols = rot_6d.reshape(*rot_6d.shape[:-1], 3, 2) + + # Extract the two column vectors + col_0 = rot_2cols[..., :, 0] + col_1 = rot_2cols[..., :, 1] + + # Normalize the columns to ensure they are unit vectors + col_0 = F.normalize(col_0, dim=-1) + col_1 = F.normalize(col_1, dim=-1) + + # Reconstruct the third column via cross product + col_2 = torch.cross(col_0, col_1, dim=-1) + + # Stack to form full rotation matrix (..., 3, 3) + rot_matrix = torch.stack([col_0, col_1, col_2], dim=-1) + + # Convert rotation matrix to quaternion (w, x, y, z format) + quat = matrix_to_quaternion(rot_matrix) + + return quat + + +def remove_bvh_base_rot(quat, w_last: bool): + """Remove BVH base rotation. BVH base = conj(SMPL base), so conj(BVH base) = SMPL base.""" + base_rot = torch.tensor([[0.5, 0.5, 0.5, 0.5]]).to(quat) + return quat_mul(quat, base_rot.repeat(quat.shape[0], 1), w_last=w_last) diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/__init__.py b/GR00T-WholeBodyControl/gear_sonic/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/__init__.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..65de7fcc9307aa0fcefd5bcd7f9fedba7366c5aa Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/__init__.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/__init__.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2ed0a2f6e4a0b0dc676cbe92044976392954185d Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/__init__.cpython-311.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/__init__.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d6c88a9185dfb7a6523da4e3d4c73246d45e443e Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/average_meters.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/average_meters.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b6e31f25e566f716561a1245fdb258789ab0618 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/average_meters.cpython-311.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/average_meters.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/average_meters.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7190084c8cb7a330f17a50047b2c86a6e5fa671d Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/average_meters.cpython-312.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/batch_normalizer.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/batch_normalizer.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..06fcf453ac7398a8cb2e3d96e2f2b07758c1c423 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/batch_normalizer.cpython-311.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/batch_normalizer.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/batch_normalizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1a22a78d9c2835c6182f864c18f244b8af48d814 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/batch_normalizer.cpython-312.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/common.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/common.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..03c1b5f526e3cc0daf315f78d2fba530d7917a68 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/common.cpython-311.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/common.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8f1e04b51e6a7de0f08c1fa471682ec8f03d6655 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/common.cpython-312.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/config_utils.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/config_utils.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f162bdb1affc1b58b05d9cf9dbb5ad9d272fb020 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/config_utils.cpython-311.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/config_utils.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/config_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..37d5195f134fba15e2f9338dc6978ab95bf0dc93 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/config_utils.cpython-312.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/logging.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/logging.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e89f4f6dc994d7d03d135b9eb6d1237003b94068 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/logging.cpython-311.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/logging.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/logging.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5bfe469f4fd62771e0349d2ca7e0b47fcf33e8df Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/logging.cpython-312.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/obs_utils.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/obs_utils.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bc3c573c502cd8e263ad08f903e6b9b9b8acccd0 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/obs_utils.cpython-311.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/obs_utils.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/obs_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2acfdf2292cae67a89847dc09e0ea34e24105e88 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/obs_utils.cpython-312.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/running_mean_std.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/running_mean_std.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7b63d02b63e4e9b5f440665d6c4655d7f883ee4c Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/running_mean_std.cpython-311.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/running_mean_std.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/running_mean_std.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a9e8c07541deb14011ac1a60dd10560296489f39 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/__pycache__/running_mean_std.cpython-312.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/average_meters.py b/GR00T-WholeBodyControl/gear_sonic/utils/average_meters.py new file mode 100644 index 0000000000000000000000000000000000000000..2d8679063a8d442b9ab49566fa02a5c173f5354f --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/average_meters.py @@ -0,0 +1,177 @@ +"""Running-average utility classes for tracking scalar and tensor metrics. + +Three classes are provided at increasing levels of generality: + +* ``AverageMeter`` – fixed-capacity windowed mean for batched tensors, + implemented as an ``nn.Module`` with a registered buffer so it moves with + the model between devices. +* ``TensorAverageMeter`` – simple accumulate-then-mean helper for a single + metric. +* ``TensorAverageMeterDict`` – dictionary wrapper around + ``TensorAverageMeter`` for tracking multiple named metrics simultaneously. +""" + +import numpy as np +import torch +import torch.nn as nn + + +class AverageMeter(nn.Module): + """Windowed running mean with a configurable sample capacity. + + Maintains a weighted mean over the last ``max_size`` samples. Implemented + as an ``nn.Module`` so ``self.mean`` is a registered buffer and follows + ``.to(device)`` / ``.cuda()`` calls automatically. + + Args: + in_shape: Shape of the per-sample value tensor (passed to + ``torch.zeros``). + max_size: Maximum effective sample count used when computing the + running mean. Older samples are down-weighted once this limit + is reached. + """ + + def __init__(self, in_shape, max_size): + super().__init__() + self.max_size = max_size + self.current_size = 0 + self.register_buffer("mean", torch.zeros(in_shape, dtype=torch.float32)) + + def update(self, values): + """Incorporate a new batch of values into the running mean. + + Args: + values: Tensor of shape ``(batch, *in_shape)``. Empty batches + are silently skipped. + """ + size = values.size()[0] + if size == 0: + return + new_mean = torch.mean(values.float(), dim=0) + size = np.clip(size, 0, self.max_size) + old_size = min(self.max_size - size, self.current_size) + size_sum = old_size + size + self.current_size = size_sum + self.mean = (self.mean * old_size + new_mean * size) / size_sum + + def clear(self): + """Reset the meter to its initial empty state.""" + self.current_size = 0 + self.mean.fill_(0) + + def __len__(self): + """Return the current effective sample count.""" + return self.current_size + + def get_mean(self): + """Return the current mean as a NumPy array on CPU. + + Returns: + np.ndarray of shape ``in_shape`` with the leading size-1 dimension + squeezed out. + """ + return self.mean.squeeze(0).cpu().numpy() + + +class TensorAverageMeter: + """Accumulate tensors and compute their mean on demand. + + Unlike ``AverageMeter`` this class keeps all accumulated tensors in a list + and concatenates them lazily when ``mean()`` is called. It is best suited + for metrics that are computed once per rollout step and averaged at the + end of an epoch. + """ + + def __init__(self): + self.tensors = [] + + def add(self, x): + """Append a tensor to the accumulator. + + Scalar tensors (0-D) are automatically unsqueezed to 1-D before + appending so that concatenation works correctly. + + Args: + x: Tensor to accumulate. + """ + if len(x.shape) == 0: + x = x.unsqueeze(0) + self.tensors.append(x) + + def mean(self): + """Return the mean of all accumulated tensors. + + Returns: + A scalar tensor if tensors have been added, or the integer ``0`` + when the accumulator is empty or contains no elements. + """ + if len(self.tensors) == 0: + return 0 + cat = torch.cat(self.tensors, dim=0) + if cat.numel() == 0: + return 0 + else: + return cat.mean() + + def clear(self): + """Discard all accumulated tensors.""" + self.tensors = [] + + def mean_and_clear(self): + """Compute the mean, clear the accumulator, and return the mean. + + Returns: + Same as ``mean()``. + """ + mean = self.mean() + self.clear() + return mean + + +class TensorAverageMeterDict: + """Dictionary of ``TensorAverageMeter`` objects for multi-metric tracking. + + Accepts batches of ``{key: tensor}`` dicts and lazily creates a + ``TensorAverageMeter`` per key. Uses plain ``dict`` internally (not + ``defaultdict``) to avoid lambda pickling issues with DDP. + """ + + def __init__(self): + self.data = {} + + def add(self, data_dict): + """Append a batch of named metric tensors to their respective meters. + + Args: + data_dict: Mapping from metric name to tensor value. New keys are + registered automatically. + """ + for k, v in data_dict.items(): + # Originally used a defaultdict, this had lambda + # pickling issues with DDP. + if k not in self.data: + self.data[k] = TensorAverageMeter() + self.data[k].add(v) + + def mean(self): + """Return a dict mapping each key to its accumulated mean. + + Returns: + Dict[str, scalar tensor | int] with the same keys as were added. + """ + mean_dict = {k: v.mean() for k, v in self.data.items()} + return mean_dict + + def clear(self): + """Discard all meters and their accumulated data.""" + self.data = {} + + def mean_and_clear(self): + """Compute means for all keys, clear the meters, and return the means. + + Returns: + Same as ``mean()``. + """ + mean = self.mean() + self.clear() + return mean diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/batch_normalizer.py b/GR00T-WholeBodyControl/gear_sonic/utils/batch_normalizer.py new file mode 100644 index 0000000000000000000000000000000000000000..1c9895d10d94d96e9cfc611b85dc2a793581b9a3 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/batch_normalizer.py @@ -0,0 +1,48 @@ +import torch +import torch.nn as nn + + +class BatchNormNormalizer(nn.Module): + def __init__(self, insize, epsilon=1e-05, per_channel=False, norm_only=False): + super().__init__() + assert len(insize) == 1, "BatchNormNormalizer only supports 1D observation spaces" + self._normalizer = nn.SyncBatchNorm(num_features=insize[0], affine=False) + + @property + def num_features(self): + return self._normalizer.num_features + + def forward(self, x: torch.Tensor) -> torch.Tensor: + input_shape = x.shape + if len(x.shape) == 3: + x = x.reshape(-1, x.shape[-1]) + + x = self._normalizer(x) + x = x.view(input_shape) + return x + + def update(self, input: torch.Tensor): + """Update running stats from input. No-op in eval mode. + + Calls SyncBatchNorm.forward() for its side effect of updating + running_mean/running_var (and multi-GPU sync). Output is discarded. + """ + if not self.training: # do nothing if in evaluation mode + return + if len(input.shape) == 3: + input = input.reshape(-1, input.shape[-1]) + with torch.no_grad(): + self._normalizer(input) + + def normalize(self, input: torch.Tensor) -> torch.Tensor: + """Normalize using current running stats without updating them.""" + input_shape = input.shape + if len(input.shape) == 3: + input = input.reshape(-1, input_shape[-1]) + y = (input - self._normalizer.running_mean) / torch.sqrt( + self._normalizer.running_var + self._normalizer.eps + ) + y = torch.clamp(y, min=-5.0, max=5.0) + if len(input_shape) == 3: + y = y.view(input_shape) + return y diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/common.py b/GR00T-WholeBodyControl/gear_sonic/utils/common.py new file mode 100644 index 0000000000000000000000000000000000000000..fab53b9b86f69cc1b0bd3448795d10ac3f63ac50 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/common.py @@ -0,0 +1,224 @@ +"""General utility functions for RL training scripts. + +Provides argument-conflict resolution for Hydra/CLI arg lists, colour-coded +console print helpers, timestamp generation, model-args loading, random-seed +initialisation, and a scalar-to-RGB colour mapper. +""" + +# Copyright (c) 2018-2022, NVIDIA Corporation +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import argparse +from datetime import datetime +import os +import random +import sys + +import numpy as np +import torch + + +def solve_argv_conflict(args_list): + """Remove entries from ``args_list`` that are overridden by ``sys.argv``. + + When programmatic defaults conflict with user-provided command-line + arguments, this function removes the duplicates from ``args_list`` so + that the command-line value wins. + + Args: + args_list: Mutable list of argument strings (modified in-place). + Entries that also appear in ``sys.argv[1:]`` (along with their + positional values) are removed. + """ + arguments_to_be_removed = [] + arguments_size = [] + + for argv in sys.argv[1:]: + if argv.startswith("-"): + size_count = 1 + for i, args in enumerate(args_list): + if args == argv: + arguments_to_be_removed.append(args) + for more_args in args_list[i + 1 :]: + if not more_args.startswith("-"): + size_count += 1 + else: + break + arguments_size.append(size_count) + break + + for args, size in zip(arguments_to_be_removed, arguments_size): + args_index = args_list.index(args) + for _ in range(size): + args_list.pop(args_index) + + +def print_error(*message): + """Print an error message in red and raise ``RuntimeError``. + + Args: + *message: Message fragments passed to ``print``. + + Raises: + RuntimeError: Always raised after printing. + """ + print("\033[91m", "ERROR ", *message, "\033[0m") + raise RuntimeError + + +def print_ok(*message): + """Print a success message in green. + + Args: + *message: Message fragments passed to ``print``. + """ + print("\033[92m", *message, "\033[0m") + + +def print_warning(*message): + """Print a warning message in yellow. + + Args: + *message: Message fragments passed to ``print``. + """ + print("\033[93m", *message, "\033[0m") + + +def print_info(*message): + """Print an informational message in cyan. + + Args: + *message: Message fragments passed to ``print``. + """ + print("\033[96m", *message, "\033[0m") + + +def get_time_stamp(): + """Return the current date-time as a formatted string. + + Returns: + String of the form ``"MM-DD-YYYY-HH-MM-SS"``. + """ + now = datetime.now() + year = now.strftime("%Y") + month = now.strftime("%m") + day = now.strftime("%d") + hour = now.strftime("%H") + minute = now.strftime("%M") + second = now.strftime("%S") + return f"{month}-{day}-{year}-{hour}-{minute}-{second}" + + +def parse_model_args(model_args_path): + """Load model arguments from a Python-literal file as an ``argparse.Namespace``. + + The file is expected to contain a single Python dict literal that is + ``eval``-ed and wrapped in ``argparse.Namespace``. + + Args: + model_args_path: Path to the model-args file. + + Returns: + ``argparse.Namespace`` with one attribute per dict key. + """ + fp = open(model_args_path) + model_args = eval(fp.read()) + model_args = argparse.Namespace(**model_args) + + return model_args + + +def seeding(seed=0, torch_deterministic=False): + """Set global random seeds for reproducibility. + + Seeds ``random``, ``numpy``, ``torch`` (CPU and all GPUs), and the + ``PYTHONHASHSEED`` environment variable. Optionally enables cuDNN + deterministic mode at the cost of performance. + + Args: + seed: Integer seed value. + torch_deterministic: If True, enables fully deterministic CUDA + operations (sets ``CUBLAS_WORKSPACE_CONFIG``, disables cuDNN + benchmarking, and calls ``torch.use_deterministic_algorithms(True)``). + + Returns: + The seed value that was applied. + """ + print(f"Setting seed: {seed}") + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + os.environ["PYTHONHASHSEED"] = str(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + if torch_deterministic: + # refer to https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + torch.use_deterministic_algorithms(True) + else: + torch.backends.cudnn.benchmark = True + torch.backends.cudnn.deterministic = False + + return seed + + +def distance_l2(root_pos, wp_pos): + """Compute the L2 distance between two position tensors. + + Args: + root_pos: Reference position tensor. + wp_pos: Waypoint position tensor of the same shape as ``root_pos``. + + Returns: + Scalar tensor with the Euclidean distance. + """ + return torch.norm(wp_pos - root_pos, dim=0) + + +def value_to_color(value, min_value, max_value): + """ + Converts a numerical value to an RGB color. + The color will range from blue (low values) to red (high values). + """ + # Ensure value is within the range [0, max_value] + value = max(min_value, min(value, max_value)) + + # Calculate the proportion of the value + red = (value - min_value) / (max_value - min_value) + + # Map the proportion to the red channel for a red gradient + # Blue for minimum value and red for maximum value + blue = 1 - red + green = 0 # Keep green constant for simplicity + + # Return the RGB color + return red, green, blue diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/config_utils.py b/GR00T-WholeBodyControl/gear_sonic/utils/config_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a91687619e57d81509c94598dece1e58a7ae7c42 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/config_utils.py @@ -0,0 +1,20 @@ +import math + +from loguru import logger +from omegaconf import OmegaConf + + +def register_rl_resolvers(): + + try: + OmegaConf.register_new_resolver("eval", eval) + OmegaConf.register_new_resolver("if", lambda pred, a, b: a if pred else b) + OmegaConf.register_new_resolver("eq", lambda x, y: x.lower() == y.lower()) + OmegaConf.register_new_resolver("sqrt", lambda x: math.sqrt(float(x))) + OmegaConf.register_new_resolver("sum", lambda x: sum(x)) + OmegaConf.register_new_resolver("ceil", lambda x: math.ceil(x)) + OmegaConf.register_new_resolver("int", lambda x: int(x)) + OmegaConf.register_new_resolver("len", lambda x: len(x)) + OmegaConf.register_new_resolver("sum_list", lambda lst: sum(lst)) + except Exception as e: + logger.warning(f"Warning: Some resolvers already registered: {e}") diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/__init__.py b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/episode_state.py b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/episode_state.py new file mode 100644 index 0000000000000000000000000000000000000000..fafbd91d3c45aa9df10d1491b3900dd4f513d6d5 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/episode_state.py @@ -0,0 +1,32 @@ +class EpisodeState: + """Episode state controller for data collection. + + Manages the state transitions for episode recording: + - IDLE: Not recording + - RECORDING: Currently recording data + - NEED_TO_SAVE: Recording stopped, waiting to save + """ + + def __init__(self): + self.RECORDING = "recording" + self.IDLE = "idle" + self.NEED_TO_SAVE = "need_to_save" + + self.state = self.IDLE + + def change_state(self): + """Cycle through states: IDLE -> RECORDING -> NEED_TO_SAVE -> IDLE""" + if self.state == self.IDLE: + self.state = self.RECORDING + elif self.state == self.RECORDING: + self.state = self.NEED_TO_SAVE + elif self.state == self.NEED_TO_SAVE: + self.state = self.IDLE + + def reset_state(self): + """Reset to IDLE state.""" + self.state = self.IDLE + + def get_state(self): + """Get current state.""" + return self.state diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/keyboard_subscriber.py b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/keyboard_subscriber.py new file mode 100644 index 0000000000000000000000000000000000000000..73054667c942e78734be2cdf17f895315684c6cf --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/keyboard_subscriber.py @@ -0,0 +1,39 @@ +"""ZMQ-based keyboard subscriber for data collection control.""" + +import zmq + +DEFAULT_ZMQ_KEYBOARD_PORT = 5580 + + +class ZMQKeyboardSubscriber: + """Receives keyboard events from ZMQ SUB socket (non-blocking).""" + + def __init__(self, port: int = DEFAULT_ZMQ_KEYBOARD_PORT, host: str = "localhost"): + self._ctx = zmq.Context() + self._socket = self._ctx.socket(zmq.SUB) + self._socket.setsockopt_string(zmq.SUBSCRIBE, "") + self._socket.setsockopt(zmq.CONFLATE, 1) + self._socket.setsockopt(zmq.RCVTIMEO, 0) + self._socket.connect(f"tcp://{host}:{port}") + self._data = None + print(f"[ZMQKeyboardSubscriber] Connected to tcp://{host}:{port}") + + def read_msg(self): + """Return the latest key press (or None).""" + try: + self._data = self._socket.recv_string(zmq.NOBLOCK) + except zmq.Again: + pass + data = self._data + self._data = None + return data + + def close(self): + self._socket.close() + self._ctx.term() + + def __del__(self): + try: + self.close() + except Exception: + pass diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/telemetry.py b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/telemetry.py new file mode 100644 index 0000000000000000000000000000000000000000..c9086f4f98bd1f819ac508cf4866820337d1805b --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/telemetry.py @@ -0,0 +1,109 @@ +"""Timing and performance monitoring for data collection loops.""" + +from collections import deque +import time + + +class Telemetry: + """Handles timing and performance monitoring for different code sections.""" + + def __init__(self, window_size: int = 100): + self.window_size = window_size + self._current_timers = {} + self._last_results = {} + self._history = {} + + def start_timer(self, name: str): + self._current_timers[name] = time.perf_counter() + + def stop_timer(self, name: str) -> float: + if name not in self._current_timers: + print(f"Warning: Telemetry Timer '{name}' stopped without being started.") + return 0.0 + + start_time = self._current_timers.pop(name) + duration = time.perf_counter() - start_time + self.record_value(name, duration) + return duration + + class Timer: + """Context manager for timing operations.""" + + def __init__(self, telemetry, name: str): + self.telemetry = telemetry + self.name = name + + def __enter__(self): + self.telemetry.start_timer(self.name) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + del exc_type, exc_val, exc_tb + self.telemetry.stop_timer(self.name) + return False + + def timer(self, name: str): + """Returns a context manager for timing an operation.""" + return self.Timer(self, name) + + def record_value(self, name: str, value: float): + self._last_results[name] = value + if name not in self._history: + self._history[name] = deque(maxlen=self.window_size) + self._history[name].append(value) + + def get_last_timing(self) -> dict[str, float]: + return self._last_results.copy() + + def clear_last_timing(self): + self._last_results.clear() + + def get_average(self, name: str) -> float | None: + if name not in self._history or not self._history[name]: + return None + return sum(self._history[name]) / len(self._history[name]) + + def log_timing_info( + self, + context: str = "", + threshold: float = 0.001, + log_averages: bool = True, + ): + current_iteration_data = self.get_last_timing() + significant_timings = {k: v for k, v in current_iteration_data.items() if v > threshold} + + should_log = bool(significant_timings) + + avg_data = {} + if log_averages: + for name in self._history: + avg = self.get_average(name) + if avg is not None: + avg_data[f"{name}_avg"] = avg + if avg_data: + should_log = True + + if not should_log: + self.clear_last_timing() + return + + log_lines = [f"\n{context} Timing breakdown:" if context else "\nTiming breakdown:"] + + if significant_timings: + log_lines.append(f" Current Iteration (> {threshold*1000:.1f}ms):") + for name, duration in sorted(significant_timings.items()): + log_lines.append(f" {name}: {duration * 1000:.2f}ms") + elif current_iteration_data: + total_key = next((k for k in current_iteration_data if "total" in k), None) + if total_key: + log_lines.append( + f" Current Iteration Total: {current_iteration_data[total_key] * 1000:.2f}ms" + ) + + if log_averages and avg_data: + log_lines.append(f" Moving Averages (last {self.window_size} iters):") + for name, avg in sorted(avg_data.items()): + log_lines.append(f" {name}: {avg * 1000:.2f}ms") + + print("\n".join(log_lines)) + self.clear_last_timing() diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/text_to_speech.py b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/text_to_speech.py new file mode 100644 index 0000000000000000000000000000000000000000..03c3a43cc6bdbf644205beb268806d1f6a601fa5 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/text_to_speech.py @@ -0,0 +1,44 @@ +"""Optional text-to-speech feedback for data collection.""" + +import threading + + +class TextToSpeech: + def __init__(self, rate: int = 150, volume: float = 1.0): + try: + import pyttsx3 + + self.engine = pyttsx3.init(driverName="espeak") + self.engine.setProperty("rate", rate) + self.engine.setProperty("volume", volume) + except Exception as e: + print(f"[Text To Speech] Initialization failed: {e}") + self.engine = None + self._speech_thread: threading.Thread | None = None + self._lock = threading.Lock() + + def say(self, message: str, blocking: bool = False): + if self.engine: + if blocking: + self._say_blocking(message) + else: + thread = threading.Thread(target=self._say_blocking, args=(message,), daemon=True) + thread.start() + self._speech_thread = thread + + def _say_blocking(self, message: str): + with self._lock: + try: + self.engine.say(message) + self.engine.runAndWait() + except RuntimeError: + pass + + def wait_for_completion(self): + if self._speech_thread and self._speech_thread.is_alive(): + self._speech_thread.join() + + def print_and_say(self, message: str, say: bool = True, blocking: bool = False): + print(message) + if say and self.engine is not None: + self.say(message, blocking=blocking) diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/transforms.py b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..17edaab008fb5469c1720ebaea17a0bd85feec86 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/transforms.py @@ -0,0 +1,94 @@ +"""Rotation and gravity transform utilities for data collection.""" + +import numpy as np +from scipy.spatial.transform import Rotation as R + + +def quat_to_rot6d(q): + """Convert scalar-first quaternion(s) (wxyz) to 6D rotation representation. + + The 6D representation consists of the first two columns of the rotation + matrix, flattened (Zhou et al., CVPR 2019). + + Accepted input shapes: + * ``(4,)`` -- single quaternion -> returns ``(6,)`` + * ``(N, 4)`` -- batch of quats -> returns ``(N, 6)`` + * ``(N*4,)`` -- flat concatenated -> returns ``(N*6,)`` + """ + q = np.asarray(q) + if q.ndim == 1 and q.shape[0] > 4: + assert q.shape[0] % 4 == 0, f"Flat quat length {q.shape[0]} is not divisible by 4" + q = q.reshape(-1, 4) + rot_6d = quat_to_rot6d(q) + return rot_6d.ravel() + + single = q.ndim == 1 + q = np.atleast_2d(q) + q_xyzw = q[:, [1, 2, 3, 0]] + rot_mat = R.from_quat(q_xyzw).as_matrix() # (N, 3, 3) + rot_6d = rot_mat[:, :, :2].transpose(0, 2, 1).reshape(-1, 6) # (N, 6) + if single: + return rot_6d[0].astype(q.dtype) + return rot_6d.astype(q.dtype) + + +def rot6d_to_quat(r): + """Convert 6D rotation representation to scalar-first quaternion(s) (wxyz). + + Accepted input shapes: + * ``(6,)`` -- single rot6d -> returns ``(4,)`` + * ``(N, 6)`` -- batch -> returns ``(N, 4)`` + * ``(N*6,)`` -- flat concat -> returns ``(N*4,)`` + (length must be divisible by 6) + + Args: + r: 6D rotation array (first two columns of rotation matrix, row-major). + + Returns: + Quaternion array in wxyz order. + """ + r = np.asarray(r, dtype=np.float64) + if r.ndim == 1 and r.shape[0] > 6: + assert r.shape[0] % 6 == 0, f"Flat rot6d length {r.shape[0]} is not divisible by 6" + r = r.reshape(-1, 6) + quats = rot6d_to_quat(r) + return quats.ravel() + + single = r.ndim == 1 + r = np.atleast_2d(r) # (N, 6) + col0 = r[:, :3] + col1 = r[:, 3:] + col0 = col0 / (np.linalg.norm(col0, axis=1, keepdims=True) + 1e-8) + dot = np.sum(col0 * col1, axis=1, keepdims=True) + col1 = col1 - dot * col0 + col1 = col1 / (np.linalg.norm(col1, axis=1, keepdims=True) + 1e-8) + col2 = np.cross(col0, col1) + rot_mat = np.stack([col0, col1, col2], axis=-1) # (N, 3, 3) + q_xyzw = R.from_matrix(rot_mat).as_quat() # (N, 4) xyzw + q_wxyz = q_xyzw[:, [3, 0, 1, 2]] + if single: + return q_wxyz[0].astype(np.float32) + return q_wxyz.astype(np.float32) + + +def compute_projected_gravity(base_quat: np.ndarray) -> np.ndarray: + """Compute projected gravity vector in robot's body frame from base quaternion. + + Projects the world gravity vector [0, 0, -1] into the robot's body frame by + rotating it by the inverse of the base quaternion. + + Args: + base_quat: Base quaternion [qw, qx, qy, qz] of shape (4,) + + Returns: + Projected gravity vector [gx, gy, gz] of shape (3,) in robot's body frame + """ + base_quat = np.asarray(base_quat, dtype=np.float64) + if base_quat.shape != (4,): + raise ValueError(f"base_quat must have shape (4,), got {base_quat.shape}") + + gravity_vec_world = np.array([0.0, 0.0, -1.0]) + base_rotation = R.from_quat(base_quat, scalar_first=True) + projected_gravity = base_rotation.inv().apply(gravity_vec_world) + + return projected_gravity.astype(np.float32) diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/zmq_state_subscriber.py b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/zmq_state_subscriber.py new file mode 100644 index 0000000000000000000000000000000000000000..5e3b900cccf211be4bb2c13c72f99fe7c7b0bd28 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/zmq_state_subscriber.py @@ -0,0 +1,138 @@ +""" +ZMQ utilities for subscribing to robot state and config from the C++ deploy process. + +Provides: +- ``ZMQStateSubscriber`` — non-blocking SUB on the ``g1_debug`` topic +- ``poll_robot_config_zmq`` — one-shot CONFIG topic reader +""" + +import time + +import msgpack +import msgpack_numpy as mnp +import numpy as np +import zmq + +STATE_ZMQ_TOPIC = "g1_debug" +CONFIG_ZMQ_TOPIC = "robot_config" +DEFAULT_STATE_ZMQ_PORT = 5557 + + +def _unpack_msgpack_zmq(raw: bytes, topic: str) -> dict: + """Strip a ZMQ topic prefix and decode the msgpack payload.""" + payload = raw[len(topic):] + return msgpack.unpackb(payload, raw=False) + + +def _convert_lists_to_numpy(data: dict) -> dict: + """Convert list values in a dict to numpy arrays.""" + if not isinstance(data, dict): + return data + result = {} + for key, value in data.items(): + if isinstance(value, (list, tuple)): + result[key] = np.array(value) + elif isinstance(value, dict): + result[key] = _convert_lists_to_numpy(value) + else: + result[key] = value + return result + + +class ZMQStateSubscriber: + """Non-blocking SUB on the ``g1_debug`` ZMQ topic for robot state. + + Uses ``zmq.CONFLATE`` so only the latest message is kept. + """ + + def __init__( + self, + host: str = "localhost", + port: int = DEFAULT_STATE_ZMQ_PORT, + topic: str = STATE_ZMQ_TOPIC, + ): + mnp.patch() + self._ctx = zmq.Context() + self._socket = self._ctx.socket(zmq.SUB) + self._socket.setsockopt_string(zmq.SUBSCRIBE, topic) + self._socket.setsockopt(zmq.CONFLATE, 1) + self._socket.setsockopt(zmq.RCVTIMEO, 0) + self._socket.connect(f"tcp://{host}:{port}") + self._topic = topic + self._msg = None + print(f"[ZMQStateSubscriber] Connected to tcp://{host}:{port} (topic: {topic})") + + def _poll(self): + """Poll for latest message (non-blocking).""" + try: + raw = self._socket.recv(zmq.NOBLOCK) + except zmq.Again: + return + + msg = _unpack_msgpack_zmq(raw, self._topic) + msg = _convert_lists_to_numpy(msg) + self._msg = msg + + def get_msg(self, clear: bool = True): + """Return the latest state message (or ``None``).""" + self._poll() + msg = self._msg + if clear: + self._msg = None + return msg + + def close(self): + self._socket.close() + self._ctx.term() + + def __del__(self): + try: + self.close() + except Exception: + pass + + +def poll_robot_config_zmq(host: str, port: int, timeout_sec: float = 0) -> dict: + """Wait for the ``robot_config`` message from the C++ ZMQ publisher. + + The publisher re-sends the config every ~2 s, so this simply polls with a + short receive timeout until a message arrives or *timeout_sec* elapses. + + Args: + timeout_sec: Max seconds to wait. ``0`` means wait indefinitely. + + Returns the decoded config dict. + + Raises: + TimeoutError: If *timeout_sec* > 0 and no message is received in time. + """ + mnp.patch() + ctx = zmq.Context() + sub = ctx.socket(zmq.SUB) + sub.setsockopt_string(zmq.SUBSCRIBE, CONFIG_ZMQ_TOPIC) + sub.setsockopt(zmq.RCVTIMEO, 500) + sub.setsockopt(zmq.CONFLATE, 1) + sub.connect(f"tcp://{host}:{port}") + + if timeout_sec > 0: + print(f"[Config] Waiting up to {timeout_sec}s for robot_config on tcp://{host}:{port} ...") + else: + print(f"[Config] Waiting for robot_config on tcp://{host}:{port} ... is gear_sonic_deploy running?") + deadline = (time.monotonic() + timeout_sec) if timeout_sec > 0 else None + try: + while True: + if deadline is not None and time.monotonic() >= deadline: + raise TimeoutError( + f"[Config] No robot_config received on tcp://{host}:{port} " + f"within {timeout_sec}s. Is the C++ deploy process running?" + ) + try: + raw = sub.recv() + config = _unpack_msgpack_zmq(raw, CONFIG_ZMQ_TOPIC) + print(f"[Config] Received robot_config ({len(config)} fields)") + return config + except zmq.Again: + pass + finally: + sub.close() + ctx.term() diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/inference/__init__.py b/GR00T-WholeBodyControl/gear_sonic/utils/inference/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/inference/__pycache__/__init__.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/inference/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ab46ff58e6f7cacb27aad18c4c480091b6aa325d Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/inference/__pycache__/__init__.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/inference/__pycache__/initial_poses.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/inference/__pycache__/initial_poses.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..427eb4629960e31933c5702a2aafab251c235e6c Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/inference/__pycache__/initial_poses.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/inference/__pycache__/vla_utils.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/inference/__pycache__/vla_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1ca10ec254821332f32d548292b0c5d923378d3b Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/inference/__pycache__/vla_utils.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/inference/dex1_head.py b/GR00T-WholeBodyControl/gear_sonic/utils/inference/dex1_head.py new file mode 100644 index 0000000000000000000000000000000000000000..0cec082f86c5a8767ada525bea91610b0af6dbb9 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/inference/dex1_head.py @@ -0,0 +1,273 @@ +"""Adapters for the deployed Dex1 grippers and two-axis servo head. + +The fine-tuned checkpoint used by the real-robot client has the following +action schema:: + + motion_token[64] + left_hand_joints[1] + right_hand_joints[1] + head_joints[2] + +Dex1 and head feedback/commands are transported as msgpack dictionaries over +dedicated ZMQ sockets. This module contains no robot connections by itself; +the main inference program creates all sockets explicitly. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import time +from typing import Any, Callable + +import numpy as np + + +ACTION_DIMS = { + "motion_token": 64, + "left_hand_joints": 1, + "right_hand_joints": 1, + "head_joints": 2, +} + + +def action_field(action: dict[str, Any], key: str) -> np.ndarray: + """Read a policy action field with or without the ``action.`` prefix.""" + for candidate in (key, f"action.{key}"): + if candidate in action: + return np.asarray(action[candidate], dtype=np.float32) + raise KeyError(f"Missing action field {key!r}; available fields: {sorted(action)}") + + +def validate_action_chunk( + action: dict[str, Any], expected_horizon: int +) -> dict[str, np.ndarray]: + """Validate one standard GR00T action chunk and normalize it to ``[T, D]``.""" + if not isinstance(action, dict): + raise TypeError(f"Policy action must be a dict, got {type(action).__name__}") + + normalized: dict[str, np.ndarray] = {} + horizon: int | None = None + for key, width in ACTION_DIMS.items(): + value = action_field(action, key) + if value.ndim == 3: + if value.shape[0] != 1: + raise ValueError(f"{key} batch must be 1, got shape {value.shape}") + value = value[0] + if value.ndim != 2 or value.shape[1] != width: + raise ValueError(f"{key} must have shape [T, {width}], got {value.shape}") + if not np.isfinite(value).all(): + raise ValueError(f"{key} contains NaN or Inf") + + if horizon is None: + horizon = value.shape[0] + elif value.shape[0] != horizon: + raise ValueError( + f"Action fields have inconsistent horizons: {horizon} and {value.shape[0]}" + ) + normalized[key] = np.ascontiguousarray(value, dtype=np.float32) + + if horizon != expected_horizon: + raise ValueError(f"Expected action horizon {expected_horizon}, got {horizon}") + if float(np.abs(normalized["motion_token"]).max()) > 1.25: + raise ValueError("motion_token exceeds the SONIC safety bound of 1.25") + + # The trained gripper targets are normalized to [0, 1]. A small margin is + # tolerated for model noise; larger excursions reject the whole chunk. + for key in ("left_hand_joints", "right_hand_joints"): + low = float(normalized[key].min()) + high = float(normalized[key].max()) + if low < -0.25 or high > 1.25: + raise ValueError(f"{key} is far outside [0, 1]: min={low:.4f}, max={high:.4f}") + + return normalized + + +def _finite_scalar(value: Any, name: str) -> float: + array = np.asarray(value, dtype=np.float64).reshape(-1) + if array.size != 1: + raise ValueError(f"{name} must contain exactly one scalar, got shape {array.shape}") + scalar = float(array[0]) + if not np.isfinite(scalar): + raise ValueError(f"{name} must be finite, got {scalar}") + return scalar + + +def parse_dex1_state(message: dict[str, Any]) -> tuple[float, float]: + """Parse one measured scalar per Dex1 gripper.""" + if not isinstance(message, dict): + raise ValueError(f"Dex1 state must be a dict, got {type(message).__name__}") + left = _finite_scalar(message.get("left_hand_joints", []), "left Dex1 state") + right = _finite_scalar(message.get("right_hand_joints", []), "right Dex1 state") + return left, right + + +def parse_head_state(message: dict[str, Any]) -> np.ndarray: + """Parse absolute ``[yaw, pitch]`` feedback in radians.""" + if not isinstance(message, dict): + raise ValueError(f"Head state must be a dict, got {type(message).__name__}") + if "head_joints" in message: + head = np.asarray(message["head_joints"], dtype=np.float64).reshape(-1) + elif "yaw_position" in message and "pitch_position" in message: + head = np.asarray( + [message["yaw_position"], message["pitch_position"]], dtype=np.float64 + ).reshape(-1) + else: + raise ValueError( + "Head state must contain head_joints[2] or yaw_position/pitch_position" + ) + if head.shape != (2,) or not np.isfinite(head).all(): + raise ValueError(f"Head state must contain two finite angles, got {head}") + return np.ascontiguousarray(head, dtype=np.float32) + + +class LatestMsgpackSubscriber: + """A conflated ZMQ subscriber that retains the latest validated message.""" + + def __init__( + self, + context: Any, + endpoint: str, + validator: Callable[[dict[str, Any]], Any], + ): + import zmq + + self.endpoint = endpoint + self.socket = context.socket(zmq.SUB) + self.socket.setsockopt(zmq.SUBSCRIBE, b"") + self.socket.setsockopt(zmq.CONFLATE, 1) + self.socket.connect(endpoint) + self.validator = validator + self.latest: dict[str, Any] | None = None + self.received_at = 0.0 + + def read(self) -> dict[str, Any] | None: + import msgpack + import zmq + + while True: + try: + raw = self.socket.recv(zmq.NOBLOCK) + except zmq.Again: + return self.latest + message = msgpack.unpackb(raw, raw=False) + self.validator(message) + self.latest = message + self.received_at = time.monotonic() + + def is_fresh(self, timeout: float) -> bool: + return ( + timeout > 0.0 + and self.latest is not None + and time.monotonic() - self.received_at <= timeout + ) + + def close(self) -> None: + self.socket.close(linger=0) + + +@dataclass +class Dex1CommandFilter: + """Clamp scalar targets and rate-limit them from the last measured command.""" + + max_step: float = 0.08 + left_safe: float = 1.0 + right_safe: float = 1.0 + warning_sink: Callable[[str], None] = print + + def __post_init__(self) -> None: + if not np.isfinite(self.max_step) or self.max_step <= 0: + raise ValueError("max_step must be positive and finite") + self.reset(self.left_safe, self.right_safe) + + def reset(self, left_measured: float, right_measured: float) -> None: + self.left_safe = float(np.clip(_finite_scalar(left_measured, "left Dex1"), 0.0, 1.0)) + self.right_safe = float( + np.clip(_finite_scalar(right_measured, "right Dex1"), 0.0, 1.0) + ) + + def _update_one(self, value: Any, previous: float, name: str) -> float: + raw = _finite_scalar(value, name) + clipped = float(np.clip(raw, 0.0, 1.0)) + if clipped != raw: + self.warning_sink(f"[warning] {name} target {raw:.4f} clipped to {clipped:.4f}") + delta = float(np.clip(clipped - previous, -self.max_step, self.max_step)) + return previous + delta + + def update(self, left: Any, right: Any) -> tuple[float, float]: + self.left_safe = self._update_one(left, self.left_safe, "left Dex1") + self.right_safe = self._update_one(right, self.right_safe, "right Dex1") + return self.left_safe, self.right_safe + + +@dataclass +class HeadCommandFilter: + """Apply mechanical limits and per-frame rate limits to head targets.""" + + yaw_limits: tuple[float, float] = (-1.2, 1.2) + pitch_limits: tuple[float, float] = (-0.6, 0.6) + max_yaw_step: float = 0.08 + max_pitch_step: float = 0.06 + last_head: np.ndarray | None = None + + def __post_init__(self) -> None: + for name, limits in (("yaw_limits", self.yaw_limits), ("pitch_limits", self.pitch_limits)): + values = np.asarray(limits, dtype=np.float64) + if values.shape != (2,) or not np.isfinite(values).all() or values[0] >= values[1]: + raise ValueError(f"{name} must contain two increasing finite values") + for name, value in ( + ("max_yaw_step", self.max_yaw_step), + ("max_pitch_step", self.max_pitch_step), + ): + if not np.isfinite(value) or value <= 0: + raise ValueError(f"{name} must be positive and finite") + + def reset(self, measured_head: Any) -> None: + measured = parse_head_state({"head_joints": measured_head}) + measured[0] = np.clip(measured[0], *self.yaw_limits) + measured[1] = np.clip(measured[1], *self.pitch_limits) + self.last_head = measured + + def update(self, head: Any) -> np.ndarray: + target = parse_head_state({"head_joints": head}) + target[0] = np.clip(target[0], *self.yaw_limits) + target[1] = np.clip(target[1], *self.pitch_limits) + if self.last_head is not None: + delta = np.clip( + target - self.last_head, + [-self.max_yaw_step, -self.max_pitch_step], + [self.max_yaw_step, self.max_pitch_step], + ) + target = self.last_head + delta + self.last_head = np.ascontiguousarray(target, dtype=np.float32) + return self.last_head.copy() + + +def pack_gripper_command( + left: float, right: float, timestamp_ns: int | None = None +) -> bytes: + """Serialize the command map consumed by the robot-side Dex1 service.""" + import msgpack + + left_value = float(np.clip(_finite_scalar(left, "left Dex1 command"), 0.0, 1.0)) + right_value = float(np.clip(_finite_scalar(right, "right Dex1 command"), 0.0, 1.0)) + return msgpack.packb( + { + "left_hand_joints": [left_value], + "right_hand_joints": [right_value], + "timestamp_ns": int(timestamp_ns if timestamp_ns is not None else time.time_ns()), + }, + use_bin_type=True, + ) + + +def pack_head_command(yaw: float, pitch: float, timestamp_ns: int | None = None) -> bytes: + """Serialize absolute yaw/pitch targets for the robot-side head service.""" + import msgpack + + return msgpack.packb( + { + "yaw_position": _finite_scalar(yaw, "yaw command"), + "pitch_position": _finite_scalar(pitch, "pitch command"), + "left_joystick": [0.0, 0.0], + "right_joystick": [0.0, 0.0], + "timestamp_ns": int(timestamp_ns if timestamp_ns is not None else time.time_ns()), + }, + use_bin_type=True, + ) diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/inference/initial_poses.py b/GR00T-WholeBodyControl/gear_sonic/utils/inference/initial_poses.py new file mode 100644 index 0000000000000000000000000000000000000000..124758e70192b9258106c542a8cb6efc6f0e2186 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/inference/initial_poses.py @@ -0,0 +1,32 @@ +"""Default initial poses for VLA inference. + +These arrays are sent to the C++ control loop when the user presses 'i' +to move the robot to a known starting configuration before inference begins. + +WARNING: The initial motion token below is specific to the SONIC checkpoint used +during training. Different SONIC checkpoints encode different latent spaces, so +this token will produce a different (and likely incorrect) pose if you switch to +a different SONIC checkpoint. When changing the SONIC checkpoint, you MUST update +LATENT_INITIAL_MOTION_TOKEN to a value that corresponds to a known safe standing +pose in the new checkpoint's latent space. +""" + +import numpy as np + +# 64-dim motion token for a stable standing pose. +# CHECKPOINT-SPECIFIC: this value must be updated if the SONIC checkpoint changes. +LATENT_INITIAL_MOTION_TOKEN = np.array( + [ + -0.0625, 0.0000, -0.0625, -0.1250, -0.1875, -0.0625, 0.1875, + 0.2500, 0.1875, -0.1250, 0.0625, -0.0625, -0.2500, -0.2500, + -0.3125, -0.0625, 0.0000, -0.0625, -0.1250, -0.1875, 0.0000, + -0.2500, 0.0000, -0.2500, -0.0625, 0.0625, 0.1250, -0.1250, + 0.2500, 0.1875, 0.2500, -0.1250, 0.1250, 0.1875, -0.0625, + 0.0000, -0.1875, -0.1875, 0.2500, 0.0000, 0.0000, -0.1250, + 0.0625, 0.0000, -0.0625, -0.0625, 0.1875, -0.0625, 0.0000, + 0.0625, 0.1250, 0.0625, 0.1250, 0.0625, 0.1250, 0.0000, + 0.1250, 0.1875, 0.0000, 0.0000, 0.0625, 0.0625, 0.1875, + 0.0625, + ], + dtype=np.float32, +) diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/inference/vla_utils.py b/GR00T-WholeBodyControl/gear_sonic/utils/inference/vla_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3165a31d8a341558868eeb9234701857bb4830af --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/inference/vla_utils.py @@ -0,0 +1,106 @@ +"""Utility functions for VLA inference. + +Includes action processing, observation preparation, latency compensation, +and inference scheduling logic. +""" + +from typing import Any, Dict + +import numpy as np + +from gear_sonic.data.robot_model.robot_model import RobotModel + + +def concat_action(robot_model: RobotModel, goal: Dict[str, Any]) -> Dict[str, Any]: + """Process the action dict from the policy into a flat dict. + + Strips ``action.`` prefixes from keys (if present) and returns the result. + + Args: + robot_model: RobotModel instance (unused for latent actions, kept for API compat). + goal: Action dict from policy. + + Returns: + Processed action dict with prefixes stripped. + """ + processed_goal = {} + for key, value in goal.items(): + processed_goal[key.replace("action.", "")] = value + return processed_goal + + +def prepare_observation_for_eval(robot_model: RobotModel, obs: dict) -> dict: + """Split whole-body ``q`` into per-joint-group state keys for the policy. + + Populates ``obs["state"]`` with ``left_arm``, ``right_arm``, ``waist``, + ``left_leg``, ``right_leg``, ``left_hand``, ``right_hand`` sub-keys + using the nested dict format expected by ``Gr00tPolicy``. + + Args: + robot_model: RobotModel instance. + obs: Observation dict containing ``"q"`` key and a ``"state"`` sub-dict. + + Returns: + Modified observation dict with ``obs["state"]`` populated. + """ + assert "q" in obs, "q is not in the observation" + + whole_q = obs["q"] + assert whole_q.shape[-1] == robot_model.num_joints, "q has wrong shape" + + if "state" not in obs: + obs["state"] = {} + + obs["state"]["left_arm"] = whole_q[..., robot_model.get_joint_group_indices("left_arm")] + obs["state"]["right_arm"] = whole_q[..., robot_model.get_joint_group_indices("right_arm")] + obs["state"]["waist"] = whole_q[..., robot_model.get_joint_group_indices("waist")] + obs["state"]["left_leg"] = whole_q[..., robot_model.get_joint_group_indices("left_leg")] + obs["state"]["right_leg"] = whole_q[..., robot_model.get_joint_group_indices("right_leg")] + obs["state"]["left_hand"] = whole_q[..., robot_model.get_joint_group_indices("left_hand")] + obs["state"]["right_hand"] = whole_q[..., robot_model.get_joint_group_indices("right_hand")] + + return obs + + +def calculate_latency_compensated_index( + inference_delay: float, control_freq: float, action_horizon: int +) -> int: + """Calculate the starting action index compensating for inference latency. + + When inference completes, some time has elapsed, so we skip the first few + actions that are now "stale" and start from a later index in the chunk. + + Args: + inference_delay: Time elapsed since inference started (seconds). + control_freq: Control loop frequency (Hz), e.g. 20. + action_horizon: Total number of actions in the chunk, e.g. 16. + + Returns: + Starting index (0 to action_horizon-1) for the action chunk. + """ + raw_index = np.round(inference_delay * control_freq) + return int(np.clip(raw_index, 0, action_horizon - 1)) + + +def should_trigger_new_inference( + cached_chunk_exists: bool, + inference_thread_running: bool, + time_since_last_inference: float, + inference_interval: float, +) -> bool: + """Determine if a new inference should be triggered. + + Args: + cached_chunk_exists: Whether we have a cached action chunk. + inference_thread_running: Whether inference is currently running. + time_since_last_inference: Time elapsed since last inference started (seconds). + inference_interval: Minimum time between inferences (seconds). + + Returns: + True if new inference should start. + """ + if not cached_chunk_exists: + return True + if inference_thread_running: + return False + return time_since_last_inference >= inference_interval diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/inference_helpers.py b/GR00T-WholeBodyControl/gear_sonic/utils/inference_helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..5ed61f2dabe7e63a867dbf487b37bdb963de40c8 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/inference_helpers.py @@ -0,0 +1,467 @@ +"""Export trained RL policies to ONNX format for deployment. + +Focuses on SONIC / universal-token model export: + - encoder+decoder pair (``export_universal_token_module_as_onnx``) + - all encoders (``export_universal_token_encoders_as_onnx``) + - decoder only (``export_universal_token_decoder_as_onnx``) + - generic policy (``export_policy_as_onnx``) +""" + +import copy +import os + +import torch +from torch import nn + + +# --------------------------------------------------------------------------- +# Generic policy export +# --------------------------------------------------------------------------- + + +def export_policy_as_onnx(inference_model, path, exported_policy_name, example_obs_dict): + """Export a PPO actor policy as an ONNX model. + + Args: + inference_model: Dict containing an ``"actor"`` key with the actor module. + path: Directory path to save the exported model. + exported_policy_name: Filename for the exported ONNX model. + example_obs_dict: Example observation dict for ONNX tracing. + """ + os.makedirs(path, exist_ok=True) + path = os.path.join(path, exported_policy_name) + + actor = copy.deepcopy(inference_model["actor"]).to("cpu") + actor.eval() + + class PPOWrapper(nn.Module): + def __init__(self, actor): + super().__init__() + self.actor = actor + + def forward(self, obs_dict): + return self.actor.act_inference(obs_dict) + + wrapper = PPOWrapper(actor) + example_input_list = {"obs_dict": example_obs_dict} + with torch.no_grad(): + torch.onnx.export( + wrapper, + example_input_list, + path, + verbose=True, + input_names=["obs_dict"], + output_names=["action"], + opset_version=13, + ) + + +# --------------------------------------------------------------------------- +# Universal-token encoder + decoder pair +# --------------------------------------------------------------------------- + + +def export_universal_token_module_as_onnx( + universal_token_module, encoder_name, decoder_name, path, exported_model_name, batch_size=1 +): + """Export UniversalTokenModule with a specific encoder and decoder as ONNX. + + The exported model accepts a single flattened 2-D tensor whose layout is:: + + [tokenizer_observations | proprioception] + + Tokenizer observations are ordered according to the union of features + required by the chosen encoder and decoder. + + Args: + universal_token_module: The UniversalTokenModule instance. + encoder_name: Name of the encoder to use. + decoder_name: Name of the decoder to use. + path: Directory path to save the ONNX model. + exported_model_name: Name of the exported ONNX file. + batch_size: Batch size for example input. + """ + os.makedirs(path, exist_ok=True) + full_path = os.path.join(path, exported_model_name) + + module = copy.deepcopy(universal_token_module).to("cpu") + module.eval() + + # Determine which tokenizer observations are needed (union of encoder and decoder inputs) + encoder_input_features = module.encoder_input_features[encoder_name] + decoder_input_features = module.decoder_input_features[decoder_name] + + special_keys = {"token", "token_flattened", "proprioception", "action", "meta_action"} + encoder_tokenizer_obs = [f for f in encoder_input_features if f not in special_keys] + decoder_tokenizer_obs = [f for f in decoder_input_features if f not in special_keys] + + # required_tokenizer_obs = list(set(encoder_tokenizer_obs + decoder_tokenizer_obs)) + required_set = set(encoder_tokenizer_obs + decoder_tokenizer_obs) + required_tokenizer_obs = [ + name for name in module.tokenizer_obs_names if name in required_set + ] + + # Calculate total input dimension + total_feature_dim = 0 + for feature_name in required_tokenizer_obs: + if feature_name in module.tokenizer_obs_dims: + feature_dims = module.tokenizer_obs_dims[feature_name] + total_feature_dim += torch.prod(torch.tensor(feature_dims)).item() + + proprioception_dim = module.obs_dim_dict.get("actor_obs", 0) + total_feature_dim += proprioception_dim + + example_input = torch.randn(batch_size, total_feature_dim, device="cpu") + + class UniversalTokenWrapper(nn.Module): + def __init__(self, module, encoder_name, decoder_name, required_tokenizer_obs): + super().__init__() + self.module = module + self.encoder_name = encoder_name + self.decoder_name = decoder_name + self.required_tokenizer_obs = required_tokenizer_obs + + def forward(self, obs_dict): + combined_input = obs_dict["actor_obs"] + combined_input = combined_input.unsqueeze(1) # add sequence dimension + + tokenizer_feature_dim = 0 + for feature_name in self.required_tokenizer_obs: + feature_dims = self.module.tokenizer_obs_dims[feature_name] + tokenizer_feature_dim += torch.prod(torch.tensor(feature_dims)).item() + + tokenizer_part = combined_input[..., :tokenizer_feature_dim] + proprioception_part = combined_input[..., tokenizer_feature_dim:] + + # Reconstruct tokenizer_obs dict from flattened tensor + tokenizer_obs = {} + index = 0 + for feature_name in self.required_tokenizer_obs: + feature_dims = tuple(self.module.tokenizer_obs_dims[feature_name]) + feature_size = torch.prod(torch.tensor(feature_dims)).item() + feature_data = tokenizer_part[..., index : index + feature_size] + tokenizer_obs[feature_name] = feature_data.reshape( + feature_data.shape[:2] + feature_dims + ) + index += feature_size + + # Get encoder-specific observations + encoder_input_features = self.module.encoder_input_features[self.encoder_name] + encoder_tokenizer_obs = { + k: tokenizer_obs[k] for k in encoder_input_features if k in tokenizer_obs + } + + # Encode + encoded_tokens, _ = self.module.encode(self.encoder_name, encoder_tokenizer_obs) + encoded_tokens = encoded_tokens.unsqueeze(1) # add back the sequence dimension + + # Prepare decode input + decode_input_dict = { + "token": encoded_tokens, + "token_flattened": encoded_tokens.flatten(start_dim=-2), + "proprioception": proprioception_part, + } + decode_input_dict.update(tokenizer_obs) + + # Decode + decoded_output = self.module.decode(self.decoder_name, decode_input_dict) + + if "action" in decoded_output: + return decoded_output["action"].squeeze(1) + else: + output = torch.cat(list(decoded_output.values()), dim=-1) + return output.squeeze(1) + + wrapper = UniversalTokenWrapper(module, encoder_name, decoder_name, required_tokenizer_obs) + + obs_dict = {"actor_obs": example_input} + example_input_list = {"obs_dict": obs_dict} + with torch.no_grad(): + torch.onnx.export( + wrapper, + example_input_list, + full_path, + verbose=True, + input_names=["obs_dict"], + output_names=["action"], + opset_version=13, + ) + + print(f"\nExported ONNX model: {encoder_name} encoder -> {decoder_name} decoder") # noqa: T201 + print(f"Saved to: {full_path}") # noqa: T201 + print(f"Required tokenizer observations: {required_tokenizer_obs}") # noqa: T201 + print( # noqa: T201 + f"Input shape: {example_input.shape} " + f"(tokenizer: {total_feature_dim - proprioception_dim}, proprioception: {proprioception_dim})" + ) + + +# --------------------------------------------------------------------------- +# Encoders-only export (all encoders, dynamic selection via encoder_index) +# --------------------------------------------------------------------------- + + +def export_universal_token_encoders_as_onnx( + universal_token_module, path, exported_model_name, batch_size=1 +): + """Export only the ENCODERS of a UniversalTokenModule as ONNX. + + Uses ``encoder_index`` to dynamically select which encoder to run. + + Input layout:: + + [encoder_index(1) | tokenizer_observations] + + Args: + universal_token_module: The UniversalTokenModule instance. + path: Directory path to save the ONNX model. + exported_model_name: Name of the exported ONNX file. + batch_size: Batch size for example input. + """ + os.makedirs(path, exist_ok=True) + full_path = os.path.join(path, exported_model_name) + + module = copy.deepcopy(universal_token_module).to("cpu") + module.eval() + + encoder_names = module.encoders_to_iterate + if not encoder_names: + raise ValueError("No encoders found in the module") + + special_keys = {"token", "token_flattened", "proprioception", "action", "meta_action"} + + # Start with encoder_index, then follow module.tokenizer_obs_names order + all_tokenizer_obs = ["encoder_index"] + + features_needed = set() + for enc_name in encoder_names: + encoder_input_features = module.encoder_input_features[enc_name] + features_needed.update([f for f in encoder_input_features if f not in special_keys]) + + for obs_name in module.tokenizer_obs_names: + if obs_name in features_needed: + all_tokenizer_obs.append(obs_name) + + required_tokenizer_obs = all_tokenizer_obs + + # Calculate tokenizer input dimension (only required observations) + tokenizer_feature_dim = 0 + for feature_name in required_tokenizer_obs: + if feature_name in module.tokenizer_obs_dims: + feature_dims = module.tokenizer_obs_dims[feature_name] + tokenizer_feature_dim += torch.prod(torch.tensor(feature_dims)).item() + + example_tokenizer_obs = torch.randn(batch_size, tokenizer_feature_dim, device="cpu") + example_encoder_index = torch.zeros((batch_size, 1), device="cpu") + + class EncodersOnlyWrapper(nn.Module): + def __init__(self, module, encoder_names, required_tokenizer_obs): + super().__init__() + self.module = module + self.encoder_names = encoder_names + self.required_tokenizer_obs = required_tokenizer_obs + + def forward(self, obs_dict): + inputs = obs_dict["actor_obs"] + encoder_index = inputs[..., 0].long() + encoder_index_onehot = torch.zeros( + (encoder_index.shape[0], len(self.encoder_names)), device="cpu" + ) + encoder_index_onehot[torch.arange(encoder_index.shape[0]), encoder_index] = 1.0 + tokenizer_input = inputs[..., 1:] + + tokenizer_input = tokenizer_input.unsqueeze(1) # add sequence dimension (B, 1, F) + + # Reconstruct required tokenizer observations only + tokenizer_obs = {} + index = 0 + for feature_name in self.required_tokenizer_obs: + feature_dims = tuple(self.module.tokenizer_obs_dims[feature_name]) + feature_size = torch.prod(torch.tensor(feature_dims)).item() + feature_data = tokenizer_input[..., index : index + feature_size] + tokenizer_obs[feature_name] = feature_data.reshape( + feature_data.shape[:2] + feature_dims + ) + index += feature_size + + all_encoded_tokens = [] + for encoder_name in self.encoder_names: + encoder_input_features = self.module.encoder_input_features[encoder_name] + encoder_tokenizer_obs = { + k: tokenizer_obs[k] for k in encoder_input_features if k in tokenizer_obs + } + encoded_tokens, _ = self.module.encode(encoder_name, encoder_tokenizer_obs) + all_encoded_tokens.append(encoded_tokens) + + # Stack all encoded tokens: (num_encoders, B, token_features...) + stacked_tokens = torch.stack(all_encoded_tokens, dim=0) + + # Use encoder_index to select the appropriate tokens via weighted sum + encoder_weights = encoder_index_onehot.t() # (num_encoders, B) + for _ in range(len(stacked_tokens.shape) - 2): + encoder_weights = encoder_weights.unsqueeze(-1) + + weighted_tokens = stacked_tokens * encoder_weights + selected_tokens = weighted_tokens.sum(dim=0) # (B, token_features...) + selected_tokens = selected_tokens.flatten(start_dim=-2) + + return selected_tokens + + wrapper = EncodersOnlyWrapper(module, encoder_names, required_tokenizer_obs) + + obs_dict = { + "actor_obs": torch.cat([example_encoder_index, example_tokenizer_obs], dim=-1), + } + example_input_dict = {"obs_dict": obs_dict} + + with torch.no_grad(): + torch.onnx.export( + wrapper, + example_input_dict, + full_path, + verbose=True, + input_names=["obs_dict"], + output_names=["encoded_tokens"], + opset_version=13, + ) + + print( # noqa: T201 + f"\nExported ENCODERS ONLY ONNX model with {len(encoder_names)} encoders: {encoder_names}" + ) + print(f"Saved to: {full_path}") # noqa: T201 + print(f"Required tokenizer observations: {required_tokenizer_obs}") # noqa: T201 + print( # noqa: T201 + f"Input shapes: tokenizer_obs={example_tokenizer_obs.shape}, " + f"encoder_index={example_encoder_index.shape}" + ) + + +# --------------------------------------------------------------------------- +# Decoder-only export +# --------------------------------------------------------------------------- + + +def export_universal_token_decoder_as_onnx( + universal_token_module, decoder_name, path, exported_model_name, batch_size=1 +): + """Export only the DECODER of a UniversalTokenModule as ONNX. + + Input layout:: + + [encoded_tokens | tokenizer_observations | proprioception] + + Args: + universal_token_module: The UniversalTokenModule instance. + decoder_name: Name of the decoder to use. + path: Directory path to save the ONNX model. + exported_model_name: Name of the exported ONNX file. + batch_size: Batch size for example input. + """ + os.makedirs(path, exist_ok=True) + full_path = os.path.join(path, exported_model_name) + + module = copy.deepcopy(universal_token_module).to("cpu") + module.eval() + + # Determine which tokenizer observations the decoder needs + decoder_input_features = module.decoder_input_features[decoder_name] + special_keys = {"token", "token_flattened", "proprioception", "action", "meta_action"} + required_tokenizer_obs = [f for f in decoder_input_features if f not in special_keys] + + # Calculate tokenizer input dimension (only required observations) + tokenizer_feature_dim = 0 + for feature_name in required_tokenizer_obs: + if feature_name in module.tokenizer_obs_dims: + feature_dims = module.tokenizer_obs_dims[feature_name] + tokenizer_feature_dim += torch.prod(torch.tensor(feature_dims)).item() + + proprioception_dim = module.obs_dim_dict.get("actor_obs", 0) + token_total_dim = module.token_total_dim + + example_encoded_tokens = torch.randn(batch_size, token_total_dim, device="cpu") + example_tokenizer = torch.randn(batch_size, tokenizer_feature_dim, device="cpu") + example_proprioception = torch.randn(batch_size, proprioception_dim, device="cpu") + + class DecoderOnlyWrapper(nn.Module): + def __init__(self, module, decoder_name, required_tokenizer_obs): + super().__init__() + self.module = module + self.decoder_name = decoder_name + self.required_tokenizer_obs = required_tokenizer_obs + + def forward(self, obs_dict): + inputs = obs_dict["actor_obs"] + + proprioception_dim = self.module.obs_dim_dict["actor_obs"] + token_total_dim = self.module.token_total_dim + + tokenizer_feature_dim = 0 + for feature_name in self.required_tokenizer_obs: + feature_dims = self.module.tokenizer_obs_dims[feature_name] + tokenizer_feature_dim += torch.prod(torch.tensor(feature_dims)).item() + + # Split input: [encoded_tokens | tokenizer_obs | proprioception] + encoded_tokens = inputs[..., :token_total_dim] + tokenizer_part = inputs[..., token_total_dim : token_total_dim + tokenizer_feature_dim] + proprioception = inputs[..., -proprioception_dim:] + + # Add sequence dimension + encoded_tokens = encoded_tokens.unsqueeze(1) + tokenizer_part = tokenizer_part.unsqueeze(1) + proprioception = proprioception.unsqueeze(1) + + # Reconstruct tokenizer observations from flat tensor + tokenizer_obs = {} + index = 0 + for feature_name in self.required_tokenizer_obs: + feature_dims = tuple(self.module.tokenizer_obs_dims[feature_name]) + feature_size = torch.prod(torch.tensor(feature_dims)).item() + feature_data = tokenizer_part[..., index : index + feature_size] + tokenizer_obs[feature_name] = feature_data.reshape( + feature_data.shape[:2] + feature_dims + ) + index += feature_size + + # Prepare decode input + decode_input_dict = { + "token_flattened": encoded_tokens, + "proprioception": proprioception, + } + decode_input_dict.update(tokenizer_obs) + + # Decode + decoded_output = self.module.decode(self.decoder_name, decode_input_dict) + + if "action" in decoded_output: + return decoded_output["action"].squeeze(1) + else: + output = torch.cat(list(decoded_output.values()), dim=-1) + return output.squeeze(1) + + wrapper = DecoderOnlyWrapper(module, decoder_name, required_tokenizer_obs) + + obs_dict = { + "actor_obs": torch.cat( + [example_encoded_tokens, example_tokenizer, example_proprioception], dim=-1 + ), + } + example_input_dict = {"obs_dict": obs_dict} + + with torch.no_grad(): + torch.onnx.export( + wrapper, + example_input_dict, + full_path, + verbose=True, + input_names=["obs_dict"], + output_names=["action"], + opset_version=13, + ) + + print(f"\nExported DECODER ONLY ONNX model with decoder: {decoder_name}") # noqa: T201 + print(f"Saved to: {full_path}") # noqa: T201 + print(f"Required tokenizer observations: {required_tokenizer_obs}") # noqa: T201 + print( # noqa: T201 + f"Input shapes: encoded_tokens={example_encoded_tokens.shape}, " + f"tokenizer={example_tokenizer.shape}, proprioception={example_proprioception.shape}" + ) diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/logging.py b/GR00T-WholeBodyControl/gear_sonic/utils/logging.py new file mode 100644 index 0000000000000000000000000000000000000000..ba16e5057f38c1f7e875e4be5d9720de7defbe25 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/logging.py @@ -0,0 +1,48 @@ +from contextlib import contextmanager +import logging +import os +import sys + +from loguru import logger + + +class HydraLoggerBridge(logging.Handler): + def emit(self, record): + # Get corresponding loguru level + try: + level = logger.level(record.levelname).name + except ValueError: + level = record.levelno + + # Find caller from where the logged message originated + frame, depth = logging.currentframe(), 2 + while frame and frame.f_code.co_filename == logging.__file__: + frame = frame.f_back + depth += 1 + + logger.opt(depth=depth, exception=record.exc_info).log(level, record.getMessage()) + + +class LoguruStream: + def write(self, message): + if message.strip(): # Only log non-empty messages + logger.info(message.strip()) # Changed to debug level + + def flush(self): + pass + + +@contextmanager +def capture_stdout_to_loguru(): + logger.remove() + logger.add(sys.stdout, level="INFO") + loguru_stream = LoguruStream() + old_stdout = sys.stdout + sys.stdout = loguru_stream + try: + yield + finally: + sys.stdout = old_stdout + logger.remove() + console_log_level = os.environ.get("LOGURU_LEVEL", "INFO").upper() + logger.add(sys.stdout, level=console_log_level, colorize=True) diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__init__.py b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/__init__.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d985d7b1c45c7bf7d4271a2e1962d9a666809ac1 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/__init__.cpython-311.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/__init__.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e539a5ad295bfc09298c9a910a8f88983808eba6 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/__init__.cpython-312.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/motion_lib_robot.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/motion_lib_robot.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2eb56ca710c82c71b6608465fcbe82244e88f790 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/motion_lib_robot.cpython-311.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/motion_lib_robot.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/motion_lib_robot.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7989159d357e1d4a3dfad0d9e35a6849aec91ba1 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/motion_lib_robot.cpython-312.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/skeleton.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/skeleton.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd7ca3aab0fcdbc2391445ef142a06dd0b57857d Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/skeleton.cpython-311.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/skeleton.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/skeleton.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5a1654c0e10278b6dacf92d0942fc46ebcb76e8a Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/skeleton.cpython-312.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/torch_humanoid_batch.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/torch_humanoid_batch.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b6a242be2a7ca55e0ba37b182dfefda2080e09b3 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/torch_humanoid_batch.cpython-311.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/torch_humanoid_batch.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/torch_humanoid_batch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..af3d1b520beb73f10ede3966cc26b3befd052f23 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/__pycache__/torch_humanoid_batch.cpython-312.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/motion_lib_base.py b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/motion_lib_base.py new file mode 100644 index 0000000000000000000000000000000000000000..2d552a7f9d69f0fbd6135f9270847eb9cfe5a5c6 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/motion_lib_base.py @@ -0,0 +1,2847 @@ +import enum +import gc +import glob +import os +import os.path as osp +from pathlib import Path +import random +import re +import resource + +import easydict +import joblib +from loguru import logger +import numpy as np +from rich import progress +from scipy.spatial import transform +import torch +import torch.multiprocessing as mp + +from gear_sonic.isaac_utils import rotations +from gear_sonic.trl.utils import common +from gear_sonic.utils.motion_lib import skeleton + + +class FixHeightMode(enum.Enum): + no_fix = 0 + full_fix = 1 + ankle_fix = 2 + + +class MotionlibMode(enum.Enum): + file = 1 + directory = 2 + + +def to_torch(tensor): + if torch.is_tensor(tensor): + return tensor + else: + return torch.from_numpy(tensor) + + +def is_navigation_motion(motion_key): + return ( + motion_key.startswith("2025") + or motion_key.startswith("walking_2025") + or motion_key.startswith("running_2025") + or motion_key.startswith("slow_walk_2025") + ) + + +def interpolate_translation_data( + data, + source_fps, + target_fps, + num_frames, + max_num_objects=1, + pad_value=0.0, +): + """Interpolate translation-like data (e.g., root_pos, contact_points) to target FPS. + + Args: + data: Tensor of shape (T, N, D) where T=frames, N=num_objects, D=dims (e.g., 3 for pos) + source_fps: Original frame rate + target_fps: Target frame rate + num_frames: Target number of frames + max_num_objects: Maximum number of objects to pad to + pad_value: Value to use for padding + + Returns: + Interpolated tensor of shape (num_frames, max_num_objects, D) + """ + from gear_sonic.trl.utils import math + + data = torch.tensor(data).float() if not torch.is_tensor(data) else data.float() + N_objects = data.shape[1] + D = data.shape[2] + + # Interpolate to target FPS if needed + if source_fps != target_fps: + # Reshape to (T, N_objects*D) for batch interpolation + data_flat = data.reshape(data.shape[0], -1) + data_interp = math.interpolate_pose( + data_flat, + source_fps=source_fps, + target_fps=target_fps, + device=data.device, + interpolation_type="linear", + ) + data = data_interp.reshape(-1, N_objects, D) + + # Trim or pad to match num_frames + if data.shape[0] > num_frames: + data = data[:num_frames] + elif data.shape[0] < num_frames: + padding = data[-1:].repeat(num_frames - data.shape[0], 1, 1) + data = torch.cat([data, padding], dim=0) + + # Pad or trim to max_num_objects + if data.shape[1] < max_num_objects: + padding = torch.full( + (data.shape[0], max_num_objects - data.shape[1], D), + pad_value, + dtype=data.dtype, + device=data.device, + ) + data = torch.cat([data, padding], dim=1) + else: + data = data[:, :max_num_objects] + + return data + + +def interpolate_quaternion_data( + data, + source_fps, + target_fps, + num_frames, + max_num_objects=1, +): + """Interpolate quaternion data (e.g., root_quat) to target FPS using slerp. + + Args: + data: Tensor of shape (T, N, 4) where T=frames, N=num_objects + source_fps: Original frame rate + target_fps: Target frame rate + num_frames: Target number of frames + max_num_objects: Maximum number of objects to pad to + + Returns: + Interpolated tensor of shape (num_frames, max_num_objects, 4) + """ + from gear_sonic.trl.utils import math + + data = torch.tensor(data).float() if not torch.is_tensor(data) else data.float() + N_objects = data.shape[1] + + # Interpolate to target FPS if needed + if source_fps != target_fps: + # Reshape to (T, N_objects*4) for batch interpolation + data_flat = data.reshape(data.shape[0], -1) + data_interp = math.interpolate_pose( + data_flat, + source_fps=source_fps, + target_fps=target_fps, + device=data.device, + interpolation_type="slerp", + rot_type="quat", + ) + data = data_interp.reshape(-1, N_objects, 4) + + # Trim or pad to match num_frames + if data.shape[0] > num_frames: + data = data[:num_frames] + elif data.shape[0] < num_frames: + padding = data[-1:].repeat(num_frames - data.shape[0], 1, 1) + data = torch.cat([data, padding], dim=0) + + # Pad or trim to max_num_objects + if data.shape[1] < max_num_objects: + padding = torch.zeros( + data.shape[0], + max_num_objects - data.shape[1], + 4, + dtype=data.dtype, + device=data.device, + ) + padding[:, :, 0] = 1.0 # w=1 for identity quaternion + data = torch.cat([data, padding], dim=1) + else: + data = data[:, :max_num_objects] + + return data + + +def interpolate_contact_center( + contact_points_dict, + source_fps, + target_fps, + num_frames, +): + """Compute contact center and in_contact label from raw contact points. + + Directly scales source frame indices to target frame space, avoiding + dense array interpolation that would blend real positions with zeros. + + Args: + contact_points_dict: Dict mapping frame_idx -> (N_points, 3) array + source_fps: Original frame rate + target_fps: Target frame rate + num_frames: Target number of frames + + Returns: + Tuple of: + contact_center: Tensor of shape (num_frames, 3) + in_contact: Tensor of shape (num_frames,) with binary labels + """ + if not contact_points_dict: + return torch.zeros(num_frames, 3), torch.zeros(num_frames) + + fps_ratio = target_fps / source_fps + contact_center = torch.zeros(num_frames, 3) + in_contact = torch.zeros(num_frames) + + for src_idx, points in contact_points_dict.items(): + if not (hasattr(points, "shape") and len(points) > 0): + continue + # Scale source frame range to target frame space + t_start = max(0, int(src_idx * fps_ratio)) + t_end = min(num_frames, int((src_idx + 1) * fps_ratio) + 1) + center = torch.from_numpy(points.mean(axis=0).astype(np.float32)) + contact_center[t_start:t_end] = center + in_contact[t_start:t_end] = 1.0 + + return contact_center, in_contact + + +class MotionLibBase: + def __init__(self, motion_lib_cfg, num_envs, device): + self.m_cfg = motion_lib_cfg + self.motion_fps_scale = self.m_cfg.get("motion_fps_scale", 1.0) + self._sim_fps = 1 / self.m_cfg.get("step_dt", 1 / 50) + self.target_fps = self.m_cfg.get("target_fps", 50) + self.adaptive_sampling_cfg = self.m_cfg.get("adaptive_sampling", {}) + self.all_motions_loaded = False + + self.debug = motion_lib_cfg.get("debug", False) + self.use_parallel_fk = motion_lib_cfg.get("use_parallel_fk", False) + self.num_envs = num_envs + self._device = device + self.mesh_parsers = None + self.has_action = False + skeleton_file = Path(self.m_cfg.asset.assetRoot) / self.m_cfg.asset.assetFileName + self.skeleton_tree = skeleton.SkeletonTree.from_mjcf(skeleton_file) + logger.info(f"Loaded skeleton from {skeleton_file}") + logger.info(f"Loading motion data from {self.m_cfg.motion_file}...") + self.load_data(self.m_cfg.motion_file) + self.use_adaptive_sampling = self.adaptive_sampling_cfg.get("enable", False) + if self.use_adaptive_sampling: + self.init_adaptive_sampling() + self.setup_constants( + fix_height=motion_lib_cfg.get("fix_height", FixHeightMode.no_fix), + multi_thread=self.m_cfg.get("multi_thread", True), + ) + + self.vid_smpl_pose = None + self.vid_smpl_joints = None + self.smpl_data = None + smpl_motion_file = motion_lib_cfg.get("smpl_motion_file", None) + self.smpl_data_keys = set() + if smpl_motion_file is not None: + if smpl_motion_file in ("dummy", "zeros"): + # Generate dummy zero SMPL data so SMPL observation terms work + # without needing to null them out in the config. + self.smpl_data = [None] * len(self._motion_data_keys) + elif osp.exists(smpl_motion_file): + if osp.isfile(smpl_motion_file): + self.smpl_data = joblib.load(smpl_motion_file) + self.smpl_data_keys = set(self.smpl_data.keys()) + self.smpl_data = [ + (self.smpl_data[k] if k in self.smpl_data else None) # noqa: SIM401 + for k in self._motion_data_keys + ] + else: + self.smpl_data = [] + smpl_pkl_files = set( + glob.glob(osp.join(smpl_motion_file, "**", "*.pkl"), recursive=True) + ) + for k in self._motion_data_keys: + seq = os.path.basename(k) + smpl_path = osp.join(smpl_motion_file, seq + ".pkl") + if self.debug or smpl_path in smpl_pkl_files: + self.smpl_data.append({"seq": seq, "path": smpl_path}) + self.smpl_data_keys.add(seq) + else: + self.smpl_data.append(None) + else: + self.smpl_data = [None] * len(self._motion_data_keys) + + self.smpl_y_up = motion_lib_cfg.get("smpl_y_up", False) + + # SOMA skeleton data loading (parallel to SMPL) + self.soma_data = None + soma_motion_file = motion_lib_cfg.get("soma_motion_file", None) + self.soma_data_keys = set() + self.soma_y_up = motion_lib_cfg.get("soma_y_up", True) # BVH is Y-up by default + self.num_soma_joints = motion_lib_cfg.get("num_soma_joints", 26) + if soma_motion_file is not None: + if soma_motion_file in ("dummy", "zeros"): + self.soma_data = [None] * len(self._motion_data_keys) + elif osp.exists(soma_motion_file): + if osp.isfile(soma_motion_file): + self.soma_data = joblib.load(soma_motion_file) + self.soma_data_keys = set(self.soma_data.keys()) + self.soma_data = [ + (self.soma_data[k] if k in self.soma_data else None) # noqa: SIM401 + for k in self._motion_data_keys + ] + else: + # Directory mode: per-motion PKL files (may be nested in subdirs) + soma_index = { + osp.splitext(osp.basename(f))[0]: f + for f in glob.glob( + osp.join(soma_motion_file, "**", "*.pkl"), recursive=True + ) + } + self.soma_data = [] + for k in self._motion_data_keys: + seq = os.path.basename(k) + soma_path = soma_index.get(seq) + if soma_path is not None or self.debug: + self.soma_data.append( + { + "seq": seq, + "path": soma_path or osp.join(soma_motion_file, seq + ".pkl"), + } + ) + self.soma_data_keys.add(seq) + else: + self.soma_data.append(None) + else: + self.soma_data = [None] * len(self._motion_data_keys) + + # Object data loading (similar to SMPL data) + self.object_data = None + object_motion_file = motion_lib_cfg.get("object_motion_file", None) + self.object_data_keys = set() + self.max_num_objects = motion_lib_cfg.get("max_num_objects", 1) + if object_motion_file is not None: + if osp.isfile(object_motion_file): + self.object_data = joblib.load(object_motion_file) + self.object_data_keys = set(self.object_data.keys()) + self.object_data = [ + (self.object_data[k] if k in self.object_data else None) # noqa: SIM401 + for k in self._motion_data_keys + ] + else: + self.object_data = [] + # TODO: osp.exists() can be very expensive, consider using a set of all object pkl files + # like in the smpl data loading above. + for k in self._motion_data_keys: + seq = os.path.basename(k) + object_path = osp.join(object_motion_file, seq + ".pkl") + if self.debug or osp.exists(object_path): + self.object_data.append({"seq": seq, "path": object_path}) + self.object_data_keys.add(seq) + else: + self.object_data.append(None) + # randomize the upper body poses condition + self.randomize_upper_body_poses = self.m_cfg.get("cat_upper_body_poses", False) + self.cat_upper_body_poses_prob = self.m_cfg.get("cat_upper_body_poses_prob", 0.0) + # The default prefixes for the upper body augmentation -- generated by the kinematic planner. + self.upper_body_augment_prefixes = self.m_cfg.get( + "upper_body_augment_prefixes", + ["2025", "walking_2025", "running_2025", "slow_walk_2025"], + ) + # Wrist joint noise augmentation config + self.randomize_wrist_poses = self.m_cfg.get("randomize_wrist_poses", False) + self.randomize_wrist_prob = self.m_cfg.get("randomize_wrist_prob", 0.3) + self.randomize_wrist_std = self.m_cfg.get("randomize_wrist_std", 0.1) + # MuJoCo DOF indices for wrist joints (L/R roll/pitch/yaw) + self.wrist_mujoco_dof_indices = [19, 20, 21, 26, 27, 28] + + def load_data(self, motion_file): + if osp.isfile(motion_file): + self.mode = MotionlibMode.file + self._motion_data_load = joblib.load(motion_file) + else: + assert osp.isdir( + motion_file + ), f"Expected motion_file to be a directory, got: {motion_file}" + self.mode = MotionlibMode.directory + if self.debug: + self._motion_data_load = {} + else: + self._motion_data_load = { + osp.splitext(osp.basename(f))[0]: {"path": f} + for f in glob.glob(osp.join(motion_file, "**", "*.pkl"), recursive=True) + if not f.endswith("metadata.pkl") + } + + metadata_files = [] + # Check for metadata.pkl directly in motion_file directory + direct_metadata = osp.join(motion_file, "metadata.pkl") + if osp.exists(direct_metadata): + metadata_files.append(direct_metadata) + # Also check subdirectories for metadata.pkl + all_sub_dirs = os.listdir(motion_file) + if self.debug: + all_sub_dirs = all_sub_dirs[:1] + for sub_dir in all_sub_dirs: + sub_dir_path = osp.join(motion_file, sub_dir) + if osp.isdir(sub_dir_path): + sub_meta = osp.join(sub_dir_path, "metadata.pkl") + if osp.exists(sub_meta): + metadata_files.append(sub_meta) + + for metadata_file in metadata_files: + metadata = joblib.load(metadata_file) + if self.debug: + metadata = { + k: v + for k, v in list(metadata.items())[:5] + if osp.exists(f"{sub_dir_path}/{k}.pkl") + } + for k, v in metadata.items(): + if self.debug: + self._motion_data_load[k] = {"path": f"{sub_dir_path}/{k}.pkl"} + if ( + k in self._motion_data_load + ): # metadata file can have more motion sequences than in the directory. Only load the necessary ones. # noqa: E501 + self._motion_data_load[k].update(v) + + print(f"Loaded {len(self._motion_data_load)} motion files") # noqa: T201 + + data_list = self._motion_data_load + + filter_motion_keys = self.m_cfg.get("filter_motion_keys", None) + if filter_motion_keys is not None: + if isinstance(filter_motion_keys, str): # noqa: SIM108 + patterns = [filter_motion_keys] + else: + patterns = list(filter_motion_keys) + + if all(pattern in data_list for pattern in patterns): + matched_keys = [pattern for pattern in patterns if pattern in data_list] + else: + compiled = [] + for pattern in patterns: + try: + compiled.append(re.compile(pattern)) + except re.error as exc: + raise ValueError(f"Invalid filter_motion_keys regex: {pattern}") from exc + + matched_keys = [ + k for k in data_list if any(regex.fullmatch(k) for regex in compiled) + ] + matched_keys.sort() + data_list = {k: data_list[k] for k in matched_keys} + + remove_motion_keys = self.m_cfg.get("remove_motion_keys", None) + if remove_motion_keys is not None: + # Remove any motion whose key starts with any of the remove_motion_keys prefixes + keys_to_remove = [ + k for k in data_list if any(k.startswith(prefix) for prefix in remove_motion_keys) + ] + for k in keys_to_remove: + del data_list[k] + + max_unique_motions = self.m_cfg.get("max_unique_motions", None) + if max_unique_motions is not None and len(data_list) > max_unique_motions: + import random + + keys = sorted(data_list.keys()) # Sort for determinism, then sample + selected = random.sample(keys, max_unique_motions) + data_list = {k: data_list[k] for k in selected} + print( # noqa: T201 + f"Limited to {max_unique_motions} random motions (from {len(keys)})" + ) # noqa: RUF100, T201 + + self._motion_data_list = np.array(list(data_list.values())) + self._motion_data_keys = np.array(list(data_list.keys())) + + # # HACK: Force specific motion only + # _FORCE_MOTION_KEY = "canned_food_31_jason_rigged_001_indoor2-v4_rand00063_000065" + # if _FORCE_MOTION_KEY in self._motion_data_keys: + # idx = list(self._motion_data_keys).index(_FORCE_MOTION_KEY) + # self._motion_data_list = np.array([self._motion_data_list[idx]]) + # self._motion_data_keys = np.array([_FORCE_MOTION_KEY]) + + self._num_unique_motions = len(self._motion_data_list) + logger.info(f"Loaded {self._num_unique_motions} motions") + + def _should_augment_upper_body(self, motion_key): + """Check if motion key matches any prefix for upper body augmentation""" # noqa: D415 + return any(motion_key.startswith(prefix) for prefix in self.upper_body_augment_prefixes) + + def setup_constants(self, fix_height=FixHeightMode.full_fix, multi_thread=True): + self.fix_height = fix_height + self.multi_thread = multi_thread + + #### Termination history + self._curr_motion_ids = None + self._termination_history = torch.zeros(self._num_unique_motions).to(self._device) + self._success_rate = torch.zeros(self._num_unique_motions).to(self._device) + self._sampling_history = torch.zeros(self._num_unique_motions).to(self._device) + self._sampling_prob = ( + torch.ones(self._num_unique_motions).to(self._device) / self._num_unique_motions + ) # For use in sampling batches + + def update_soft_sampling_weight(self, failed_keys): + # sampling weight based on evaluation, only "mostly" trained on "failed" sequences. Auto PMCP. + if len(failed_keys) > 0: + all_keys = self._motion_data_keys.tolist() + indexes = [all_keys.index(k) for k in failed_keys] + self._termination_history[indexes] += 1 + self.update_sampling_prob(self._termination_history) + + print( # noqa: T201 + "############################################################ Auto PMCP ############################################################" # noqa: E501 + ) + print( # noqa: T201 + f"Training mostly on {len(self._sampling_prob.cpu().nonzero())} seqs " + ) # noqa: RUF100, T201 + print( # noqa: T201 + self._motion_data_keys[self._sampling_prob.cpu().nonzero()].flatten() + ) # noqa: RUF100, T201 + print( # noqa: T201 + "###############################################################################################################################" + ) + else: + all_keys = self._motion_data_keys.tolist() + self._sampling_prob = ( + torch.ones(self._num_unique_motions).to(self._device) / self._num_unique_motions + ) # For use in sampling batches + + def update_sampling_prob(self, termination_history): + if ( + len(termination_history) == len(self._termination_history) + and termination_history.sum() > 0 + ): + self._sampling_prob[:] = termination_history / termination_history.sum() + if self._sampling_prob[self._curr_motion_ids].sum() == 0: + self._sampling_prob[self._curr_motion_ids] += 1e-6 + self._sampling_prob[:] = self._sampling_prob[:] / self._sampling_prob[:].sum() + self._sampling_batch_prob = ( + self._sampling_prob[self._curr_motion_ids] + / self._sampling_prob[self._curr_motion_ids].sum() + ) + self._termination_history = termination_history + return True + else: + return False + + def get_motion_actions(self, motion_ids, motion_times): + motion_len = self._motion_lengths[motion_ids] + num_frames = self._motion_num_frames[motion_ids] + dt = self._motion_dt[motion_ids] + # import ipdb; ipdb.set_trace() + frame_idx0, frame_idx1, blend = self._calc_frame_blend( + motion_times, motion_len, num_frames, dt + ) + f0l = frame_idx0 + self.length_starts[motion_ids] + f1l = frame_idx1 + self.length_starts[motion_ids] # noqa: F841 + + action = self._motion_actions[f0l] + return action + + def get_time_step_total(self, motion_ids): + return self._motion_num_frames[motion_ids] + + @property + def body_indexes(self): + return self.m_cfg.get("body_indexes_data", None) + + def get_dof_pos(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + + return self.dof_pos[motion_steps + length_starts] + + def get_dof_vel(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + return self.dof_vel[motion_steps + length_starts] + + def get_hand_dof_pos(self, motion_ids, motion_steps): + """Get hand DOF positions if available (for 43-DOF motion on 43-DOF robot).""" + if not hasattr(self, "hand_dof_pos") or self.hand_dof_pos is None: + return None + length_starts = self.length_starts[motion_ids] + return self.hand_dof_pos[motion_steps + length_starts] + + def get_body_pos_w(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + + return self.body_pos_w[motion_steps + length_starts] + + def get_body_quat_w(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + return self.body_quat_w[motion_steps + length_starts] + + def get_body_lin_vel_w(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + return self.body_lin_vel_w[motion_steps + length_starts] + + def get_body_ang_vel_w(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + return self.body_ang_vel_w[motion_steps + length_starts] + + # Full body data getters (all bodies, not sliced by body_indexes) + def get_body_pos_w_full(self, motion_ids, motion_steps): + """Get full body positions (all bodies, IsaacLab order).""" + length_starts = self.length_starts[motion_ids] + return self.body_pos_w_full[motion_steps + length_starts] + + def get_body_quat_w_full(self, motion_ids, motion_steps): + """Get full body quaternions (all bodies, IsaacLab order, wxyz).""" + length_starts = self.length_starts[motion_ids] + return self.body_quat_w_full[motion_steps + length_starts] + + def get_body_lin_vel_w_full(self, motion_ids, motion_steps): + """Get full body linear velocities (all bodies, IsaacLab order).""" + length_starts = self.length_starts[motion_ids] + return self.body_lin_vel_w_full[motion_steps + length_starts] + + def get_body_ang_vel_w_full(self, motion_ids, motion_steps): + """Get full body angular velocities (all bodies, IsaacLab order).""" + length_starts = self.length_starts[motion_ids] + return self.body_ang_vel_w_full[motion_steps + length_starts] + + def get_root_pos_w(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + return self.body_pos_w[motion_steps + length_starts, 0, :] + + def get_root_quat_w(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + return self.body_quat_w[motion_steps + length_starts, 0, :] + + def get_root_lin_vel_w(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + return self.body_lin_vel_w[motion_steps + length_starts, 0, :] + + def get_root_ang_vel_w(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + return self.body_ang_vel_w[motion_steps + length_starts, 0, :] + + def get_smpl_pose(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + return self._motion_smpl_poses[motion_steps + length_starts] + + def get_smpl_joints(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + return self._motion_smpl_joints[motion_steps + length_starts] + + def get_smpl_transl(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + return self._motion_smpl_transl[motion_steps + length_starts] + + @staticmethod + def _resample_soma_tensor(data, fps_source, fps_target): + """Resample a SOMA tensor along dim 0 matching interploate_pose frame count. + + Uses the same arange(0, duration, 1/fps_target) formula as + torch_humanoid_batch.interploate_pose so robot and SOMA frame counts align. + """ + n_src = data.shape[0] + duration = (n_src - 1) / fps_source + tgt_times = torch.arange(0, duration, 1 / fps_target, dtype=torch.float32) + n_tgt = len(tgt_times) + if n_tgt <= 1: + return data[:1] + # Compute blend weights (same logic as _compute_frame_blend) + phase = tgt_times / duration + idx0 = (phase * (n_src - 1)).floor().long() + idx1 = torch.minimum(idx0 + 1, torch.tensor(n_src - 1)) + blend = (phase * (n_src - 1) - idx0).float() + # Reshape blend for broadcasting with arbitrary trailing dims + for _ in range(data.dim() - 1): + blend = blend.unsqueeze(-1) + return data[idx0] * (1 - blend) + data[idx1] * blend + + def get_soma_joints(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + return self._motion_soma_joints[motion_steps + length_starts] + + def get_soma_root_quat(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + return self._motion_soma_root_quat[motion_steps + length_starts] + + def get_soma_transl(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + return self._motion_soma_transl[motion_steps + length_starts] + + def get_object_root_pos(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + return self._motion_object_root_pos[motion_steps + length_starts] + + def get_object_root_quat(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + return self._motion_object_root_quat[motion_steps + length_starts] + + def get_object_lin_vel(self, motion_ids, motion_steps): + """Get object linear velocity from motion library.""" + length_starts = self.length_starts[motion_ids] + return self._motion_object_lin_vel[motion_steps + length_starts] + + def get_object_ang_vel(self, motion_ids, motion_steps): + """Get object angular velocity from motion library.""" + length_starts = self.length_starts[motion_ids] + return self._motion_object_ang_vel[motion_steps + length_starts] + + def get_object_contact_center(self, motion_ids, motion_steps, hand="right_hand"): + """Get object contact center from motion library. + + Contact center is the mean of all contact points per frame for the given hand. + + Args: + motion_ids: (N,) tensor of motion indices + motion_steps: (N,) tensor of frame indices within each motion + hand: Which hand's contact center to return ("left_hand" or "right_hand") + + Returns: + Tensor of shape (N, 3) with contact center positions in object-local frame, + or None if not available. + """ + attr = f"_motion_object_contact_center_{'left' if hand == 'left_hand' else 'right'}" + if not hasattr(self, attr): + return None + length_starts = self.length_starts[motion_ids] + return getattr(self, attr)[motion_steps + length_starts] + + def get_object_in_contact(self, motion_ids, motion_steps, hand="right_hand"): + """Get binary in_contact label for the given hand. + + Args: + motion_ids: (N,) tensor of motion indices + motion_steps: (N,) tensor of frame indices within each motion + hand: Which hand ("left_hand" or "right_hand") + + Returns: + Tensor of shape (N,) with 1.0 if in contact, 0.0 otherwise, + or None if not available. + """ + attr = f"_motion_object_in_contact_{'left' if hand == 'left_hand' else 'right'}" + if not hasattr(self, attr): + return None + length_starts = self.length_starts[motion_ids] + return getattr(self, attr)[motion_steps + length_starts] + + def get_hand_action(self, motion_ids, motion_steps, hand="right_hand"): + """Get discrete hand action (open/closed) for the given hand. + + Args: + motion_ids: (N,) tensor of motion indices + motion_steps: (N,) tensor of frame indices within each motion + hand: Which hand ("left_hand" or "right_hand") + + Returns: + Tensor of shape (N,) with -1.0 = open, +1.0 = closed, + or None if not available. + """ + attr = f"_motion_hand_action_{'left' if hand == 'left_hand' else 'right'}" + if not hasattr(self, attr): + return None + length_starts = self.length_starts[motion_ids] + return getattr(self, attr)[motion_steps + length_starts] + + def get_feet_l(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + return self.feet_l[motion_steps + length_starts] + + def get_feet_r(self, motion_ids, motion_steps): + length_starts = self.length_starts[motion_ids] + return self.feet_r[motion_steps + length_starts] + + def get_motion_state(self, motion_ids, motion_times, offset=None): + motion_len = self._motion_lengths[motion_ids] + num_frames = self._motion_num_frames[motion_ids] + dt = self._motion_dt[motion_ids] + + frame_idx0, frame_idx1, blend = self._calc_frame_blend( + motion_times, motion_len, num_frames, dt + ) + f0l = frame_idx0 + self.length_starts[motion_ids] + f1l = frame_idx1 + self.length_starts[motion_ids] + + if "dof_pos" in self.__dict__: + local_rot0 = self.dof_pos[f0l] + local_rot1 = self.dof_pos[f1l] + else: + local_rot0 = self.body_pos_b[f0l] + local_rot1 = self.body_pos_b[f1l] + + body_lin_vel_w0 = self.body_lin_vel_w[f0l] + body_lin_vel_w1 = self.body_lin_vel_w[f1l] + + body_ang_vel0 = self.body_ang_vel_w[f0l] + body_ang_vel1 = self.body_ang_vel_w[f1l] + + body_pos_w0 = self.body_pos_w[f0l, :] + body_pos_w1 = self.body_pos_w[f1l, :] + + dof_vel0 = self.dof_vel[f0l] + dof_vel1 = self.dof_vel[f1l] + + vals = [ + local_rot0, + local_rot1, + body_lin_vel_w0, + body_lin_vel_w1, + body_ang_vel0, + body_ang_vel1, + body_pos_w0, + body_pos_w1, + dof_vel0, + dof_vel1, + ] + for v in vals: + assert v.dtype != torch.float64 + + blend = blend.unsqueeze(-1) + + blend_exp = blend.unsqueeze(-1) + + if offset is None: + body_pos_w = ( + 1.0 - blend_exp + ) * body_pos_w0 + blend_exp * body_pos_w1 # ZL: apply offset + else: + body_pos_w = ( + (1.0 - blend_exp) * body_pos_w0 + blend_exp * body_pos_w1 + offset[..., None, :] + ) # ZL: apply offset + + body_lin_vel_w = (1.0 - blend_exp) * body_lin_vel_w0 + blend_exp * body_lin_vel_w1 + body_ang_vel_w = (1.0 - blend_exp) * body_ang_vel0 + blend_exp * body_ang_vel1 + + if "dof_pos" in self.__dict__: # Robot Joints + dof_vel = (1.0 - blend) * dof_vel0 + blend * dof_vel1 + dof_pos = (1.0 - blend) * local_rot0 + blend * local_rot1 + else: + dof_vel = (1.0 - blend_exp) * dof_vel0 + blend_exp * dof_vel1 + local_rot = rotations.slerp(local_rot0, local_rot1, torch.unsqueeze(blend, axis=-1)) + dof_pos = self._local_rotation_to_dof_smpl(local_rot) + + body_quat_w0 = self.body_quat_w[f0l] + body_quat_w1 = self.body_quat_w[f1l] + body_quat_w = rotations.slerp(body_quat_w0, body_quat_w1, blend_exp) + return_dict = {} + + if "gts_t" in self.__dict__: + body_pos_w_t0 = self.body_pos_t_w[f0l] + body_pos_w_t1 = self.body_pos_t_w[f1l] + + body_quat_t0 = self.body_quat_t_w[f0l] + body_quat_t1 = self.body_quat_t_w[f1l] + + body_lin_vel_w_t0 = self.body_lin_vel_t_w[f0l] + body_lin_vel_w_t1 = self.body_lin_vel_t_w[f1l] + + body_ang_vel_t0 = self.body_ang_vel_t_w[f0l] + body_ang_vel_t1 = self.body_ang_vel_t_w[f1l] + if offset is None: + body_pos_t_w = (1.0 - blend_exp) * body_pos_w_t0 + blend_exp * body_pos_w_t1 + else: + body_pos_t_w = ( + (1.0 - blend_exp) * body_pos_w_t0 + + blend_exp * body_pos_w_t1 + + offset[..., None, :] + ) + body_quat_t_w = rotations.slerp(body_quat_t0, body_quat_t1, blend_exp) + body_lin_vel_t_w = (1.0 - blend_exp) * body_lin_vel_w_t0 + blend_exp * body_lin_vel_w_t1 + body_ang_vel_t_w = (1.0 - blend_exp) * body_ang_vel_t0 + blend_exp * body_ang_vel_t1 + else: + body_pos_t_w = body_pos_w + body_quat_t_w = body_quat_w + body_lin_vel_t_w = body_lin_vel_w + body_ang_vel_t_w = body_ang_vel_w + + if self.smpl_data is not None: + smpl_pose0 = self._motion_smpl_poses[f0l] + smpl_pose1 = self._motion_smpl_poses[f1l] + smpl_pose = (1.0 - blend) * smpl_pose0 + blend * smpl_pose1 + return_dict.update({"smpl_pose": smpl_pose.clone()}) + + if hasattr(self, "_motion_smpl_joints"): + smpl_joints0 = self._motion_smpl_joints[f0l] + smpl_joints1 = self._motion_smpl_joints[f1l] + smpl_joints = (1.0 - blend_exp) * smpl_joints0 + blend_exp * smpl_joints1 + return_dict.update({"smpl_joints": smpl_joints.clone()}) + + if hasattr(self, "_motion_smpl_transl"): + smpl_transl0 = self._motion_smpl_transl[f0l] + smpl_transl1 = self._motion_smpl_transl[f1l] + smpl_transl = (1.0 - blend_exp) * smpl_transl0 + blend_exp * smpl_transl1 + return_dict.update({"smpl_transl": smpl_transl.clone()}) + + if self.soma_data is not None: + if hasattr(self, "_motion_soma_joints"): + soma_joints0 = self._motion_soma_joints[f0l] + soma_joints1 = self._motion_soma_joints[f1l] + soma_joints = (1.0 - blend_exp) * soma_joints0 + blend_exp * soma_joints1 + return_dict.update({"soma_joints": soma_joints.clone()}) + + if hasattr(self, "_motion_soma_root_quat"): + # For quaternions, use slerp (approximate with linear blend + normalize) + soma_rq0 = self._motion_soma_root_quat[f0l] + soma_rq1 = self._motion_soma_root_quat[f1l] + soma_root_quat = (1.0 - blend_exp) * soma_rq0 + blend_exp * soma_rq1 + soma_root_quat = soma_root_quat / (soma_root_quat.norm(dim=-1, keepdim=True) + 1e-8) + return_dict.update({"soma_root_quat": soma_root_quat.clone()}) + + if hasattr(self, "_motion_soma_transl"): + soma_transl0 = self._motion_soma_transl[f0l] + soma_transl1 = self._motion_soma_transl[f1l] + soma_transl = (1.0 - blend_exp) * soma_transl0 + blend_exp * soma_transl1 + return_dict.update({"soma_transl": soma_transl.clone()}) + + if self.object_data is not None: + if hasattr(self, "_motion_object_root_pos"): + object_root_pos0 = self._motion_object_root_pos[f0l] + object_root_pos1 = self._motion_object_root_pos[f1l] + object_root_pos = ( + 1.0 - blend_exp + ) * object_root_pos0 + blend_exp * object_root_pos1 + return_dict.update({"object_root_pos": object_root_pos.clone()}) + + if hasattr(self, "_motion_object_root_quat"): + object_root_quat0 = self._motion_object_root_quat[f0l] + object_root_quat1 = self._motion_object_root_quat[f1l] + # Use slerp for quaternion interpolation + object_root_quat = rotations.slerp(object_root_quat0, object_root_quat1, blend_exp) + return_dict.update({"object_root_quat": object_root_quat.clone()}) + + return_dict.update( + { + "root_pos": body_pos_w[..., 0, :].clone(), + "root_rot": body_quat_w[..., 0, :].clone(), + "dof_pos": dof_pos.clone(), + "root_vel": body_lin_vel_w[..., 0, :].clone(), + "root_ang_vel": body_ang_vel_w[..., 0, :].clone(), + "dof_vel": dof_vel.clone(), + "motion_aa": self._motion_aa[f0l].clone(), + "motion_bodies": self._motion_bodies[motion_ids].clone(), + "body_pos_w": body_pos_w.clone(), + "body_quat_w": body_quat_w.clone(), + "body_lin_vel_w": body_lin_vel_w.clone(), + "body_ang_vel_w": body_ang_vel_w.clone(), + "body_pos_w_t": body_pos_t_w.clone(), + "body_quat_t": body_quat_t_w.clone(), + "body_lin_vel_w_t": body_lin_vel_t_w.clone(), + "body_ang_vel_t": body_ang_vel_t_w.clone(), + } + ) + if "feet_l" in self.__dict__: + blend_int = blend.round().int() + feet_l = torch.where(blend_int == 0, self.feet_l[f0l], self.feet_l[f1l]) + feet_r = torch.where(blend_int == 0, self.feet_r[f0l], self.feet_r[f1l]) + return_dict.update( + { + "feet_l": feet_l.clone().bool(), + "feet_r": feet_r.clone().bool(), + } + ) + return return_dict + + def load_all_motions(self): + self.all_motions_loaded = True + self.load_motions(random_sample=False, num_motions_to_load=self._num_unique_motions) + + def load_motions_for_training(self, max_num_seqs=None): + if self.all_motions_loaded: + print("All motions already loaded!!! No need to resample.") # noqa: T201 + return False + + if self.m_cfg.get("override_num_motions_to_load", None) is not None: + max_num_seqs = self.m_cfg.override_num_motions_to_load + + # Option to load unique motions (no duplicates) - useful for replay/evaluation + load_unique = self.m_cfg.get("load_unique_motions", False) + + if ( + max_num_seqs is None + ): # if not specified, load all motions, can OOM if the dataset is too large. + max_num_seqs = self._num_unique_motions + self.all_motions_loaded = True + self.load_motions(random_sample=False, num_motions_to_load=self._num_unique_motions) + elif ( + max_num_seqs >= self._num_unique_motions + ): # if specified but more than the number of unique motions, load all motions as well. + self.all_motions_loaded = True + self.load_motions(random_sample=False, num_motions_to_load=self._num_unique_motions) + else: # if there are more motions than specified, then randomly sample the requested number of motions. + self.all_motions_loaded = False + # Use random_sample=False when load_unique=True to avoid duplicates + self.load_motions(random_sample=not load_unique, num_motions_to_load=max_num_seqs) + if load_unique: + print( # noqa: T201 + f"[MotionLib] Loaded {max_num_seqs} unique motions (no duplicates)" + ) # noqa: RUF100, T201 + return True + + def load_motions_for_evaluation(self, start_idx=0): + # disable this check to avoid upper body poses randomization in evaluation + # if self.all_motions_loaded: + # print("All motions already loaded!!! No need to resample.") + # return + + if ( + self._num_unique_motions > self.num_envs + ): # if number of motions is more than number of envs, then we should only partially load the motions. + self.all_motions_loaded = False + self.load_motions( + random_sample=False, + num_motions_to_load=self.num_envs, + start_idx=start_idx, + is_evaluation=True, + ) + else: + self.all_motions_loaded = True + self.load_motions( + random_sample=False, + num_motions_to_load=self._num_unique_motions, + start_idx=start_idx, + is_evaluation=True, + ) + + def load_motions( + self, + random_sample=True, + start_idx=0, + max_len=-1, + target_heading=None, + num_motions_to_load=None, + is_evaluation=False, + ): + + if "gts" in self.__dict__: + del ( + self.body_pos_w, + self.body_quat_w, + self.body_pos_b, + self.root_linv_vel_w, + self.root_ang_vel_w, + self.body_ang_vel_w, + self.body_lin_vel_w, + self.dof_vels, + self.dof_pos, + ) + if "gts_t" in self.__dict__: + del ( + self.body_pos_t_w, + self.body_quat_t_w, + self.body_lin_vel_t_w, + self.body_ang_vel_t_w, + ) + + motions = [] + _motion_lengths = [] + _motion_fps = [] + _motion_dt = [] + _motion_num_frames = [] + _motion_bodies = [] + _motion_aa = [] + has_action = False # noqa: F841 + _motion_actions = [] + _motion_smpl_poses = [] + _motion_smpl_joints = [] + _motion_smpl_transl = [] + _motion_soma_joints = [] + _motion_soma_root_quat = [] + _motion_soma_transl = [] + _motion_object_root_pos = [] + _motion_object_root_quat = [] + _motion_object_contact_center_left = [] + _motion_object_contact_center_right = [] + _motion_object_in_contact_left = [] + _motion_object_in_contact_right = [] + _motion_hand_action_left = [] + _motion_hand_action_right = [] + + total_len = 0.0 + self.num_joints = len(self.skeleton_tree.node_names) + if num_motions_to_load is None: # noqa: SIM108 + num_motion_to_load = self.num_envs + else: + num_motion_to_load = num_motions_to_load + + if self.use_adaptive_sampling: + self.update_adaptive_sampling_motion_sequences() + + if random_sample: + sample_idxes = torch.multinomial( + self._sampling_prob, num_samples=num_motion_to_load, replacement=True + ).to(self._device) + else: # start_idx only used for non-random sampling. + sample_idxes = torch.clamp( + torch.arange(num_motion_to_load) + start_idx, max=self._num_unique_motions - 1 + ).to(self._device) + + # sample_idxes = torch.tensor([self._motion_data_keys.tolist().index("0-KIT_8_WalkInClockwiseCircle04_poses")]).to(self._device) # noqa: E501 + self._curr_motion_ids = sample_idxes + self.curr_motion_keys = ( + [self._motion_data_keys[sample_idxes.cpu()]] + if sample_idxes.numel() == 1 + else self._motion_data_keys[sample_idxes.cpu()].tolist() + ) + self._sampling_batch_prob = ( + self._sampling_prob[self._curr_motion_ids] + / self._sampling_prob[self._curr_motion_ids].sum() + ) + + logger.info(f"Loading {num_motion_to_load} motions...") + logger.info(f"Sampling motion: {sample_idxes[:10]}, ....") + logger.info(f"Current motion keys: {self.curr_motion_keys[:10]}, ....") + + motion_data_list = self._motion_data_list[sample_idxes.cpu().numpy()] + if self.smpl_data is not None: + smpl_data_list = [self.smpl_data[idx] for idx in sample_idxes.cpu().numpy()] + else: + smpl_data_list = None + if self.object_data is not None: + object_data_list = [self.object_data[idx] for idx in sample_idxes.cpu().numpy()] + else: + object_data_list = None + if self.soma_data is not None: + soma_data_list = [self.soma_data[idx] for idx in sample_idxes.cpu().numpy()] + else: + soma_data_list = None + torch.set_num_threads(1) + + # Increase file descriptor limit to prevent "too many open files" error + try: + soft_limit, hard_limit = resource.getrlimit(resource.RLIMIT_NOFILE) + target_limit = 1048576 + + # Try to set both soft and hard limits + if soft_limit < target_limit: + try: + # First try to increase hard limit (requires root) + resource.setrlimit(resource.RLIMIT_NOFILE, (target_limit, target_limit)) + logger.info( + f"Increased file descriptor limits from {soft_limit}/{hard_limit} to {target_limit}/{target_limit}" # noqa: E501 + ) + except PermissionError: + # Fallback to increasing only soft limit up to hard limit + new_soft = min(target_limit, hard_limit) + resource.setrlimit(resource.RLIMIT_NOFILE, (new_soft, hard_limit)) + logger.info( + f"Increased soft file descriptor limit from {soft_limit} to {new_soft} (hard limit: {hard_limit})" # noqa: E501 + ) + except Exception as e: # noqa: BLE001 + logger.warning(f"Could not increase file descriptor limit: {e}") + + manager = mp.Manager() + queue = manager.Queue() + num_jobs = min(min(mp.cpu_count(), 32), len(motion_data_list)) # noqa: PLW3301 + + if num_jobs <= 8 or not self.multi_thread or len(motion_data_list) <= 128: + num_jobs = 1 + + logger.info(f"Loading motions with {num_jobs} jobs...") + self.res_non_nav_dataset = {} + res_acc = {} # using dictionary ensures order of the results. + workers = [] + + # if self.randomize_upper_body_poses: + # self.cat_upper_body_poses_prob = 1.0 + if self.randomize_upper_body_poses and not is_evaluation: + + # get indices that are in navigation dataset + nav_indices = [ + i + for i in range(len(motion_data_list)) + if self._should_augment_upper_body(self.curr_motion_keys[i]) + ] + other_indices = [ + i + for i in range(len(motion_data_list)) + if not self._should_augment_upper_body(self.curr_motion_keys[i]) + ] + + nav_motion_data_list = [motion_data_list[i] for i in nav_indices] + other_motion_data_list = [motion_data_list[i] for i in other_indices] + + if self.smpl_data is not None: + nav_smpl_data_list = [smpl_data_list[i] for i in nav_indices] + other_smpl_data_list = [smpl_data_list[i] for i in other_indices] + else: + nav_smpl_data_list = None + other_smpl_data_list = None + + if self.object_data is not None: + nav_object_data_list = [object_data_list[i] for i in nav_indices] + other_object_data_list = [object_data_list[i] for i in other_indices] + else: + nav_object_data_list = None + other_object_data_list = None + + if soma_data_list is not None: + nav_soma_data_list = [soma_data_list[i] for i in nav_indices] + other_soma_data_list = [soma_data_list[i] for i in other_indices] + else: + nav_soma_data_list = None + other_soma_data_list = None + + # load non-navigation dataset first + if len(other_motion_data_list) > 0: + jobs = other_motion_data_list + chunk = np.ceil(len(jobs) / num_jobs).astype(int) + ids = np.array(other_indices) # Use original indices, not sequential + + jobs = [ + ( + ids[i : i + chunk], + jobs[i : i + chunk], + ( + None + if other_smpl_data_list is None + else other_smpl_data_list[i : i + chunk] + ), + ( + None + if other_object_data_list is None + else other_object_data_list[i : i + chunk] + ), + ( + None + if other_soma_data_list is None + else other_soma_data_list[i : i + chunk] + ), + self.fix_height, + target_heading, + max_len, + is_evaluation, + ) + for i in range(0, len(jobs), chunk) + ] + + job_args = [jobs[i] for i in range(len(jobs))] + for i in range(1, len(jobs)): + worker_args = (*job_args[i], queue, i) + worker = mp.Process(target=self.load_motion_with_skeleton, args=worker_args) + worker.start() + workers.append(worker) + res_acc.update(self.load_motion_with_skeleton(*jobs[0], None, 0)) + + # Wait for all workers to complete and clean them up + for worker in workers: + worker.join() + worker.close() + workers = [] + + for i in progress.track( # noqa: B007 + range(len(jobs) - 1), "Gathering results for non-navigation dataset..." + ): + res = queue.get() + res_acc.update(res) + + self.res_non_nav_dataset = res_acc.copy() + + # load navigation dataset + if len(nav_motion_data_list) > 0: + + jobs = nav_motion_data_list + chunk = np.ceil(len(jobs) / num_jobs).astype(int) + ids = np.array(nav_indices) # Use original indices, not sequential + + jobs = [ + ( + ids[i : i + chunk], + jobs[i : i + chunk], + nav_smpl_data_list[ + i : i + chunk + ], # navigation dataset would never have smpl data. This is always empty. + ( + None + if nav_object_data_list is None + else nav_object_data_list[i : i + chunk] + ), + (None if nav_soma_data_list is None else nav_soma_data_list[i : i + chunk]), + self.fix_height, + target_heading, + max_len, + is_evaluation, + ) + for i in range(0, len(jobs), chunk) + ] + job_args = [jobs[i] for i in range(len(jobs))] + for i in range(1, len(jobs)): + worker_args = (*job_args[i], queue, i) + worker = mp.Process(target=self.load_motion_with_skeleton, args=worker_args) + worker.start() + workers.append(worker) + res_acc.update(self.load_motion_with_skeleton(*jobs[0], None, 0)) + + for i in progress.track( # noqa: B007 + range(len(jobs) - 1), "Gathering results for navigation dataset..." + ): + res = queue.get() + res_acc.update(res) + + # Wait for all workers to complete and clean them up + for worker in workers: + worker.join() + worker.close() + workers = [] + + else: + jobs = motion_data_list + chunk = np.ceil(len(jobs) / num_jobs).astype(int) + ids = np.arange(len(jobs)) + + jobs = [ + ( + ids[i : i + chunk], + jobs[i : i + chunk], + None if smpl_data_list is None else smpl_data_list[i : i + chunk], + None if object_data_list is None else object_data_list[i : i + chunk], + None if soma_data_list is None else soma_data_list[i : i + chunk], + self.fix_height, + target_heading, + max_len, + is_evaluation, + ) + for i in range(0, len(jobs), chunk) + ] + + job_args = [jobs[i] for i in range(len(jobs))] + for i in range(1, len(jobs)): + worker_args = (*job_args[i], queue, i) + worker = mp.Process(target=self.load_motion_with_skeleton, args=worker_args) + worker.start() + workers.append(worker) + res_acc.update(self.load_motion_with_skeleton(*jobs[0], None, 0)) + + for i in progress.track(range(len(jobs) - 1), "Gathering results..."): # noqa: B007 + res = queue.get() + res_acc.update(res) + + nav_indices = [] + other_indices = list(range(len(motions))) + + # Wait for all workers to complete and clean them up + for worker in workers: + worker.join() + worker.close() + workers = [] + + for f in progress.track(range(len(res_acc)), description="Processing motions..."): + motion_file_data, curr_motion = res_acc[f] + motion_fps = int(curr_motion.fps * self.motion_fps_scale) + curr_dt = 1.0 / motion_fps + num_frames = curr_motion.global_rotation.shape[0] + + curr_len = 1.0 / motion_fps * (num_frames - 1) + + if "beta" in motion_file_data: + _motion_aa.append(motion_file_data["pose_aa"].reshape(-1, self.num_joints * 3)) + _motion_bodies.append(curr_motion.gender_beta) + else: + _motion_aa.append(np.zeros((num_frames, self.num_joints * 3))) + _motion_bodies.append(torch.zeros(17)) + + _motion_fps.append(motion_fps) + _motion_dt.append(curr_dt) + _motion_num_frames.append(num_frames) + motions.append(curr_motion) + _motion_lengths.append(curr_len) + if self.has_action: + _motion_actions.append(curr_motion.action) + if self.smpl_data is not None: + _motion_smpl_poses.append(curr_motion["smpl_pose"]) + if "smpl_joints" in curr_motion: + _motion_smpl_joints.append(curr_motion["smpl_joints"]) + if "smpl_transl" in curr_motion: + _motion_smpl_transl.append(curr_motion["smpl_transl"]) + if self.soma_data is not None: + if "soma_joints" in curr_motion: + _motion_soma_joints.append(curr_motion["soma_joints"]) + if "soma_root_quat" in curr_motion: + _motion_soma_root_quat.append(curr_motion["soma_root_quat"]) + if "soma_transl" in curr_motion: + _motion_soma_transl.append(curr_motion["soma_transl"]) + if self.object_data is not None: + if "object_root_pos" in curr_motion: + _motion_object_root_pos.append(curr_motion["object_root_pos"]) + if "object_root_quat" in curr_motion: + _motion_object_root_quat.append(curr_motion["object_root_quat"]) + if "object_contact_center_left" in curr_motion: + _motion_object_contact_center_left.append( + curr_motion["object_contact_center_left"] + ) + if "object_in_contact_left" in curr_motion: + _motion_object_in_contact_left.append(curr_motion["object_in_contact_left"]) + if "object_contact_center_right" in curr_motion: + _motion_object_contact_center_right.append( + curr_motion["object_contact_center_right"] + ) + if "object_in_contact_right" in curr_motion: + _motion_object_in_contact_right.append(curr_motion["object_in_contact_right"]) + if "hand_action_left" in motion_file_data: + raw_action = motion_file_data["hand_action_left"] + # Nearest-neighbor interpolation to match target fps + src_len = len(raw_action) + if src_len != num_frames: + indices = np.round(np.linspace(0, src_len - 1, num_frames)).astype(int) + raw_action = raw_action[indices] + _motion_hand_action_left.append(raw_action) + if "hand_action_right" in motion_file_data: + raw_action = motion_file_data["hand_action_right"] + # Nearest-neighbor interpolation to match target fps + src_len = len(raw_action) + if src_len != num_frames: + indices = np.round(np.linspace(0, src_len - 1, num_frames)).astype(int) + raw_action = raw_action[indices] + _motion_hand_action_right.append(raw_action) + del curr_motion + + self._motion_lengths = torch.tensor( + _motion_lengths, device=self._device, dtype=torch.float32 + ) + self._motion_fps = torch.tensor(_motion_fps, device=self._device, dtype=torch.float32) + self._motion_bodies = torch.stack(_motion_bodies).to(self._device).type(torch.float32) + self._motion_aa = torch.tensor( + np.concatenate(_motion_aa), device=self._device, dtype=torch.float32 + ) + + if self.smpl_data is not None: + self._motion_smpl_poses = torch.cat(_motion_smpl_poses, dim=0).float().to(self._device) + if len(_motion_smpl_joints) > 0: + self._motion_smpl_joints = ( + torch.cat(_motion_smpl_joints, dim=0).float().to(self._device) + ) + if len(_motion_smpl_transl) > 0: + self._motion_smpl_transl = ( + torch.cat(_motion_smpl_transl, dim=0).float().to(self._device) + ) + if self.soma_data is not None: + if len(_motion_soma_joints) > 0: + self._motion_soma_joints = ( + torch.cat(_motion_soma_joints, dim=0).float().to(self._device) + ) + if len(_motion_soma_root_quat) > 0: + self._motion_soma_root_quat = ( + torch.cat(_motion_soma_root_quat, dim=0).float().to(self._device) + ) + if len(_motion_soma_transl) > 0: + self._motion_soma_transl = ( + torch.cat(_motion_soma_transl, dim=0).float().to(self._device) + ) + if self.object_data is not None: + if len(_motion_object_root_pos) > 0: + self._motion_object_root_pos = ( + torch.cat(_motion_object_root_pos, dim=0).float().to(self._device) + ) + if len(_motion_object_root_quat) > 0: + self._motion_object_root_quat = ( + torch.cat(_motion_object_root_quat, dim=0).float().to(self._device) + ) + # Store per-hand contact centers and in_contact labels + if len(_motion_object_contact_center_left) > 0: + self._motion_object_contact_center_left = ( + torch.cat(_motion_object_contact_center_left, dim=0).float().to(self._device) + ) + if len(_motion_object_in_contact_left) > 0: + self._motion_object_in_contact_left = ( + torch.cat(_motion_object_in_contact_left, dim=0).float().to(self._device) + ) + if len(_motion_object_contact_center_right) > 0: + self._motion_object_contact_center_right = ( + torch.cat(_motion_object_contact_center_right, dim=0).float().to(self._device) + ) + if len(_motion_object_in_contact_right) > 0: + self._motion_object_in_contact_right = ( + torch.cat(_motion_object_in_contact_right, dim=0).float().to(self._device) + ) + if len(_motion_hand_action_left) > 0: + self._motion_hand_action_left = ( + torch.from_numpy(np.concatenate(_motion_hand_action_left, axis=0)) + .float() + .to(self._device) + ) + if len(_motion_hand_action_right) > 0: + self._motion_hand_action_right = ( + torch.from_numpy(np.concatenate(_motion_hand_action_right, axis=0)) + .float() + .to(self._device) + ) + self._motion_dt = torch.tensor(_motion_dt, device=self._device, dtype=torch.float32) + + # Compute object velocities from position/quaternion using finite differences + if self.object_data is not None and hasattr(self, "_motion_object_root_pos"): + self._compute_object_velocities(_motion_num_frames, _motion_dt) + self._motion_num_frames = torch.tensor(_motion_num_frames, device=self._device) + + if self.has_action: + self._motion_actions = torch.cat(_motion_actions, dim=0).float().to(self._device) + self._num_motions = len(motions) + + self.body_pos_w = ( + torch.cat([m.global_translation for m in motions], dim=0).float().to(self._device) + ) + self.body_quat_w = ( + torch.cat([m.global_rotation for m in motions], dim=0).float().to(self._device) + ) + self.body_pos_b = ( + torch.cat([m.local_rotation for m in motions], dim=0).float().to(self._device) + ) + self.root_linv_vel_w = ( + torch.cat([m.global_root_velocity for m in motions], dim=0).float().to(self._device) + ) + self.root_ang_vel_w = ( + torch.cat([m.global_root_angular_velocity for m in motions], dim=0) + .float() + .to(self._device) + ) + self.body_ang_vel_w = ( + torch.cat([m.global_angular_velocity for m in motions], dim=0).float().to(self._device) + ) + self.body_lin_vel_w = ( + torch.cat([m.global_velocity for m in motions], dim=0).float().to(self._device) + ) + self.dof_vel = torch.cat([m.dof_vels for m in motions], dim=0).float().to(self._device) + self.feet_l = torch.cat([m.feet_l for m in motions], dim=0).float().to(self._device) + self.feet_r = torch.cat([m.feet_r for m in motions], dim=0).float().to(self._device) + + # if "global_translation_extend" in motions[0].__dict__: + # self.body_pos_t_w = torch.cat([m.global_translation_extend for m in motions], dim=0).float().to(self._device) # noqa: E501 + # self.body_quat_t_w = torch.cat([m.global_rotation_extend for m in motions], dim=0).float().to(self._device) # noqa: E501 + # self.body_lin_vel_t_w = torch.cat([m.global_velocity_extend for m in motions], dim=0).float().to(self._device) # noqa: E501 + # self.body_ang_vel_t_w = torch.cat([m.global_angular_velocity_extend for m in motions], dim=0).float().to(self._device) # noqa: E501 + # self.feet_l = torch.cat([m.feet_l for m in motions], dim=0).float().to(self._device) + # self.feet_r = torch.cat([m.feet_r for m in motions], dim=0).float().to(self._device) + + if "dof_pos" in motions[0].__dict__: + self.dof_pos = torch.cat([m.dof_pos for m in motions], dim=0).float().to(self._device) + + # Store hand DOF positions if available (for 43-DOF motion) + if "hand_dof_pos" in motions[0].__dict__: + self.hand_dof_pos = ( + torch.cat([m.hand_dof_pos for m in motions], dim=0).float().to(self._device) + ) + else: + self.hand_dof_pos = None + + lengths = self._motion_num_frames + lengths_shifted = lengths.roll(1) + lengths_shifted[0] = 0 + self.length_starts = lengths_shifted.cumsum(0) + + # Zero out initial root XY so all motions start at origin + if self.m_cfg.get("zero_root_xy", False): + print( # noqa: T201 + f"[zero_root_xy] Zeroing initial root XY for {len(motions)} motions" + ) # noqa: RUF100, T201 + for i in range(len(motions)): + start = self.length_starts[i] + end = start + self._motion_num_frames[i] + init_xy = self.body_pos_w[start, 0, :2].clone() # root body, XY + print( # noqa: T201 + f" Motion {i}: init_xy=[{init_xy[0]:.3f}, {init_xy[1]:.3f}], frames={self._motion_num_frames[i]}" # noqa: E501 + ) + self.body_pos_w[start:end, :, :2] -= init_xy + if ( + hasattr(self, "_motion_object_root_pos") + and self._motion_object_root_pos is not None + ): + self._motion_object_root_pos[start:end, :, :2] -= init_xy + + self.motion_ids = torch.arange(len(motions), dtype=torch.long, device=self._device) + + motion_has_smpl = [ + self.curr_motion_keys[i] in self.smpl_data_keys for i in range(len(motions)) + ] + self.motion_has_smpl = torch.tensor(motion_has_smpl, dtype=torch.bool, device=self._device) + + motion_has_soma = [ + self.curr_motion_keys[i] in self.soma_data_keys for i in range(len(motions)) + ] + self.motion_has_soma = torch.tensor(motion_has_soma, dtype=torch.bool, device=self._device) + + motion_has_object = [ + self.curr_motion_keys[i] in self.object_data_keys for i in range(len(motions)) + ] + self.motion_has_object = torch.tensor( + motion_has_object, dtype=torch.bool, device=self._device + ) + + motion = motions[0] # noqa: F841 + self.num_bodies = self.num_joints + + num_motions = self.num_motions() + total_len = self.get_total_length() + + if self.use_adaptive_sampling: + self.update_adaptive_sampling_motion_frames() + + logger.info( + f"Loaded {num_motions:d} motions with a total length of {total_len:.3f}s and {self.body_pos_w.shape[0]} frames." # noqa: E501 + ) + + del ( + motions, + _motion_lengths, + _motion_fps, + _motion_dt, + _motion_num_frames, + _motion_bodies, + _motion_aa, + _motion_actions, + _motion_smpl_poses, + _motion_smpl_joints, + _motion_smpl_transl, + _motion_object_root_pos, + _motion_object_root_quat, + ) + + if "mujoco_to_isaaclab_body" in self.m_cfg.keys(): # noqa: SIM118 + self.dof_pos = self.dof_pos[:, self.m_cfg.mujoco_to_isaaclab_dof] + self.dof_vel = self.dof_vel[:, self.m_cfg.mujoco_to_isaaclab_dof] + + # Keep full body data (all bodies, IsaacLab order) before slicing + self.num_bodies_full = len(self.m_cfg.mujoco_to_isaaclab_body) + self.body_pos_w_full = self.body_pos_w[:, self.m_cfg.mujoco_to_isaaclab_body] + self.body_quat_w_full = rotations.xyzw_to_wxyz( + self.body_quat_w[:, self.m_cfg.mujoco_to_isaaclab_body] + ) + self.body_lin_vel_w_full = self.body_lin_vel_w[:, self.m_cfg.mujoco_to_isaaclab_body] + self.body_ang_vel_w_full = self.body_ang_vel_w[:, self.m_cfg.mujoco_to_isaaclab_body] + + # Slice to only selected body_indexes + self.body_pos_w = self.body_pos_w_full[:, self.body_indexes] + self.body_quat_w = self.body_quat_w_full[:, self.body_indexes] + self.body_lin_vel_w = self.body_lin_vel_w_full[:, self.body_indexes] + self.body_ang_vel_w = self.body_ang_vel_w_full[:, self.body_indexes] + assert ( + self.m_cfg.get("anchor_body_idx_full", 0) == 0 and self.body_indexes[0] == 0 + ), "The anchor body has to be 0; otherwise will cause issues in the sliced body_indexes data's anchor." + else: + # No body reordering — full body data is the same as the original data + self.body_pos_w_full = self.body_pos_w + self.body_quat_w_full = self.body_quat_w + self.body_lin_vel_w_full = self.body_lin_vel_w + self.body_ang_vel_w_full = self.body_ang_vel_w + self.num_bodies_full = self.body_pos_w.shape[2] + + # Run cleanup after slicing so temporary fragments do not live through the next cycle. + gc.collect() + torch.cuda.empty_cache() + + def foot_detect(self, positions, vel_thres, height_thresh): + fid_l = self.m_cfg.get("left_foot_body_idx", [6]) + fid_r = self.m_cfg.get("right_foot_body_idx", [12]) + # fid_l, fid_r = [6], [12] + velfactor = torch.tensor( + [vel_thres] * len(fid_l), device=positions.device, dtype=positions.dtype + ) + heightfactor = torch.tensor( + [height_thresh] * len(fid_l), device=positions.device, dtype=positions.dtype + ) + + feet_l_xyz = (positions[1:, fid_l] - positions[:-1, fid_l]) ** 2 + feet_l_xyz = torch.cat([feet_l_xyz, feet_l_xyz[[-1]]], dim=0) + feet_l_h = positions[:, fid_l, 2] + feet_l = torch.logical_and( + (feet_l_xyz.sum(dim=-1)) < velfactor, feet_l_h < heightfactor + ).float() + # feet_l = ((feet_l_x + feet_l_y + feet_l_z) < velfactor).float() + + feet_r_xyz = (positions[1:, fid_r] - positions[:-1, fid_r]) ** 2 + feet_r_xyz = torch.cat([feet_r_xyz, feet_r_xyz[[-1]]], dim=0) + feet_r_h = positions[:, fid_r, 2] + feet_r = torch.logical_and( + (feet_r_xyz.sum(dim=-1)) < velfactor, feet_r_h < heightfactor + ).float() + # feet_r = (((feet_r_x + feet_r_y + feet_r_z) < velfactor)).float() + return feet_l, feet_r + + def _compute_object_velocities(self, motion_num_frames, motion_dt): + """Compute object linear and angular velocities from position and quaternion data. + Uses finite differences: v = (p_{t+1} - p_t) / dt + Handles motion boundaries properly (first frame uses forward difference). + """ # noqa: D205 + total_frames = self._motion_object_root_pos.shape[0] + num_objects = self._motion_object_root_pos.shape[1] + + # Initialize velocity tensors + self._motion_object_lin_vel = torch.zeros_like(self._motion_object_root_pos) + self._motion_object_ang_vel = torch.zeros( + total_frames, num_objects, 3, device=self._device, dtype=torch.float32 + ) + + # Compute length_starts for indexing (cumsum of frame counts) + num_frames_tensor = torch.tensor(motion_num_frames, device=self._device) + lengths_shifted = num_frames_tensor.roll(1) + lengths_shifted[0] = 0 + length_starts = lengths_shifted.cumsum(0) + + # Compute velocities for each motion sequence separately + for i, (start, num_frames, dt) in enumerate( # noqa: B007 + zip(length_starts, motion_num_frames, motion_dt) # noqa: B905 + ): + start = start.item() # noqa: PLW2901 + end = start + num_frames + + if num_frames < 2: + continue # Cannot compute velocity with less than 2 frames + + # Get position and quaternion for this motion + pos = self._motion_object_root_pos[start:end] # (T, N_obj, 3) + quat = self._motion_object_root_quat[start:end] # (T, N_obj, 4) + + # Compute linear velocity: v = (p_{t+1} - p_t) / dt + lin_vel = (pos[1:] - pos[:-1]) / dt + # First frame uses same velocity as second frame + lin_vel = torch.cat([lin_vel[:1], lin_vel], dim=0) + self._motion_object_lin_vel[start:end] = lin_vel + + # Compute angular velocity from quaternion difference using same method as robot body + # ω = axis * angle / dt (same as _compute_angular_velocity in torch_humanoid_batch.py) + q_curr = quat[:-1] # (T-1, N_obj, 4) + q_next = quat[1:] # (T-1, N_obj, 4) + + # Compute quaternion difference: q_diff = q_next * q_curr^{-1} + # Using quat_mul_norm and quat_inverse (w_last=False for xyzw format) + diff_quat = rotations.quat_mul_norm( + q_next, rotations.quat_inverse(q_curr, w_last=False), w_last=False + ) + + # Extract angle and axis from quaternion difference + diff_angle, diff_axis = rotations.quat_angle_axis(diff_quat, w_last=False) + + # Angular velocity: ω = axis * angle / dt + ang_vel = diff_axis * diff_angle.unsqueeze(-1) / dt + # First frame uses same velocity as second frame + ang_vel = torch.cat([ang_vel[:1], ang_vel], dim=0) + self._motion_object_ang_vel[start:end] = ang_vel + + logger.info(f"Computed object velocities for {len(motion_num_frames)} motions") + + def fix_trans_height(self, pose_aa, trans, fix_height_mode): + if fix_height_mode == FixHeightMode.no_fix: + return trans, 0 + with torch.no_grad(): + + mesh_obj = self.mesh_parsers.mesh_fk(pose_aa[None, :1], trans[None, :1]) + height_diff = np.asarray(mesh_obj.vertices)[..., 2].min() + trans[..., 2] -= height_diff + + return trans, height_diff + + def load_motion_with_skeleton( + self, + ids, + motion_data_list, + smpl_data_list, + object_data_list, + soma_data_list, + fix_height, + target_heading, # noqa: ARG002 + max_len, + is_evaluation, + queue, + pid, + ): + # loading motion with the specified skeleton. Perfoming forward kinematics to get the joint positions + res = {} + + if pid == 0: # noqa: SIM108 + pbar = progress.track(range(len(ids)), description="Loading motions...") + else: + pbar = range(len(ids)) + + for f in pbar: + + curr_id = ids[f] # id for this datasample + + curr_file = motion_data_list[f] + if "path" in curr_file: + curr_file, *_ = joblib.load( + curr_file["path"] + ).values() # First value since it's a single item dictionary + + seq_len = curr_file["root_trans_offset"].shape[0] + if max_len == -1 or seq_len < max_len: + start, end = 0, seq_len + else: + start = random.randint(0, seq_len - max_len) + end = start + max_len + + trans = to_torch(curr_file["root_trans_offset"]).clone()[start:end] + pose_aa = to_torch(curr_file["pose_aa"][start:end]).clone() + + # import ipdb; ipdb.set_trace() + if "action" in curr_file.keys(): # noqa: SIM118 + self.has_action = True + + if "fps" not in curr_file.keys(): # noqa: SIM118 + curr_file["fps"] = 30.0 + dt = 1 / curr_file["fps"] # noqa: F841 + + B, J, N = pose_aa.shape + freeze_frame_aug, freeze_idx = False, 0 + + # self.m_cfg.freeze_frame_aug=True; is_evaluation=False; self.m_cfg.freeze_frame_prob=1 + # Debugging, force freeze frame augmentation + + if not is_evaluation and self.m_cfg.get("freeze_frame_aug", False): + freeze_prob = self.m_cfg.get("freeze_frame_prob", 0.1) + if np.random.random() < freeze_prob: # noqa: NPY002 + # Freeze the sequence at a random index + freeze_frame_aug = True + freeze_idx = np.random.randint(0, B) # noqa: NPY002 + # Repeat the frozen frame for all subsequent frames + pose_aa[freeze_idx:] = pose_aa[freeze_idx : freeze_idx + 1].clone() + trans[freeze_idx:] = trans[freeze_idx : freeze_idx + 1].clone() + + if not is_evaluation and self.m_cfg.get("randomize_heading", False): + # ZL: this randomization is not combatiable with SMPL + random_rot = np.zeros(3) + random_rot[2] = np.pi * (2 * np.random.random() - 1.0) # noqa: NPY002 + random_heading_rot = transform.Rotation.from_euler("xyz", random_rot) + pose_aa = pose_aa.reshape(B, -1) + pose_aa[:, :3] = torch.tensor( + ( + random_heading_rot * transform.Rotation.from_rotvec(pose_aa[:, :3]) + ).as_rotvec() + ) + trans = torch.matmul( + trans, torch.from_numpy(random_heading_rot.as_matrix().T).float() + ) + pose_aa = pose_aa.reshape(B, J, N) + + # self.cat_upper_body_poses_prob of the time, randomize the upper body poses and only for the motions are generated kinematically. # noqa: E501 + randomize_upper_body_poses = ( + self.randomize_upper_body_poses + and random.random() < self.cat_upper_body_poses_prob + and (self._should_augment_upper_body(self.curr_motion_keys[curr_id])) + ) + + # only randomize the upper body poses if the non-navigation dataset is loaded + if ( + randomize_upper_body_poses + and self.res_non_nav_dataset is not None + and len(self.res_non_nav_dataset) > 0 + ): + # ZL: this randomization is not combatiable with SMPL, so only for kinematic generated data. + # find the index for the upper body, skip the first index in pose_aa as it is the root. + upper_body_indices = [ + i for i in range(1, J) if i - 1 not in self.m_cfg.lower_joint_indices_mujoco + ] + # randomly select a motion from the non-navigation dataset + selected_file, selected_motion = random.choice( + list(self.res_non_nav_dataset.values()) + ) + selected_pose_aa = to_torch(selected_file["pose_aa"]) + + # Sample a matching slice from the selected motion + # Use the same method as main code to determine sequence length + selected_seq_len = selected_file["root_trans_offset"].shape[0] + current_seq_len = pose_aa.shape[0] + if selected_seq_len >= current_seq_len: + selected_start = random.randint(0, selected_seq_len - current_seq_len) + selected_end = selected_start + current_seq_len + selected_slice = selected_pose_aa[selected_start:selected_end] + else: + # If selected motion is shorter, create a ping-pong (forward then backward) sequence + forward = selected_pose_aa + backward = selected_pose_aa.flip(dims=[0]) # reverse the sequence + # Concatenate forward and backward, excluding the last frame of forward to avoid duplication + extended = torch.cat([forward, backward[1:]], dim=0) + + # If still not long enough, repeat the extended sequence + if extended.shape[0] < current_seq_len: + repeats = (current_seq_len + extended.shape[0] - 1) // extended.shape[ + 0 + ] # ceiling division + extended = extended.repeat(repeats, 1, 1) + + selected_slice = extended[:current_seq_len] + + pose_aa[:, upper_body_indices] = selected_slice[:, upper_body_indices] + + # Wrist joint noise augmentation + if ( + not is_evaluation + and self.randomize_wrist_poses + and random.random() < self.randomize_wrist_prob + ): + wrist_pose_aa_indices = [d + 1 for d in self.wrist_mujoco_dof_indices] + noise = torch.randn(B, len(wrist_pose_aa_indices), N) * self.randomize_wrist_std + pose_aa[:, wrist_pose_aa_indices] = pose_aa[:, wrist_pose_aa_indices] + noise + + if self.mesh_parsers is not None: + trans, trans_fix = self.fix_trans_height(pose_aa, trans, fix_height_mode=fix_height) + curr_motion = self.mesh_parsers.fk_batch( + pose_aa[None,], + trans[None,], + return_full=True, + fps=curr_file["fps"], + target_fps=self.target_fps, + interpolate_data=True, + use_parallel_fk=self.use_parallel_fk, + ) + if self.smpl_data is not None: + curr_smpl_data = smpl_data_list[f] + if curr_smpl_data is not None: + if "path" in curr_smpl_data: + curr_smpl_data = joblib.load(curr_smpl_data["path"]) + + if curr_smpl_data["fps"] != self.target_fps: + smpl_pose = torch.tensor(curr_smpl_data["pose_aa"][start:end]).float() + smpl_pose[:, -6:] = 0.0 + curr_motion["smpl_pose"] = self.mesh_parsers.interploate_pose( + None, smpl_pose[None,], curr_smpl_data["fps"], self.target_fps + )[1][0] + else: + smpl_pose = torch.tensor(curr_smpl_data["pose_aa"]).float() + smpl_pose[:, -6:] = 0.0 + # new_seq_len = curr_motion['global_translation'].shape[1] + curr_motion["smpl_pose"] = smpl_pose + if "smpl_joints" in curr_smpl_data: + smpl_joints = torch.tensor(curr_smpl_data["smpl_joints"]).float() + curr_motion["smpl_joints"] = smpl_joints + if ( + curr_motion["smpl_joints"].shape[0] + != curr_motion["global_translation"].shape[1] + ): + print( # noqa: T201 + f"Length mismatch: smpl_joints={curr_motion['smpl_joints'].shape[0]}, " + f"global_translation={curr_motion['global_translation'].shape[1]}" + ) + print(smpl_data_list[f], motion_data_list[f]) # noqa: T201 + + assert ( + curr_motion["smpl_joints"].shape[0] + == curr_motion["global_translation"].shape[1] + ) + else: + num_frames = curr_motion["global_translation"].shape[1] + curr_motion["smpl_joints"] = torch.zeros(num_frames, 24, 3).to( + curr_motion["global_translation"] + ) + if "transl" in curr_smpl_data: + transl = torch.tensor(curr_smpl_data["transl"]).float() + curr_motion["smpl_transl"] = transl + assert ( + curr_motion["smpl_transl"].shape[0] + == curr_motion["global_translation"].shape[1] + ) + else: + num_frames = curr_motion["global_translation"].shape[1] + curr_motion["smpl_transl"] = torch.zeros(num_frames, 3).to( + curr_motion["global_translation"] + ) + assert ( + curr_motion["smpl_pose"].shape[0] + == curr_motion["global_translation"].shape[1] + ) + + if freeze_frame_aug: + freeze_idx_new_fps = int( + freeze_idx * self.target_fps / curr_file["fps"] + ) + curr_motion["smpl_pose"][freeze_idx_new_fps:] = curr_motion[ + "smpl_pose" + ][freeze_idx_new_fps : freeze_idx_new_fps + 1].clone() + curr_motion["smpl_joints"][freeze_idx_new_fps:] = curr_motion[ + "smpl_joints" + ][freeze_idx_new_fps : freeze_idx_new_fps + 1].clone() + curr_motion["smpl_transl"][freeze_idx_new_fps:] = curr_motion[ + "smpl_transl" + ][freeze_idx_new_fps : freeze_idx_new_fps + 1].clone() + else: + curr_motion["smpl_pose"] = torch.zeros( + curr_motion["global_translation"].shape[1], 72 + ).to(curr_motion["global_translation"]) + curr_motion["smpl_joints"] = torch.zeros( + curr_motion["global_translation"].shape[1], 24, 3 + ).to(curr_motion["global_translation"]) + curr_motion["smpl_transl"] = torch.zeros( + curr_motion["global_translation"].shape[1], 3 + ).to(curr_motion["global_translation"]) + # print(curr_motion['smpl_pose'].shape, curr_motion['global_translation'].shape) + + # Load SOMA skeleton data if available + if soma_data_list is not None: + curr_soma_data = soma_data_list[f] + if curr_soma_data is not None: + if "path" in curr_soma_data: + loaded = joblib.load(curr_soma_data["path"]) + curr_soma_data, *_ = loaded.values() + + num_frames = curr_motion["global_translation"].shape[1] + n_soma = self.num_soma_joints + + # Resample SOMA data using the canonical interploate_pose formula + # to match robot frame count from fk_batch. + soma_fps = curr_soma_data.get("fps", self.target_fps) + + if "soma_joints" in curr_soma_data: + soma_joints = torch.tensor(curr_soma_data["soma_joints"]).float() + soma_joints_orig_len = soma_joints.shape[0] + if soma_fps != self.target_fps: + soma_joints = self._resample_soma_tensor( + soma_joints, soma_fps, self.target_fps + ) + curr_motion["soma_joints"] = soma_joints + assert soma_joints.shape[0] == num_frames, ( + f"SOMA soma_joints length {soma_joints.shape[0]} != " + f"robot frames {num_frames} " + f"(soma_orig={soma_joints_orig_len} @ {soma_fps}fps, " + f"robot_orig={seq_len} @ {curr_file['fps']}fps, " + f"target_fps={self.target_fps})" + ) + else: + curr_motion["soma_joints"] = torch.zeros(num_frames, n_soma, 3).to( + curr_motion["global_translation"] + ) + + if "soma_root_quat" in curr_soma_data: + soma_root_quat = torch.tensor(curr_soma_data["soma_root_quat"]).float() + if soma_fps != self.target_fps: + soma_root_quat = self._resample_soma_tensor( + soma_root_quat, soma_fps, self.target_fps + ) + # Renormalize quaternions after linear interpolation + soma_root_quat = soma_root_quat / ( + soma_root_quat.norm(dim=-1, keepdim=True) + 1e-8 + ) + curr_motion["soma_root_quat"] = soma_root_quat + assert soma_root_quat.shape[0] == num_frames, ( + f"SOMA soma_root_quat length {soma_root_quat.shape[0]} != " + f"robot frames {num_frames}" + ) + else: + curr_motion["soma_root_quat"] = torch.zeros(num_frames, 4).to( + curr_motion["global_translation"] + ) + curr_motion["soma_root_quat"][:, 0] = 1.0 # identity in wxyz + + if "soma_transl" in curr_soma_data: + soma_transl = torch.tensor(curr_soma_data["soma_transl"]).float() + if soma_fps != self.target_fps: + soma_transl = self._resample_soma_tensor( + soma_transl, soma_fps, self.target_fps + ) + curr_motion["soma_transl"] = soma_transl + assert soma_transl.shape[0] == num_frames, ( + f"SOMA soma_transl length {soma_transl.shape[0]} != " + f"robot frames {num_frames}" + ) + else: + curr_motion["soma_transl"] = torch.zeros(num_frames, 3).to( + curr_motion["global_translation"] + ) + + if freeze_frame_aug: + freeze_idx_new_fps = int( + freeze_idx * self.target_fps / curr_file["fps"] + ) + for key in ("soma_joints", "soma_root_quat", "soma_transl"): + curr_motion[key][freeze_idx_new_fps:] = curr_motion[key][ + freeze_idx_new_fps : freeze_idx_new_fps + 1 + ].clone() + else: + num_frames = curr_motion["global_translation"].shape[1] + n_soma = self.num_soma_joints + curr_motion["soma_joints"] = torch.zeros(num_frames, n_soma, 3).to( + curr_motion["global_translation"] + ) + curr_motion["soma_root_quat"] = torch.zeros(num_frames, 4).to( + curr_motion["global_translation"] + ) + curr_motion["soma_root_quat"][:, 0] = 1.0 # identity in wxyz + curr_motion["soma_transl"] = torch.zeros(num_frames, 3).to( + curr_motion["global_translation"] + ) + + # Load object data if available + if self.object_data is not None: + curr_object_data = object_data_list[f] if object_data_list is not None else None + if curr_object_data is not None: + if "path" in curr_object_data: + loaded = joblib.load(curr_object_data["path"]) + curr_object_data, *_ = loaded.values() + + num_frames = curr_motion["global_translation"].shape[1] + original_fps = curr_object_data.get("fps", curr_file["fps"]) + + if "root_pos" in curr_object_data: + curr_motion["object_root_pos"] = interpolate_translation_data( + curr_object_data["root_pos"][start:end], + source_fps=original_fps, + target_fps=self.target_fps, + num_frames=num_frames, + max_num_objects=self.max_num_objects, + pad_value=0.0, + ) + + if "root_quat" in curr_object_data: + curr_motion["object_root_quat"] = interpolate_quaternion_data( + curr_object_data["root_quat"][start:end], + source_fps=original_fps, + target_fps=self.target_fps, + num_frames=num_frames, + max_num_objects=self.max_num_objects, + ) + + # Load per-hand contact points, compute contact centers and in_contact labels + for hand in ("left_hand", "right_hand"): + key = f"contact_points_{hand}" + side = hand.split("_")[0] # "left" or "right" + if key in curr_object_data: + # Remap dict keys to [start:end] slice so contact frames + # align with the sliced root_pos/root_quat data + raw_dict = curr_object_data[key] + sliced_dict = { + k - start: v for k, v in raw_dict.items() if start <= k < end + } + center, label = interpolate_contact_center( + sliced_dict, + source_fps=original_fps, + target_fps=self.target_fps, + num_frames=num_frames, + ) + curr_motion[f"object_contact_center_{side}"] = center + curr_motion[f"object_in_contact_{side}"] = label + else: + # Fill with zeros if no object data available + num_frames = curr_motion["global_translation"].shape[1] + curr_motion["object_root_pos"] = torch.zeros( + num_frames, self.max_num_objects, 3 + ).to(curr_motion["global_translation"]) + curr_motion["object_root_quat"] = torch.zeros( + num_frames, self.max_num_objects, 4 + ).to(curr_motion["global_translation"]) + curr_motion["object_root_quat"][ + :, :, 0 + ] = 1.0 # w=1 for identity quaternion + + curr_motion = easydict.EasyDict( + { + k: v.squeeze(dim=-1).squeeze(dim=0) if torch.is_tensor(v) else v + for k, v in curr_motion.items() + } + ) + # add "action" to curr_motion + if self.has_action: + curr_motion.action = to_torch(curr_file["action"]).clone()[start:end] + + # Extract hand DOFs if motion file has more than 29 DOFs + hand_dof_count = self.m_cfg.get("hand_dof_count", 0) + if hand_dof_count > 0 and "dof" in curr_file: + raw_dof = to_torch(curr_file["dof"]).clone()[start:end] + if raw_dof.shape[-1] > 29: + # Extract hand DOFs (indices 29 onwards) and interpolate to target FPS + hand_dof = raw_dof[:, 29 : 29 + hand_dof_count] + if curr_file["fps"] != self.target_fps: + # Simple linear interpolation for hand DOFs + num_target_frames = curr_motion["dof_pos"].shape[0] + hand_dof_interp = ( + torch.nn.functional.interpolate( + hand_dof.T.unsqueeze(0), # (1, C, T) + size=num_target_frames, + mode="linear", + align_corners=True, + ) + .squeeze(0) + .T + ) # (T, C) + curr_motion.hand_dof_pos = hand_dof_interp + else: + curr_motion.hand_dof_pos = hand_dof + + if self.vid_smpl_pose is not None: # for cross embodiment tracking + vid_smpl_pose = self.vid_smpl_pose[f] + vid_smpl_pose = self.mesh_parsers.interploate_pose( + None, vid_smpl_pose[None,], 30.0, self.target_fps + )[1][0] + if curr_motion["smpl_pose"].shape[0] < vid_smpl_pose.shape[0]: + for key in curr_motion.keys(): # noqa: SIM118 + if isinstance(curr_motion[key], torch.Tensor): + curr_motion[key] = torch.cat( + [ + curr_motion[key], + torch.zeros( + vid_smpl_pose.shape[0] - curr_motion[key].shape[0], + *curr_motion[key].shape[1:], + ).to(curr_motion[key]), + ], + dim=0, + ) + curr_motion["smpl_pose"][:, :] = vid_smpl_pose[ + : curr_motion["smpl_pose"].shape[0], : + ] + feet_l, feet_r = self.foot_detect(curr_motion["global_translation"], 0.0005, 0.05) + curr_motion["feet_l"] = feet_l + curr_motion["feet_r"] = feet_r + res[curr_id] = (curr_file, curr_motion) + else: + logger.error("No mesh parser found") + + if queue is not None: + queue.put(res) + else: + return res + + def num_motions(self): + return self._num_motions + + def get_total_length(self): + return sum(self._motion_lengths) + + def get_motion_num_steps(self, motion_ids=None): + if motion_ids is None: + return ( + (self._motion_num_frames * self._sim_fps / self._motion_fps).floor().int() + ) # don't use ceil as it will cause frames to be missed. + else: + return ( + (self._motion_num_frames[motion_ids] * self._sim_fps / self._motion_fps[motion_ids]) + .floor() + .int() + ) + + def sample_time(self, motion_ids, truncate_time=None): + n = len(motion_ids) # noqa: F841 + phase = torch.rand(motion_ids.shape, device=self._device) + motion_len = self._motion_lengths[motion_ids] + if truncate_time is not None: + assert truncate_time >= 0.0 + motion_len -= truncate_time + + motion_time = phase * motion_len + return motion_time.to(self._device) + + def sample_time_steps(self, motion_ids, truncate_time=None): + motion_time = self.sample_time(motion_ids, truncate_time) + motion_time_steps = (motion_time * self._sim_fps).floor().int() + return motion_time_steps + + def sample_motions(self, n): + motion_ids = torch.multinomial( + self._sampling_batch_prob, num_samples=n, replacement=True + ).to(self._device) + + return motion_ids + + def get_motion_ids_in_dataset(self, motion_ids): + return self._curr_motion_ids[motion_ids] + + def get_motion_length(self, motion_ids=None): + if motion_ids is None: + return self._motion_lengths + else: + return self._motion_lengths[motion_ids] + + def _calc_frame_blend(self, time, len, num_frames, dt): # noqa: A002 + time = time.clone() + phase = time / len + phase = torch.clip(phase, 0.0, 1.0) # clip time to be within motion length. + time[time < 0] = 0 + + frame_idx0 = (phase * (num_frames - 1)).long() + frame_idx1 = torch.min(frame_idx0 + 1, num_frames - 1) + blend = torch.clip( + (time - frame_idx0 * dt) / dt, 0.0, 1.0 + ) # clip blend to be within 0 and 1 + + return frame_idx0, frame_idx1, blend + + def _get_num_bodies(self): + return self.num_bodies + + def _local_rotation_to_dof_smpl(self, local_rot): + B, J, _ = local_rot.shape + dof_pos = rotations.quat_to_exp_map(local_rot[:, 1:]) + return dof_pos.reshape(B, -1) + + def init_adaptive_sampling(self): + """Initialize adaptive sampling data structures over all unique motions. + + Divides every motion clip into fixed-size bins (``bin_size`` frames each) and + creates per-bin tracking tensors for failure rates and sampling probabilities. + This enables fine-grained, time-segment-level curriculum learning: bins with + higher failure rates are sampled more frequently during training. + + NOTE: This operates over ALL unique motions in the dataset (not just the + currently loaded batch), so bin indices are stable across reloads. + """ + self.adp_samp_num_frames = torch.zeros( + self._num_unique_motions, device=self._device, dtype=torch.long + ) + # Compute motion lengths and frame counts for all unique motions using self._motion_data_keys + for i, motion_key in enumerate(self._motion_data_keys): + motion_data = self._motion_data_load[motion_key] + + # Compute motion length and frame count similar to how it's done in load_motions + # Need to account for interpolation from original fps to target_fps + if "fps" not in motion_data.keys(): # noqa: SIM118 + motion_data["fps"] = 30.0 + + original_fps = motion_data["fps"] + # Get frame count: prefer metadata 'length', else 'root_trans_offset' shape, + # else lazy-load from pkl file (directory mode without metadata) + if "length" in motion_data: + original_num_frames = motion_data["length"] + elif "root_trans_offset" in motion_data: + original_num_frames = motion_data["root_trans_offset"].shape[0] + elif "path" in motion_data: + # Directory mode: lazy-load the pkl file to get frame count and fps + loaded_data, *_ = joblib.load(motion_data["path"]).values() + original_num_frames = loaded_data["root_trans_offset"].shape[0] + if "fps" in loaded_data: + original_fps = loaded_data["fps"] + else: + raise KeyError( + f"Cannot determine frame count for motion '{motion_key}': no 'length', 'root_trans_offset', or 'path' key" # noqa: E501 + ) + + original_duration = (original_num_frames - 1) / original_fps + + # Match fk_batch behavior: when fps == target_fps, interpolation is + # skipped and raw frames are used. Otherwise use the canonical + # interploate_pose formula (arange with exclusive end). + if original_fps == self.target_fps: + num_frames = original_num_frames + else: + num_frames = len(torch.arange(0, original_duration, 1 / self.target_fps)) + self.adp_samp_num_frames[i] = num_frames + + # Compute length_starts similar to how it's done in load_motions (using num_frames, not lengths) + lengths = self.adp_samp_num_frames + lengths_shifted = lengths.roll(1) + lengths_shifted[0] = 0 + self.adp_samp_length_starts = lengths_shifted.cumsum(0) + self.adp_samp_total_frames = self.adp_samp_num_frames.sum() + self.adp_samp_length_starts_mask = torch.zeros( + self.adp_samp_total_frames, device=self._device, dtype=torch.bool + ) + self.adp_samp_length_starts_mask[self.adp_samp_length_starts] = True + + # init bins - batch version + self.adp_samp_bin_size = self.adaptive_sampling_cfg.get("bin_size", 50) + self.adp_samp_frame_to_bin = torch.zeros( + self.adp_samp_total_frames, device=self._device, dtype=torch.long + ) + + # Pre-compute all bin information in batch + all_bins = [] + all_bin_motion_lengths = [] + all_bin_new_motion_masks = [] + all_num_peer_bins = [] + all_motion_to_bins = [] + + cur_bin_idx = 0 + for orig_motion_id in range(self._num_unique_motions): + num_frames = self.adp_samp_num_frames[orig_motion_id] + frame_start = self.adp_samp_length_starts[orig_motion_id] + frame_end = ( + self.adp_samp_length_starts[orig_motion_id + 1] + if orig_motion_id < self._num_unique_motions - 1 + else self.adp_samp_total_frames + ) + + # Create bin starts and ends in batch + bin_starts = torch.arange( + 0, num_frames, self.adp_samp_bin_size, device=self._device, dtype=torch.long + ) + bin_ends = torch.minimum(bin_starts + self.adp_samp_bin_size, num_frames) + num_bins = len(bin_starts) + motion_ids = torch.full( + (num_bins,), orig_motion_id, device=self._device, dtype=torch.long + ) + motion_bins = torch.stack([motion_ids, bin_starts, bin_ends], dim=1) + all_bins.append(motion_bins) + + # Calculate bin lengths + bin_lengths = bin_ends - bin_starts + all_bin_motion_lengths.append(bin_lengths) + + # Create new motion mask (first bin of each motion is True) + new_motion_mask = torch.zeros(num_bins, device=self._device, dtype=torch.bool) + new_motion_mask[0] = True + all_bin_new_motion_masks.append(new_motion_mask) + + # Number of peer bins (same for all bins in this motion) + peer_bins = torch.full((num_bins,), num_bins, device=self._device, dtype=torch.long) + all_num_peer_bins.append(peer_bins) + + bin_ids = torch.zeros(num_frames, device=self._device, dtype=torch.long) + bin_ids[bin_starts[1:]] = 1 + bin_ids = bin_ids.cumsum(0) + cur_bin_idx + self.adp_samp_frame_to_bin[frame_start:frame_end] = bin_ids + + # Store motion to bins mapping + motion_bin_indices = torch.arange( + cur_bin_idx, cur_bin_idx + num_bins, device=self._device, dtype=torch.long + ) + all_motion_to_bins.append(motion_bin_indices) + + cur_bin_idx += num_bins + + # Concatenate all batch results + self.adp_samp_bins = torch.cat(all_bins, dim=0) + self.adp_samp_bin_motion_length = torch.cat(all_bin_motion_lengths, dim=0) + self.adp_samp_bin_new_motion_mask = torch.cat(all_bin_new_motion_masks, dim=0) + self.adp_samp_num_peer_bins = torch.cat(all_num_peer_bins, dim=0) + self.orig_motion_id_to_bins = all_motion_to_bins + self.adp_samp_num_bins = len(self.adp_samp_bins) + + self.adp_samp_bin_weights = ( + self.adp_samp_bin_motion_length / self.adp_samp_bin_motion_length.float().mean() + ) + # this will make sure each sequence is sampled equally. + if self.adaptive_sampling_cfg.get("sequence_length_agnostic", True): + self.adp_samp_bin_weights = self.adp_samp_bin_weights / self.adp_samp_num_peer_bins + + init_num_failures = self.adaptive_sampling_cfg.get("init_num_failures", 1) + self.adp_samp_failure_rate_max_over_mean = self.adaptive_sampling_cfg.get( + "adp_samp_failure_rate_max_over_mean", 50.0 + ) + self.uniform_sampling_rate = self.adaptive_sampling_cfg.get("uniform_sampling_rate", 0.1) + + # Max probability constraints (None = skip, "auto" = use failure_rate_max_over_mean) + # These prevent over-concentration on challenging motions. See update_adaptive_sampling_probabilities(). + self.max_prob_per_bin_cfg = self.adaptive_sampling_cfg.get("max_prob_per_bin", None) + self.max_prob_per_motion_cfg = self.adaptive_sampling_cfg.get("max_prob_per_motion", None) + self.adp_samp_num_failures = ( + torch.ones(self.adp_samp_num_bins, device=self._device, dtype=torch.float32) + * init_num_failures + ) + self.adp_samp_num_episodes = ( + torch.ones(self.adp_samp_num_bins, device=self._device, dtype=torch.float32) + * init_num_failures + ) + self.adp_samp_failure_rate = torch.ones( + self.adp_samp_num_bins, device=self._device, dtype=torch.float32 + ) + self.adp_samp_failure_rate_raw = torch.ones( + self.adp_samp_num_bins, device=self._device, dtype=torch.float32 + ) + self.adp_sampling_prob = ( + torch.ones(self.adp_samp_num_bins, device=self._device, dtype=torch.float64) + / self.adp_samp_num_bins + ) + + def get_state_dict(self): + """Return a serializable state dict for checkpointing adaptive sampling stats. + + Returns: + Dict containing ``adp_samp_num_episodes`` and ``adp_samp_num_failures`` + tensors if adaptive sampling is enabled, otherwise an empty dict. + """ + state_dict = {} + if self.use_adaptive_sampling: + state_dict.update( + { + "adp_samp_num_episodes": self.adp_samp_num_episodes, + "adp_samp_num_failures": self.adp_samp_num_failures, + } + ) + return state_dict + + def load_state_dict(self, state_dict): + """Restore adaptive sampling statistics from a checkpoint. + + Validates that the bin count matches before restoring. If it does not match + (e.g. dataset changed between runs), the load is silently skipped. + + Args: + state_dict: Dict previously returned by ``get_state_dict()``. + """ + if self.use_adaptive_sampling and "adp_samp_num_episodes" in state_dict: + if len(self.adp_samp_num_failures) != len(state_dict["adp_samp_num_failures"]): + print("Adaptive sampling state dict does not match. Skipping load.") # noqa: T201 + return + + self.adp_samp_num_episodes[:] = state_dict["adp_samp_num_episodes"].to(self._device) + self.adp_samp_num_failures[:] = state_dict["adp_samp_num_failures"].to(self._device) + self.sync_and_compute_adaptive_sampling(sync_across_gpus=False) + return + + def update_adaptive_sampling(self, failure, motion_ids, motion_time_steps): + """Update adaptive sampling statistics based on training outcomes. + + Increments episode counts for all sampled bins, and failure counts for bins + where the policy terminated early. Uses bincount for efficient batched updates + when multiple environments hit the same bin. + + Args: + failure: Boolean tensor of shape ``(N,)`` indicating which environments + terminated due to failure (not timeout). + motion_ids: Tensor of shape ``(N,)`` with batch-local motion indices. + motion_time_steps: Tensor of shape ``(N,)`` with the simulation time step + at which the episode ended (or was sampled). + """ + # Convert motion_ids to dataset motion ids if needed + dataset_motion_ids = self.get_motion_ids_in_dataset(motion_ids) + + time_steps = self.adp_samp_length_starts[dataset_motion_ids] + motion_time_steps + + # Handle non-unique dataset_motion_ids by counting occurrences + if len(time_steps) > 0: + # Use bincount to count occurrences of each unique ID + bin_ids = self.adp_samp_frame_to_bin[time_steps] + counts = torch.bincount(bin_ids, minlength=self.adp_samp_num_bins) + counts = counts / self.adp_samp_bin_motion_length + self.adp_samp_num_episodes += counts + + # Update failure counts for failed motions + if failure.any(): + failed_time_steps = time_steps[failure] + # Handle non-unique failed motion IDs by counting occurrences + if len(failed_time_steps) > 0: + bin_ids = self.adp_samp_frame_to_bin[failed_time_steps] + failure_counts = torch.bincount(bin_ids, minlength=self.adp_samp_num_bins) + failure_counts_multiplier = self.adaptive_sampling_cfg.get( + "failure_counts_multiplier", 1 + ) + self.adp_samp_num_failures += failure_counts * failure_counts_multiplier + + def sync_and_compute_adaptive_sampling(self, accelerator=None, sync_across_gpus=False): + """Synchronize adaptive sampling stats across GPUs and recompute probabilities. + + In multi-GPU training, averages episode/failure counts across all processes + before recomputing the per-bin sampling distribution. Optionally applies + failure-rate decay to propagate difficulty information to preceding bins. + + Args: + accelerator: HuggingFace Accelerator instance for multi-GPU gather. + Required when ``sync_across_gpus=True``. + sync_across_gpus: Whether to synchronize statistics across distributed + processes before computing probabilities. + """ + if not self.use_adaptive_sampling: + return + + if sync_across_gpus: + with common.Timer("sync_adaptive_sampling_across_gpus"): + adp_samp_stats = torch.cat( + [self.adp_samp_num_episodes, self.adp_samp_num_failures], dim=-1 + ) + adp_samp_stats_all = accelerator.gather(adp_samp_stats).reshape( + -1, *adp_samp_stats.shape + ) + adp_samp_stats_all = adp_samp_stats_all.mean(dim=0) + self.adp_samp_num_episodes, self.adp_samp_num_failures = adp_samp_stats_all.chunk( + 2, dim=-1 + ) + + with common.Timer("compute_sampling_prob"): + failure_rate = self.adp_samp_num_failures / self.adp_samp_num_episodes + self.adp_samp_failure_rate_raw = failure_rate.clone() + self.adp_samp_failure_rate = failure_rate + # This is to compute the failure rate with decay. + # However, this is very slow and not necessary. We can just sample an offset before the failure happens. # noqa: E501 + if self.adaptive_sampling_cfg.get("use_failure_rate_decay", False): + gamma = self.adaptive_sampling_cfg.get("decay_gamma", 0.99) + num_steps = self.adp_samp_num_episodes.shape[0] + failure_rate_w_decay = torch.zeros_like(failure_rate) + for step in reversed(range(num_steps)): + if step == num_steps - 1: + next_failure_rate = 0 + next_is_not_terminal = 0.0 + else: + next_failure_rate = failure_rate_w_decay[step + 1] + next_is_not_terminal = ( + 1.0 - self.adp_samp_bin_new_motion_mask[step + 1].float() + ) + failure_rate_w_decay[step] = ( + failure_rate[step] + next_is_not_terminal * gamma * next_failure_rate + ) + self.adp_samp_failure_rate = failure_rate_w_decay + + # Compute the sampling probability based on the failure rate + self.update_adaptive_sampling_probabilities() + return + + def update_adaptive_sampling_probabilities(self): + """Recompute per-bin sampling probabilities for the currently loaded motion batch. + + Blends failure-rate-based probabilities with a uniform baseline (controlled by + ``uniform_sampling_rate``), then applies optional max-probability constraints + per bin and per motion to prevent over-concentration on outlier sequences. + See the inline comments for detailed rationale on the constraint design. + """ + self.adp_samp_failure_rate = self.adp_samp_failure_rate.double() + self.adp_samp_active_failure_rate = self.adp_samp_failure_rate[ + self.adp_samp_active_motion_bins + ] + adp_samp_failure_rate_upper_bound = ( + self.adp_samp_active_failure_rate.mean() * self.adp_samp_failure_rate_max_over_mean + ) + adp_samp_active_failure_rate_clipped = torch.clip( + self.adp_samp_active_failure_rate, 0.0, adp_samp_failure_rate_upper_bound + ) + failure_based_sampling_prob = ( + adp_samp_active_failure_rate_clipped / adp_samp_active_failure_rate_clipped.sum() + ) + uniform_sampling_prob = torch.ones_like(failure_based_sampling_prob) / len( + failure_based_sampling_prob + ) + self.adp_sampling_active_prob = ( + failure_based_sampling_prob * (1 - self.uniform_sampling_rate) + + uniform_sampling_prob * self.uniform_sampling_rate + ) + self.adp_sampling_active_prob *= self.adp_samp_bin_weights[self.adp_samp_active_motion_bins] + self.adp_sampling_active_prob = ( + self.adp_sampling_active_prob / self.adp_sampling_active_prob.sum() + ) + + # ========================================================================== + # MAX PROBABILITY CONSTRAINTS: Prevent over-concentration on challenging motions + # ========================================================================== + # WHY THESE CONSTRAINTS EXIST: + # --------------------------- + # Adaptive sampling focuses training on motions with higher failure rates. + # Without constraints, this can cause several problems: + # + # 1. CATASTROPHIC FORGETTING: If one motion has 90% failure rate while others + # have 10%, it could dominate sampling → policy forgets "easy" motions. + # + # 2. TRAINING INSTABILITY: Narrow sample distribution causes high gradient + # variance, leading to unstable training dynamics. + # + # 3. OVERFITTING TO OUTLIERS: Some motions may be impossible (bad mocap data, + # kinematic infeasibility) but still get sampled heavily, wasting compute. + # + # 4. DIVERSITY LOSS: For policies that need to generalize across many motions + # (e.g., CHIP_token compliance training with 18k+ clips), diversity is critical. + # + # EFFECTS OF THESE CONSTRAINTS: + # ----------------------------- + # - max_prob_per_bin: No single time-segment can exceed N× its fair share. + # Prevents over-sampling one specific "hard moment" in a motion. + # + # - max_prob_per_motion: No single motion clip can exceed N× its fair share. + # Prevents a single broken/impossible motion from dominating training. + # + # CONFIGURATION: + # -------------- + # - "auto": Uses adp_samp_failure_rate_max_over_mean to set the multiplier + # (e.g., if failure_rate_max=2, then max_prob = 2x uniform) + # - null/not set: SKIP these constraints entirely (for legacy configs) + # - 0: Explicitly disable the constraint + # - float value: Set exact max probability threshold + # + # For CHIP_token compliance training, we use conservative values (2x) to maintain + # motion diversity. For other training, higher values (50x+) may be acceptable. + # ========================================================================== + + # Skip all max_prob constraints if neither is configured (legacy behavior) + if self.max_prob_per_bin_cfg is None and self.max_prob_per_motion_cfg is None: + self.adp_sampling_active_prob = self.adp_sampling_active_prob.float() + assert (self.adp_sampling_active_prob >= 0).all() + return + + num_active_bins = len(self.adp_samp_active_motion_bins) + active_orig_motion_ids = self.adp_samp_bins[self.adp_samp_active_motion_bins, 0] + num_active_motions = len(active_orig_motion_ids.unique()) + + # 1. Max probability per bin: no single bin can exceed this fraction of total samples + if self.max_prob_per_bin_cfg is not None: + if self.max_prob_per_bin_cfg == "auto": + # Auto: use adp_samp_failure_rate_max_over_mean as multiplier + multiplier = self.adp_samp_failure_rate_max_over_mean + max_prob_per_bin = multiplier / num_active_bins if num_active_bins > 0 else 1.0 + else: + max_prob_per_bin = ( + float(self.max_prob_per_bin_cfg) if self.max_prob_per_bin_cfg else 0.0 + ) + + # Only apply if constraint is meaningful (more bins than 1/max_prob) + if max_prob_per_bin > 0 and num_active_bins > 1.0 / max_prob_per_bin: + self.adp_sampling_active_prob = torch.clamp( + self.adp_sampling_active_prob, max=max_prob_per_bin + ) + self.adp_sampling_active_prob = ( + self.adp_sampling_active_prob / self.adp_sampling_active_prob.sum() + ) + + # 2. Max probability per motion: aggregate bins per motion and cap total + if self.max_prob_per_motion_cfg is not None: + if self.max_prob_per_motion_cfg == "auto": + # Auto: use adp_samp_failure_rate_max_over_mean as multiplier + multiplier = self.adp_samp_failure_rate_max_over_mean + max_prob_per_motion = ( + multiplier / num_active_motions if num_active_motions > 0 else 1.0 + ) + else: + max_prob_per_motion = ( + float(self.max_prob_per_motion_cfg) if self.max_prob_per_motion_cfg else 0.0 + ) + + # Only apply if constraint is meaningful (more motions than 1/max_prob) + if max_prob_per_motion > 0 and num_active_motions > 1.0 / max_prob_per_motion: + unique_motions = active_orig_motion_ids.unique() + + for motion_id in unique_motions: + motion_mask = active_orig_motion_ids == motion_id + motion_total_prob = self.adp_sampling_active_prob[motion_mask].sum() + + if motion_total_prob > max_prob_per_motion: + # Scale down all bins belonging to this motion + scale_factor = max_prob_per_motion / motion_total_prob + self.adp_sampling_active_prob[motion_mask] *= scale_factor + + # Re-normalize after capping + self.adp_sampling_active_prob = ( + self.adp_sampling_active_prob / self.adp_sampling_active_prob.sum() + ) + # ========================================================================== + + self.adp_sampling_active_prob = self.adp_sampling_active_prob.float() + assert (self.adp_sampling_active_prob >= 0).all() + + def update_adaptive_sampling_motion_sequences(self): + """Recompute global (full-dataset) motion-level sampling probabilities. + + Called before ``load_motions()`` to determine which motions to load next. + Aggregates per-bin failure rates into per-motion probabilities and applies + the same max-probability constraints as the batch-level update. + """ + self.adp_samp_failure_rate = self.adp_samp_failure_rate.double() + adp_samp_failure_rate_upper_bound = ( + self.adp_samp_failure_rate.mean() * self.adp_samp_failure_rate_max_over_mean + ) + adp_samp_failure_rate_clipped = torch.clip( + self.adp_samp_failure_rate, 0.0, adp_samp_failure_rate_upper_bound + ) + failure_based_sampling_prob = ( + adp_samp_failure_rate_clipped / adp_samp_failure_rate_clipped.sum() + ) + uniform_sampling_prob = torch.ones_like(failure_based_sampling_prob) / len( + failure_based_sampling_prob + ) + self.adp_sampling_prob = ( + failure_based_sampling_prob * (1 - self.uniform_sampling_rate) + + uniform_sampling_prob * self.uniform_sampling_rate + ) + self.adp_sampling_prob *= self.adp_samp_bin_weights + self.adp_sampling_prob = self.adp_sampling_prob / self.adp_sampling_prob.sum() + + # ========================================================================== + # MAX PROBABILITY CONSTRAINTS (applied to global bin probabilities) + # See update_adaptive_sampling_probabilities() for detailed explanation. + # Skip if neither constraint is configured (legacy behavior). + # ========================================================================== + if self.max_prob_per_bin_cfg is None and self.max_prob_per_motion_cfg is None: + # Sum up the adp_sampling_prob for each motion's frames (no constraints) + motion_sampling_probs = torch.zeros(self._num_unique_motions, device=self._device) + for orig_motion_id in range(self._num_unique_motions): + motion_sampling_probs[orig_motion_id] = self.adp_sampling_prob[ + self.orig_motion_id_to_bins[orig_motion_id] + ].sum() + self._sampling_prob = motion_sampling_probs / motion_sampling_probs.sum() + return + + num_bins = self.adp_samp_num_bins + num_motions = self._num_unique_motions + + # Apply max_prob_per_bin constraint if configured + if self.max_prob_per_bin_cfg is not None: + if self.max_prob_per_bin_cfg == "auto": + # Auto: use adp_samp_failure_rate_max_over_mean as multiplier + multiplier = self.adp_samp_failure_rate_max_over_mean + max_prob_per_bin = multiplier / num_bins if num_bins > 0 else 1.0 + else: + max_prob_per_bin = ( + float(self.max_prob_per_bin_cfg) if self.max_prob_per_bin_cfg else 0.0 + ) + + # Only apply if constraint is meaningful (more bins than 1/max_prob) + if max_prob_per_bin > 0 and num_bins > 1.0 / max_prob_per_bin: + self.adp_sampling_prob = torch.clamp(self.adp_sampling_prob, max=max_prob_per_bin) + self.adp_sampling_prob = self.adp_sampling_prob / self.adp_sampling_prob.sum() + + # Sum up the adp_sampling_prob for each motion's frames + motion_sampling_probs = torch.zeros(self._num_unique_motions, device=self._device) + for orig_motion_id in range(self._num_unique_motions): + motion_sampling_probs[orig_motion_id] = self.adp_sampling_prob[ + self.orig_motion_id_to_bins[orig_motion_id] + ].sum() + + # Apply max_prob_per_motion constraint if configured + if self.max_prob_per_motion_cfg is not None: + if self.max_prob_per_motion_cfg == "auto": + # Auto: use adp_samp_failure_rate_max_over_mean as multiplier + multiplier = self.adp_samp_failure_rate_max_over_mean + max_prob_per_motion = multiplier / num_motions if num_motions > 0 else 1.0 + else: + max_prob_per_motion = ( + float(self.max_prob_per_motion_cfg) if self.max_prob_per_motion_cfg else 0.0 + ) + + # Only apply if constraint is meaningful (more motions than 1/max_prob) + if max_prob_per_motion > 0 and num_motions > 1.0 / max_prob_per_motion: + motion_sampling_probs = torch.clamp(motion_sampling_probs, max=max_prob_per_motion) + + self._sampling_prob = motion_sampling_probs / motion_sampling_probs.sum() + + def update_adaptive_sampling_motion_frames(self): + """Build the active-bin index for the currently loaded motion batch. + + Maps each loaded motion to its corresponding global bins, creating + ``adp_samp_active_motion_bins`` which is used by + ``sample_motion_ids_and_time_steps()`` and + ``update_adaptive_sampling_probabilities()`` to sample and update only + the bins that correspond to currently loaded motions. + """ + self.adp_samp_active_motion_bins = [] + self.orig_motion_id_to_motion_ids = torch.zeros( + self._num_unique_motions, device=self._device, dtype=torch.long + ) + for motion_id, orig_motion_id in enumerate(self._curr_motion_ids): + bins = self.orig_motion_id_to_bins[orig_motion_id.item()] + self.adp_samp_active_motion_bins.append(bins) + self.orig_motion_id_to_motion_ids[orig_motion_id.item()] = motion_id + + # Validate adaptive sampling frame count matches actual loaded frames + adp_frames = self.adp_samp_num_frames[orig_motion_id].item() + loaded_frames = self._motion_num_frames[motion_id].item() + assert adp_frames == loaded_frames, ( + f"Adaptive sampling frame count mismatch for motion " + f"{orig_motion_id.item()} (key={self._motion_data_keys[orig_motion_id]}): " + f"adp_samp={adp_frames}, loaded={loaded_frames}. " + f"This means init_adaptive_sampling computed a different frame count " + f"than fk_batch produced at load time." + ) + + self.adp_samp_active_motion_bins = torch.cat(self.adp_samp_active_motion_bins, dim=0) + self.update_adaptive_sampling_probabilities() + + def sample_motion_ids_and_time_steps(self, n): + """Sample motion IDs and time steps using adaptive sampling probabilities. + + Draws bins from the active-bin distribution, then samples a random frame + within each selected bin. Optionally shifts the sampled frame backward by + a random offset (``pre_failure_sample_window``) so the policy starts + practicing before the difficult segment. + + Args: + n: Number of (motion_id, time_step) pairs to sample. + + Returns: + Tuple of (motion_ids, motion_time_steps) where: + - motion_ids: ``(n,)`` long tensor with batch-local motion indices. + - motion_time_steps: ``(n,)`` int tensor with frame indices. + """ + sampled_bin_ids = torch.multinomial( + self.adp_sampling_active_prob, num_samples=n, replacement=True + ).to(self._device) + bin_ids = self.adp_samp_active_motion_bins[sampled_bin_ids] + bins = self.adp_samp_bins[bin_ids] + orig_motion_ids, bin_start, bin_end = bins[:, 0], bins[:, 1], bins[:, 2] + motion_ids = self.orig_motion_id_to_motion_ids[orig_motion_ids] + + motion_time_steps = ( + torch.rand(len(bin_start), device=bin_start.device) * (bin_end - bin_start) + ).floor().long() + bin_start + # Sample motion time steps before failures makes more sense since we need to take actions before the failure happens. # noqa: E501 + pre_failure_sample_window = self.adaptive_sampling_cfg.get("pre_failure_sample_window", 0) + if pre_failure_sample_window > 0: + offset = torch.randint(pre_failure_sample_window, (n,), device=self._device) + motion_time_steps = (motion_time_steps - offset).clamp_min(0) + return motion_ids, motion_time_steps.int() diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/motion_lib_robot.py b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/motion_lib_robot.py new file mode 100644 index 0000000000000000000000000000000000000000..030e6f415b005986d9eedfaafd277c1eaa41b663 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/motion_lib_robot.py @@ -0,0 +1,9 @@ +from gear_sonic.utils.motion_lib.motion_lib_base import MotionLibBase +from gear_sonic.utils.motion_lib.torch_humanoid_batch import Humanoid_Batch + + +class MotionLibRobot(MotionLibBase): + def __init__(self, motion_lib_cfg, num_envs, device): + super().__init__(motion_lib_cfg=motion_lib_cfg, num_envs=num_envs, device=device) + self.mesh_parsers = Humanoid_Batch(motion_lib_cfg) + return diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/skeleton.py b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/skeleton.py new file mode 100644 index 0000000000000000000000000000000000000000..daf8ddfae0f444ba4f977dbf3ebd2c901c755a95 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/skeleton.py @@ -0,0 +1,1477 @@ +# from abc import ABCMeta, abstractmethod, classmethod +from abc import abstractmethod +from collections import OrderedDict +import json +import os +import xml.etree.ElementTree as ET + +import numpy as np +import scipy.ndimage.filters as filters +import torch + +from gear_sonic.isaac_utils.rotations import ( + quat_identity_like, + quat_inverse, + quat_mul_norm, + transform_from_rotation_translation, + transform_mul, + transform_rotation, + transform_translation, +) + + +class NumpyEncoder(json.JSONEncoder): + """Special json encoder for numpy types""" + + def default(self, obj): + if isinstance( + obj, + np.int_ + | np.intc + | np.intp + | np.int8 + | np.int16 + | np.int32 + | np.int64 + | np.uint8 + | np.uint16 + | np.uint32 + | np.uint64, + ): + return int(obj) + elif isinstance(obj, np.float_ | np.float16 | np.float32 | np.float64): + return float(obj) + elif isinstance(obj, np.ndarray): + return {"__ndarray__": obj.tolist(), "dtype": str(obj.dtype), "shape": obj.shape} + return json.JSONEncoder.default(self, obj) + + +def json_numpy_obj_hook(dct): + if isinstance(dct, dict) and "__ndarray__" in dct: + data = np.asarray(dct["__ndarray__"], dtype=dct["dtype"]) + return data.reshape(dct["shape"]) + return dct + + +class Serializable: + """Implementation to read/write to file. + All class the is inherited from this class needs to implement to_dict() and + from_dict() + """ + + @classmethod + def from_dict(cls, dict_repr, *args, **kwargs): + """Read the object from an ordered dictionary + + :param dict_repr: the ordered dictionary that is used to construct the object + :type dict_repr: OrderedDict + :param args, kwargs: the arguments that need to be passed into from_dict() + :type args, kwargs: additional arguments + """ + pass + + @abstractmethod + def to_dict(self): + """Construct an ordered dictionary from the object + + :rtype: OrderedDict + """ + pass + + @classmethod + def from_file(cls, path, *args, **kwargs): + """Read the object from a file (either .npy or .json) + + :param path: path of the file + :type path: string + :param args, kwargs: the arguments that need to be passed into from_dict() + :type args, kwargs: additional arguments + """ + if path.endswith(".json"): + with open(path) as f: + d = json.load(f, object_hook=json_numpy_obj_hook) + elif path.endswith(".npy"): + d = np.load(path, allow_pickle=True).item() + else: + assert False, f"failed to load {cls.__name__} from {path}" + assert d["__name__"] == cls.__name__, "the file belongs to {}, not {}".format( + d["__name__"], cls.__name__ + ) + return cls.from_dict(d, *args, **kwargs) + + def to_file(self, path: str) -> None: + """Write the object to a file (either .npy or .json) + + :param path: path of the file + :type path: string + """ + if os.path.dirname(path) != "" and not os.path.exists(os.path.dirname(path)): + os.makedirs(os.path.dirname(path)) + d = self.to_dict() + d["__name__"] = self.__class__.__name__ + if path.endswith(".json"): + with open(path, "w") as f: + json.dump(d, f, cls=NumpyEncoder, indent=4) + elif path.endswith(".npy"): + np.save(path, d) + + +class TensorUtils(Serializable): + @classmethod + def from_dict(cls, dict_repr, *args, **kwargs): + """Read the object from an ordered dictionary + + :param dict_repr: the ordered dictionary that is used to construct the object + :type dict_repr: OrderedDict + :param kwargs: the arguments that need to be passed into from_dict() + :type kwargs: additional arguments + """ + return torch.from_numpy(dict_repr["arr"].astype(dict_repr["context"]["dtype"])) + + def to_dict(self): + """Construct an ordered dictionary from the object + + :rtype: OrderedDict + """ + return NotImplemented + + +def tensor_to_dict(x): + """Construct an ordered dictionary from the object + + :rtype: OrderedDict + """ + x_np = x.numpy() + return {"arr": x_np, "context": {"dtype": x_np.dtype.name}} + + +class SkeletonTree(Serializable): + """ + A skeleton tree gives a complete description of a rigid skeleton. It describes a tree structure + over a list of nodes with their names indicated by strings. Each edge in the tree has a local + translation associated with it which describes the distance between the two nodes that it + connects. + + Basic Usage: + >>> t = SkeletonTree.from_mjcf(SkeletonTree.__example_mjcf_path__) + >>> t + SkeletonTree( + node_names=['torso', 'front_left_leg', 'aux_1', 'front_left_foot', 'front_right_leg', 'aux_2', 'front_right_foot', 'left_back_leg', 'aux_3', 'left_back_foot', 'right_back_leg', 'aux_4', 'right_back_foot'], + parent_indices=tensor([-1, 0, 1, 2, 0, 4, 5, 0, 7, 8, 0, 10, 11]), + local_translation=tensor([[ 0.0000, 0.0000, 0.7500], + [ 0.0000, 0.0000, 0.0000], + [ 0.2000, 0.2000, 0.0000], + [ 0.2000, 0.2000, 0.0000], + [ 0.0000, 0.0000, 0.0000], + [-0.2000, 0.2000, 0.0000], + [-0.2000, 0.2000, 0.0000], + [ 0.0000, 0.0000, 0.0000], + [-0.2000, -0.2000, 0.0000], + [-0.2000, -0.2000, 0.0000], + [ 0.0000, 0.0000, 0.0000], + [ 0.2000, -0.2000, 0.0000], + [ 0.2000, -0.2000, 0.0000]]) + ) + >>> t.node_names + ['torso', 'front_left_leg', 'aux_1', 'front_left_foot', 'front_right_leg', 'aux_2', 'front_right_foot', 'left_back_leg', 'aux_3', 'left_back_foot', 'right_back_leg', 'aux_4', 'right_back_foot'] + >>> t.parent_indices + tensor([-1, 0, 1, 2, 0, 4, 5, 0, 7, 8, 0, 10, 11]) + >>> t.local_translation + tensor([[ 0.0000, 0.0000, 0.7500], + [ 0.0000, 0.0000, 0.0000], + [ 0.2000, 0.2000, 0.0000], + [ 0.2000, 0.2000, 0.0000], + [ 0.0000, 0.0000, 0.0000], + [-0.2000, 0.2000, 0.0000], + [-0.2000, 0.2000, 0.0000], + [ 0.0000, 0.0000, 0.0000], + [-0.2000, -0.2000, 0.0000], + [-0.2000, -0.2000, 0.0000], + [ 0.0000, 0.0000, 0.0000], + [ 0.2000, -0.2000, 0.0000], + [ 0.2000, -0.2000, 0.0000]]) + >>> t.parent_of('front_left_leg') + 'torso' + >>> t.index('front_right_foot') + 6 + >>> t[2] + 'aux_1' + """ + + __example_mjcf_path__ = os.path.join( + os.path.dirname(os.path.realpath(__file__)), "tests/ant.xml" + ) + + def __init__(self, node_names, parent_indices, local_translation): + """ + :param node_names: a list of names for each tree node + :type node_names: list[str] + :param parent_indices: an int32-typed tensor that represents the edge to its parent.\ + -1 represents the root node + :type parent_indices: Tensor + :param local_translation: a 3d vector that gives local translation information + :type local_translation: Tensor + """ + ln, lp, ll = len(node_names), len(parent_indices), len(local_translation) + assert len({ln, lp, ll}) == 1 + self._node_names = node_names + self._parent_indices = parent_indices.long() + self._local_translation = local_translation + self._node_indices = {self.node_names[i]: i for i in range(len(self))} + + def __len__(self): + """number of nodes in the skeleton tree""" + return len(self.node_names) + + def __iter__(self): + """iterator that iterate through the name of each node""" + yield from self.node_names + + def __getitem__(self, item): + """get the name of the node given the index""" + return self.node_names[item] + + def __repr__(self): + return ( + f"SkeletonTree(\n node_names={self._indent(repr(self.node_names))},\n parent_indices={self._indent(repr(self.parent_indices))}," + f"\n local_translation={self._indent(repr(self.local_translation))}\n)" + ) + + def _indent(self, s): + return "\n ".join(s.split("\n")) + + @property + def node_names(self): + return self._node_names + + @property + def parent_indices(self): + return self._parent_indices + + @property + def local_translation(self): + return self._local_translation + + @property + def num_joints(self): + """number of nodes in the skeleton tree""" + return len(self) + + @classmethod + def from_dict(cls, dict_repr, *args, **kwargs): + return cls( + list(map(str, dict_repr["node_names"])), + TensorUtils.from_dict(dict_repr["parent_indices"], *args, **kwargs), + TensorUtils.from_dict(dict_repr["local_translation"], *args, **kwargs), + ) + + def to_dict(self): + return OrderedDict( + [ + ("node_names", self.node_names), + ("parent_indices", tensor_to_dict(self.parent_indices)), + ("local_translation", tensor_to_dict(self.local_translation)), + ] + ) + + @classmethod + def from_mjcf(cls, path: str) -> "SkeletonTree": + """ + Parses a mujoco xml scene description file and returns a Skeleton Tree. + We use the model attribute at the root as the name of the tree. + + :param path: + :type path: string + :return: The skeleton tree constructed from the mjcf file + :rtype: SkeletonTree + """ + tree = ET.parse(path) + xml_doc_root = tree.getroot() + xml_world_body = xml_doc_root.find("worldbody") + if xml_world_body is None: + raise ValueError("MJCF parsed incorrectly please verify it.") + # assume this is the root + xml_body_root = xml_world_body.find("body") + if xml_body_root is None: + raise ValueError("MJCF parsed incorrectly please verify it.") + + node_names = [] + parent_indices = [] + local_translation = [] + + # recursively adding all nodes into the skel_tree + def _add_xml_node(xml_node, parent_index, node_index): + node_name = xml_node.attrib.get("name") + # parse the local translation into float list + pos = np.fromstring(xml_node.attrib.get("pos", "0 0 0"), dtype=float, sep=" ") + node_names.append(node_name) + parent_indices.append(parent_index) + local_translation.append(pos) + curr_index = node_index + node_index += 1 + for next_node in xml_node.findall("body"): + node_index = _add_xml_node(next_node, curr_index, node_index) + return node_index + + _add_xml_node(xml_body_root, -1, 0) + + return cls( + node_names, + torch.from_numpy(np.array(parent_indices, dtype=np.int32)), + torch.from_numpy(np.array(local_translation, dtype=np.float32)), + ) + + def parent_of(self, node_name): + """get the name of the parent of the given node + + :param node_name: the name of the node + :type node_name: string + :rtype: string + """ + return self[int(self.parent_indices[self.index(node_name)].item())] + + def index(self, node_name): + """get the index of the node + + :param node_name: the name of the node + :type node_name: string + :rtype: int + """ + return self._node_indices[node_name] + + def drop_nodes_by_names( + self, node_names: list[str], pairwise_translation=None + ) -> "SkeletonTree": + new_length = len(self) - len(node_names) + new_node_names = [] + new_local_translation = torch.zeros(new_length, 3, dtype=self.local_translation.dtype) + new_parent_indices = torch.zeros(new_length, dtype=self.parent_indices.dtype) + parent_indices = self.parent_indices.numpy() + new_node_indices: dict = {} + new_node_index = 0 + for node_index in range(len(self)): + if self[node_index] in node_names: + continue + tb_node_index = parent_indices[node_index] + if tb_node_index != -1: + local_translation = self.local_translation[node_index, :] + while tb_node_index != -1 and self[tb_node_index] in node_names: + local_translation += self.local_translation[tb_node_index, :] + tb_node_index = parent_indices[tb_node_index] + assert tb_node_index != -1, "the root node cannot be dropped" + + if pairwise_translation is not None: + local_translation = pairwise_translation[tb_node_index, node_index, :] + else: + local_translation = self.local_translation[node_index, :] + + new_node_names.append(self[node_index]) + new_local_translation[new_node_index, :] = local_translation + if tb_node_index == -1: + new_parent_indices[new_node_index] = -1 + else: + new_parent_indices[new_node_index] = new_node_indices[self[tb_node_index]] + new_node_indices[self[node_index]] = new_node_index + new_node_index += 1 + + return SkeletonTree(new_node_names, new_parent_indices, new_local_translation) + + def keep_nodes_by_names( + self, node_names: list[str], pairwise_translation=None + ) -> "SkeletonTree": + nodes_to_drop = list(filter(lambda x: x not in node_names, self)) + return self.drop_nodes_by_names(nodes_to_drop, pairwise_translation) + + +class SkeletonState(Serializable): + """ + A skeleton state contains all the information needed to describe a static state of a skeleton. + It requires a skeleton tree, local/global rotation at each joint and the root translation. + + Example: + >>> t = SkeletonTree.from_mjcf(SkeletonTree.__example_mjcf_path__) + >>> zero_pose = SkeletonState.zero_pose(t) + >>> plot_skeleton_state(zero_pose) # can be imported from `.visualization.base_interface` + [plot of the ant at zero pose + >>> local_rotation = zero_pose.local_rotation.clone() + >>> local_rotation[2] = torch.tensor([0, 0, 1, 0]) + >>> new_pose = SkeletonState.from_rotation_and_root_translation( + ... skeleton_tree=t, + ... r=local_rotation, + ... t=zero_pose.root_translation, + ... is_local=True + ... ) + >>> new_pose.local_rotation + tensor([[0., 0., 0., 1.], + [0., 0., 0., 1.], + [0., 1., 0., 0.], + [0., 0., 0., 1.], + [0., 0., 0., 1.], + [0., 0., 0., 1.], + [0., 0., 0., 1.], + [0., 0., 0., 1.], + [0., 0., 0., 1.], + [0., 0., 0., 1.], + [0., 0., 0., 1.], + [0., 0., 0., 1.], + [0., 0., 0., 1.]]) + >>> plot_skeleton_state(new_pose) # you should be able to see one of ant's leg is bent + [plot of the ant with the new pose + >>> new_pose.global_rotation # the local rotation is propagated to the global rotation at joint #3 + tensor([[0., 0., 0., 1.], + [0., 0., 0., 1.], + [0., 1., 0., 0.], + [0., 1., 0., 0.], + [0., 0., 0., 1.], + [0., 0., 0., 1.], + [0., 0., 0., 1.], + [0., 0., 0., 1.], + [0., 0., 0., 1.], + [0., 0., 0., 1.], + [0., 0., 0., 1.], + [0., 0., 0., 1.], + [0., 0., 0., 1.]]) + + Global/Local Representation (cont. from the previous example) + >>> new_pose.is_local + True + >>> new_pose.tensor # this will return the local rotation followed by the root translation + tensor([0., 0., 0., 1., 0., 0., 0., 1., 0., 1., 0., 0., 0., 0., 0., 1., 0., 0., + 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., + 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., + 0.]) + >>> new_pose.tensor.shape # 4 * 13 (joint rotation) + 3 (root translatio + torch.Size([55]) + >>> new_pose.global_repr().is_local + False + >>> new_pose.global_repr().tensor # this will return the global rotation followed by the root translation instead + tensor([0., 0., 0., 1., 0., 0., 0., 1., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0., + 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., + 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., + 0.]) + >>> new_pose.global_repr().tensor.shape # 4 * 13 (joint rotation) + 3 (root translation + torch.Size([55]) + """ + + def __init__(self, tensor_backend, skeleton_tree, is_local): + self._skeleton_tree = skeleton_tree + self._is_local = is_local + self.tensor = tensor_backend.clone() + + def __len__(self): + return self.tensor.shape[0] + + @property + def rotation(self): + if not hasattr(self, "_rotation"): + self._rotation = self.tensor[..., : self.num_joints * 4].reshape( + *(self.tensor.shape[:-1] + (self.num_joints, 4)) + ) + return self._rotation + + @property + def _local_rotation(self): + if self._is_local: + return self.rotation + else: + return None + + @property + def _global_rotation(self): + if not self._is_local: + return self.rotation + else: + return None + + @property + def is_local(self): + """is the rotation represented in local frame? + + :rtype: bool + """ + return self._is_local + + @property + def invariant_property(self): + return {"skeleton_tree": self.skeleton_tree, "is_local": self.is_local} + + @property + def num_joints(self): + """number of joints in the skeleton tree + + :rtype: int + """ + return self.skeleton_tree.num_joints + + @property + def skeleton_tree(self): + """skeleton tree + + :rtype: SkeletonTree + """ + return self._skeleton_tree + + @property + def root_translation(self): + """root translation + + :rtype: Tensor + """ + if not hasattr(self, "_root_translation"): + self._root_translation = self.tensor[..., self.num_joints * 4 : self.num_joints * 4 + 3] + return self._root_translation + + @property + def global_transformation(self): + """global transformation of each joint (transform from joint frame to global frame)""" + # Forward Kinematics + if not hasattr(self, "_global_transformation"): + local_transformation = self.local_transformation + global_transformation = [] + parent_indices = self.skeleton_tree.parent_indices.numpy() + # global_transformation = local_transformation.identity_like() + for node_index in range(len(self.skeleton_tree)): + parent_index = parent_indices[node_index] + if parent_index == -1: + global_transformation.append(local_transformation[..., node_index, :]) + else: + global_transformation.append( + transform_mul( + global_transformation[parent_index], + local_transformation[..., node_index, :], + ) + ) + self._global_transformation = torch.stack(global_transformation, axis=-2) + return self._global_transformation + + @property + def global_rotation(self): + """global rotation of each joint (rotation matrix to rotate from joint's F.O.R to global + F.O.R)""" + if self._global_rotation is None: + if not hasattr(self, "_comp_global_rotation"): + self._comp_global_rotation = transform_rotation(self.global_transformation) + return self._comp_global_rotation + else: + return self._global_rotation + + @property + def global_translation(self): + """global translation of each joint""" + if not hasattr(self, "_global_translation"): + self._global_translation = transform_translation(self.global_transformation) + return self._global_translation + + @property + def global_translation_xy(self): + """global translation in xy""" + trans_xy_data = self.global_translation.zeros_like() + trans_xy_data[..., 0:2] = self.global_translation[..., 0:2] + return trans_xy_data + + @property + def global_translation_xz(self): + """global translation in xz""" + trans_xz_data = self.global_translation.zeros_like() + trans_xz_data[..., 0:1] = self.global_translation[..., 0:1] + trans_xz_data[..., 2:3] = self.global_translation[..., 2:3] + return trans_xz_data + + @property + def local_rotation(self): + """the rotation from child frame to parent frame given in the order of child nodes appeared + in `.skeleton_tree.node_names`""" + if self._local_rotation is None: + if not hasattr(self, "_comp_local_rotation"): + local_rotation = quat_identity_like(self.global_rotation) + for node_index in range(len(self.skeleton_tree)): + parent_index = self.skeleton_tree.parent_indices[node_index] + if parent_index == -1: + local_rotation[..., node_index, :] = self.global_rotation[ + ..., node_index, : + ] + else: + local_rotation[..., node_index, :] = quat_mul_norm( + quat_inverse(self.global_rotation[..., parent_index, :]), + self.global_rotation[..., node_index, :], + ) + self._comp_local_rotation = local_rotation + return self._comp_local_rotation + else: + return self._local_rotation + + @property + def local_transformation(self): + """local translation + local rotation. It describes the transformation from child frame to + parent frame given in the order of child nodes appeared in `.skeleton_tree.node_names` + """ + if not hasattr(self, "_local_transformation"): + self._local_transformation = transform_from_rotation_translation( + r=self.local_rotation, t=self.local_translation + ) + return self._local_transformation + + @property + def local_translation(self): + """local translation of the skeleton state. It is identical to the local translation in + `.skeleton_tree.local_translation` except the root translation. The root translation is + identical to `.root_translation`""" + if not hasattr(self, "_local_translation"): + broadcast_shape = ( + tuple(self.tensor.shape[:-1]) + + (len(self.skeleton_tree),) + + tuple(self.skeleton_tree.local_translation.shape[-1:]) + ) + local_translation = self.skeleton_tree.local_translation.broadcast_to( + *broadcast_shape + ).clone() + local_translation[..., 0, :] = self.root_translation + self._local_translation = local_translation + return self._local_translation + + # Root Properties + @property + def root_translation_xy(self): + """root translation on xy""" + if not hasattr(self, "_root_translation_xy"): + self._root_translation_xy = self.global_translation_xy[..., 0, :] + return self._root_translation_xy + + @property + def global_root_rotation(self): + """root rotation""" + if not hasattr(self, "_global_root_rotation"): + self._global_root_rotation = self.global_rotation[..., 0, :] + return self._global_root_rotation + + @property + def global_root_yaw_rotation(self): + """root yaw rotation""" + if not hasattr(self, "_global_root_yaw_rotation"): + self._global_root_yaw_rotation = self.global_root_rotation.yaw_rotation() + return self._global_root_yaw_rotation + + # Properties relative to root + @property + def local_translation_to_root(self): + """The 3D translation from joint frame to the root frame.""" + if not hasattr(self, "_local_translation_to_root"): + self._local_translation_to_root = ( + self.global_translation - self.root_translation.unsqueeze(-1) + ) + return self._local_translation_to_root + + @property + def local_rotation_to_root(self): + """The 3D rotation from joint frame to the root frame. It is equivalent to + The root_R_world * world_R_node""" + return quat_inverse(self.global_root_rotation).unsqueeze(-1) * self.global_rotation + + def compute_forward_vector( + self, + left_shoulder_index, + right_shoulder_index, + left_hip_index, + right_hip_index, + gaussian_filter_width=20, + ): + """Computes forward vector based on cross product of the up vector with + average of the right->left shoulder and hip vectors""" + global_positions = self.global_translation + # Perpendicular to the forward direction. + # Uses the shoulders and hips to find this. + side_direction = ( + global_positions[:, left_shoulder_index].numpy() + - global_positions[:, right_shoulder_index].numpy() + + global_positions[:, left_hip_index].numpy() + - global_positions[:, right_hip_index].numpy() + ) + side_direction = side_direction / np.sqrt((side_direction**2).sum(axis=-1))[..., np.newaxis] + + # Forward direction obtained by crossing with the up direction. + forward_direction = np.cross(side_direction, np.array([[0, 1, 0]])) + + # Smooth the forward direction with a Gaussian. + # Axis 0 is the time/frame axis. + forward_direction = filters.gaussian_filter1d( + forward_direction, gaussian_filter_width, axis=0, mode="nearest" + ) + forward_direction = ( + forward_direction / np.sqrt((forward_direction**2).sum(axis=-1))[..., np.newaxis] + ) + + return torch.from_numpy(forward_direction) + + @staticmethod + def _to_state_vector(rot, rt): + state_shape = rot.shape[:-2] + vr = rot.reshape(*(state_shape + (-1,))) + vt = rt.broadcast_to(*state_shape + rt.shape[-1:]).reshape(*(state_shape + (-1,))) + v = torch.cat([vr, vt], axis=-1) + return v + + @classmethod + def from_dict( + cls: type["SkeletonState"], dict_repr: OrderedDict, *args, **kwargs + ) -> "SkeletonState": + rot = TensorUtils.from_dict(dict_repr["rotation"], *args, **kwargs) + rt = TensorUtils.from_dict(dict_repr["root_translation"], *args, **kwargs) + return cls( + SkeletonState._to_state_vector(rot, rt), + SkeletonTree.from_dict(dict_repr["skeleton_tree"], *args, **kwargs), + dict_repr["is_local"], + ) + + def to_dict(self) -> OrderedDict: + return OrderedDict( + [ + ("rotation", tensor_to_dict(self.rotation)), + ("root_translation", tensor_to_dict(self.root_translation)), + ("skeleton_tree", self.skeleton_tree.to_dict()), + ("is_local", self.is_local), + ] + ) + + @classmethod + def from_rotation_and_root_translation(cls, skeleton_tree, r, t, is_local=True): + """ + Construct a skeleton state from rotation and root translation + + :param skeleton_tree: the skeleton tree + :type skeleton_tree: SkeletonTree + :param r: rotation (either global or local) + :type r: Tensor + :param t: root translation + :type t: Tensor + :param is_local: to indicate that whether the rotation is local or global + :type is_local: bool, optional, default=True + """ + assert r.dim() > 0, f"the rotation needs to have at least 1 dimension (dim = {r.dim})" + state_vec = SkeletonState._to_state_vector(r, t) + + return cls( + state_vec, + skeleton_tree=skeleton_tree, + is_local=is_local, + ) + + @classmethod + def zero_pose(cls, skeleton_tree): + """ + Construct a zero-pose skeleton state from the skeleton tree by assuming that all the local + rotation is 0 and root translation is also 0. + + :param skeleton_tree: the skeleton tree as the rigid body + :type skeleton_tree: SkeletonTree + """ + return cls.from_rotation_and_root_translation( + skeleton_tree=skeleton_tree, + r=quat_identity([skeleton_tree.num_joints]), + t=torch.zeros(3, dtype=skeleton_tree.local_translation.dtype), + is_local=True, + ) + + def local_repr(self): + """ + Convert the skeleton state into local representation. This will only affects the values of + .tensor. If the skeleton state already has `is_local=True`. This method will do nothing. + + :rtype: SkeletonState + """ + if self.is_local: + return self + return SkeletonState.from_rotation_and_root_translation( + self.skeleton_tree, + r=self.local_rotation, + t=self.root_translation, + is_local=True, + ) + + def global_repr(self): + """ + Convert the skeleton state into global representation. This will only affects the values of + .tensor. If the skeleton state already has `is_local=False`. This method will do nothing. + + :rtype: SkeletonState + """ + if not self.is_local: + return self + return SkeletonState.from_rotation_and_root_translation( + self.skeleton_tree, + r=self.global_rotation, + t=self.root_translation, + is_local=False, + ) + + def _get_pairwise_average_translation(self): + global_transform_inv = transform_inverse(self.global_transformation) + p1 = global_transform_inv.unsqueeze(-2) + p2 = self.global_transformation.unsqueeze(-3) + + pairwise_translation = ( + transform_translation(transform_mul(p1, p2)) + .reshape(-1, len(self.skeleton_tree), len(self.skeleton_tree), 3) + .mean(axis=0) + ) + return pairwise_translation + + def _transfer_to(self, new_skeleton_tree: SkeletonTree): + old_indices = list(map(self.skeleton_tree.index, new_skeleton_tree)) + return SkeletonState.from_rotation_and_root_translation( + new_skeleton_tree, + r=self.global_rotation[..., old_indices, :], + t=self.root_translation, + is_local=False, + ) + + def drop_nodes_by_names( + self, node_names: list[str], estimate_local_translation_from_states: bool = True + ) -> "SkeletonState": + """ + Drop a list of nodes from the skeleton and re-compute the local rotation to match the + original joint position as much as possible. + + :param node_names: a list node names that specifies the nodes need to be dropped + :type node_names: List of strings + :param estimate_local_translation_from_states: the boolean indicator that specifies whether\ + or not to re-estimate the local translation from the states (avg.) + :type estimate_local_translation_from_states: boolean + :rtype: SkeletonState + """ + if estimate_local_translation_from_states: + pairwise_translation = self._get_pairwise_average_translation() + else: + pairwise_translation = None + new_skeleton_tree = self.skeleton_tree.drop_nodes_by_names(node_names, pairwise_translation) + return self._transfer_to(new_skeleton_tree) + + def keep_nodes_by_names( + self, node_names: list[str], estimate_local_translation_from_states: bool = True + ) -> "SkeletonState": + """ + Keep a list of nodes and drop all other nodes from the skeleton and re-compute the local + rotation to match the original joint position as much as possible. + + :param node_names: a list node names that specifies the nodes need to be dropped + :type node_names: List of strings + :param estimate_local_translation_from_states: the boolean indicator that specifies whether\ + or not to re-estimate the local translation from the states (avg.) + :type estimate_local_translation_from_states: boolean + :rtype: SkeletonState + """ + return self.drop_nodes_by_names( + list(filter(lambda x: (x not in node_names), self)), + estimate_local_translation_from_states, + ) + + def _remapped_to(self, joint_mapping: dict[str, str], target_skeleton_tree: SkeletonTree): + joint_mapping_inv = {target: source for source, target in joint_mapping.items()} + reduced_target_skeleton_tree = target_skeleton_tree.keep_nodes_by_names( + list(joint_mapping_inv) + ) + n_joints = ( + len(joint_mapping), + len(self.skeleton_tree), + len(reduced_target_skeleton_tree), + ) + assert ( + len(set(n_joints)) == 1 + ), "the joint mapping is not consistent with the skeleton trees" + source_indices = [ + self.skeleton_tree.index(joint_mapping_inv[x]) for x in reduced_target_skeleton_tree + ] + target_local_rotation = self.local_rotation[..., source_indices, :] + return SkeletonState.from_rotation_and_root_translation( + skeleton_tree=reduced_target_skeleton_tree, + r=target_local_rotation, + t=self.root_translation, + is_local=True, + ) + + def retarget_to( + self, + joint_mapping: dict[str, str], + source_tpose_local_rotation, + source_tpose_root_translation: np.ndarray, + target_skeleton_tree: SkeletonTree, + target_tpose_local_rotation, + target_tpose_root_translation: np.ndarray, + rotation_to_target_skeleton, + scale_to_target_skeleton: float, + z_up: bool = True, + ) -> "SkeletonState": + """ + Retarget the skeleton state to a target skeleton tree. This is a naive retarget + implementation with rough approximations. The function follows the procedures below. + + Steps: + 1. Drop the joints from the source (self) that do not belong to the joint mapping\ + with an implementation that is similar to "keep_nodes_by_names()" - take a\ + look at the function doc for more details (same for source_tpose) + + 2. Rotate the source state and the source tpose by "rotation_to_target_skeleton"\ + to align the source with the target orientation + + 3. Extract the root translation and normalize it to match the scale of the target\ + skeleton + + 4. Extract the global rotation from source state relative to source tpose and\ + re-apply the relative rotation to the target tpose to construct the global\ + rotation after retargetting + + 5. Combine the computed global rotation and the root translation from 3 and 4 to\ + complete the retargeting. + + 6. Make feet on the ground (global translation z) + + :param joint_mapping: a dictionary of that maps the joint node from the source skeleton to \ + the target skeleton + :type joint_mapping: dict[str, str] + + :param source_tpose_local_rotation: the local rotation of the source skeleton + :type source_tpose_local_rotation: Tensor + + :param source_tpose_root_translation: the root translation of the source tpose + :type source_tpose_root_translation: np.ndarray + + :param target_skeleton_tree: the target skeleton tree + :type target_skeleton_tree: SkeletonTree + + :param target_tpose_local_rotation: the local rotation of the target skeleton + :type target_tpose_local_rotation: Tensor + + :param target_tpose_root_translation: the root translation of the target tpose + :type target_tpose_root_translation: Tensor + + :param rotation_to_target_skeleton: the rotation that needs to be applied to the source\ + skeleton to align with the target skeleton. Essentially the rotation is t_R_s, where t is\ + the frame of reference of the target skeleton and s is the frame of reference of the source\ + skeleton + :type rotation_to_target_skeleton: Tensor + :param scale_to_target_skeleton: the factor that needs to be multiplied from source\ + skeleton to target skeleton (unit in distance). For example, to go from `cm` to `m`, the \ + factor needs to be 0.01. + :type scale_to_target_skeleton: float + :rtype: SkeletonState + """ + + # STEP 0: Preprocess + source_tpose = SkeletonState.from_rotation_and_root_translation( + skeleton_tree=self.skeleton_tree, + r=source_tpose_local_rotation, + t=source_tpose_root_translation, + is_local=True, + ) + target_tpose = SkeletonState.from_rotation_and_root_translation( + skeleton_tree=target_skeleton_tree, + r=target_tpose_local_rotation, + t=target_tpose_root_translation, + is_local=True, + ) + + # STEP 1: Drop the irrelevant joints + pairwise_translation = self._get_pairwise_average_translation() + node_names = list(joint_mapping) + new_skeleton_tree = self.skeleton_tree.keep_nodes_by_names(node_names, pairwise_translation) + + # TODO: combine the following steps before STEP 3 + source_tpose = source_tpose._transfer_to(new_skeleton_tree) + source_state = self._transfer_to(new_skeleton_tree) + + source_tpose = source_tpose._remapped_to(joint_mapping, target_skeleton_tree) + source_state = source_state._remapped_to(joint_mapping, target_skeleton_tree) + + # STEP 2: Rotate the source to align with the target + new_local_rotation = source_tpose.local_rotation.clone() + new_local_rotation[..., 0, :] = quat_mul_norm( + rotation_to_target_skeleton, source_tpose.local_rotation[..., 0, :] + ) + + source_tpose = SkeletonState.from_rotation_and_root_translation( + skeleton_tree=source_tpose.skeleton_tree, + r=new_local_rotation, + t=quat_rotate(rotation_to_target_skeleton, source_tpose.root_translation), + is_local=True, + ) + + new_local_rotation = source_state.local_rotation.clone() + new_local_rotation[..., 0, :] = quat_mul_norm( + rotation_to_target_skeleton, source_state.local_rotation[..., 0, :] + ) + source_state = SkeletonState.from_rotation_and_root_translation( + skeleton_tree=source_state.skeleton_tree, + r=new_local_rotation, + t=quat_rotate(rotation_to_target_skeleton, source_state.root_translation), + is_local=True, + ) + + # STEP 3: Normalize to match the target scale + root_translation_diff = ( + source_state.root_translation - source_tpose.root_translation + ) * scale_to_target_skeleton + + # STEP 4: the global rotation from source state relative to source tpose and + # re-apply to the target + current_skeleton_tree = source_state.skeleton_tree + target_tpose_global_rotation = source_state.global_rotation[0, :].clone() + for current_index, name in enumerate(current_skeleton_tree): + if name in target_tpose.skeleton_tree: + target_tpose_global_rotation[current_index, :] = target_tpose.global_rotation[ + target_tpose.skeleton_tree.index(name), : + ] + + global_rotation_diff = quat_mul_norm( + source_state.global_rotation, quat_inverse(source_tpose.global_rotation) + ) + new_global_rotation = quat_mul_norm(global_rotation_diff, target_tpose_global_rotation) + + # STEP 5: Putting 3 and 4 together + current_skeleton_tree = source_state.skeleton_tree + shape = source_state.global_rotation.shape[:-1] + shape = shape[:-1] + target_tpose.global_rotation.shape[-2:-1] + new_global_rotation_output = quat_identity(shape) + for current_index, name in enumerate(target_skeleton_tree): + while name not in current_skeleton_tree: + name = target_skeleton_tree.parent_of(name) + parent_index = current_skeleton_tree.index(name) + new_global_rotation_output[:, current_index, :] = new_global_rotation[ + :, parent_index, : + ] + + source_state = SkeletonState.from_rotation_and_root_translation( + skeleton_tree=target_skeleton_tree, + r=new_global_rotation_output, + t=target_tpose.root_translation + root_translation_diff, + is_local=False, + ).local_repr() + + return source_state + + def retarget_to_by_tpose( + self, + joint_mapping: dict[str, str], + source_tpose: "SkeletonState", + target_tpose: "SkeletonState", + rotation_to_target_skeleton, + scale_to_target_skeleton: float, + ) -> "SkeletonState": + """ + Retarget the skeleton state to a target skeleton tree. This is a naive retarget + implementation with rough approximations. See the method `retarget_to()` for more information + + :param joint_mapping: a dictionary of that maps the joint node from the source skeleton to \ + the target skeleton + :type joint_mapping: dict[str, str] + + :param source_tpose: t-pose of the source skeleton + :type source_tpose: SkeletonState + + :param target_tpose: t-pose of the target skeleton + :type target_tpose: SkeletonState + + :param rotation_to_target_skeleton: the rotation that needs to be applied to the source\ + skeleton to align with the target skeleton. Essentially the rotation is t_R_s, where t is\ + the frame of reference of the target skeleton and s is the frame of reference of the source\ + skeleton + :type rotation_to_target_skeleton: Tensor + :param scale_to_target_skeleton: the factor that needs to be multiplied from source\ + skeleton to target skeleton (unit in distance). For example, to go from `cm` to `m`, the \ + factor needs to be 0.01. + :type scale_to_target_skeleton: float + :rtype: SkeletonState + """ + assert ( + len(source_tpose.shape) == 0 and len(target_tpose.shape) == 0 + ), "the retargeting script currently doesn't support vectorized operations" + return self.retarget_to( + joint_mapping, + source_tpose.local_rotation, + source_tpose.root_translation, + target_tpose.skeleton_tree, + target_tpose.local_rotation, + target_tpose.root_translation, + rotation_to_target_skeleton, + scale_to_target_skeleton, + ) + + +class SkeletonMotion(SkeletonState): + + def __init__(self, tensor_backend, skeleton_tree, is_local, fps, *args, **kwargs): + self._fps = fps + super().__init__(tensor_backend, skeleton_tree, is_local, *args, **kwargs) + + def clone(self): + return SkeletonMotion(self.tensor.clone(), self.skeleton_tree, self._is_local, self._fps) + + @property + def invariant_property(self): + return { + "skeleton_tree": self.skeleton_tree, + "is_local": self.is_local, + "fps": self.fps, + } + + @property + def global_velocity(self): + """global velocity""" + curr_index = self.num_joints * 4 + 3 + return self.tensor[..., curr_index : curr_index + self.num_joints * 3].reshape( + *(self.tensor.shape[:-1] + (self.num_joints, 3)) + ) + + @property + def global_angular_velocity(self): + """global angular velocity""" + curr_index = self.num_joints * 7 + 3 + return self.tensor[..., curr_index : curr_index + self.num_joints * 3].reshape( + *(self.tensor.shape[:-1] + (self.num_joints, 3)) + ) + + @property + def fps(self): + """number of frames per second""" + return self._fps + + @property + def time_delta(self): + """time between two adjacent frames""" + return 1.0 / self.fps + + @property + def global_root_velocity(self): + """global root velocity""" + return self.global_velocity[..., 0, :] + + @property + def global_root_angular_velocity(self): + """global root angular velocity""" + return self.global_angular_velocity[..., 0, :] + + @classmethod + def from_state_vector_and_velocity( + cls, + skeleton_tree, + state_vector, + global_velocity, + global_angular_velocity, + is_local, + fps, + ): + """ + Construct a skeleton motion from a skeleton state vector, global velocity and angular + velocity at each joint. + + :param skeleton_tree: the skeleton tree that the motion is based on + :type skeleton_tree: SkeletonTree + :param state_vector: the state vector from the skeleton state by `.tensor` + :type state_vector: Tensor + :param global_velocity: the global velocity at each joint + :type global_velocity: Tensor + :param global_angular_velocity: the global angular velocity at each joint + :type global_angular_velocity: Tensor + :param is_local: if the rotation ins the state vector is given in local frame + :type is_local: boolean + :param fps: number of frames per second + :type fps: int + + :rtype: SkeletonMotion + """ + state_shape = state_vector.shape[:-1] + v = global_velocity.reshape(*(state_shape + (-1,))) + av = global_angular_velocity.reshape(*(state_shape + (-1,))) + new_state_vector = torch.cat([state_vector, v, av], axis=-1) + return cls( + new_state_vector, + skeleton_tree=skeleton_tree, + is_local=is_local, + fps=fps, + ) + + @classmethod + def from_skeleton_state(cls: type["SkeletonMotion"], skeleton_state: SkeletonState, fps: int): + """ + Construct a skeleton motion from a skeleton state. The velocities are estimated using second + order guassian filter along the last axis. The skeleton state must have at least .dim >= 1 + + :param skeleton_state: the skeleton state that the motion is based on + :type skeleton_state: SkeletonState + :param fps: number of frames per second + :type fps: int + + :rtype: SkeletonMotion + """ + assert ( + type(skeleton_state) == SkeletonState + ), f"expected type of {SkeletonState}, got {type(skeleton_state)}" + global_velocity = SkeletonMotion._compute_velocity( + p=skeleton_state.global_translation, time_delta=1 / fps + ) + global_angular_velocity = SkeletonMotion._compute_angular_velocity( + r=skeleton_state.global_rotation, time_delta=1 / fps + ) + return cls.from_state_vector_and_velocity( + skeleton_tree=skeleton_state.skeleton_tree, + state_vector=skeleton_state.tensor, + global_velocity=global_velocity, + global_angular_velocity=global_angular_velocity, + is_local=skeleton_state.is_local, + fps=fps, + ) + + @staticmethod + def _to_state_vector(rot, rt, vel, avel): + state_shape = rot.shape[:-2] + skeleton_state_v = SkeletonState._to_state_vector(rot, rt) + v = vel.reshape(*(state_shape + (-1,))) + av = avel.reshape(*(state_shape + (-1,))) + skeleton_motion_v = torch.cat([skeleton_state_v, v, av], axis=-1) + return skeleton_motion_v + + @classmethod + def from_dict( + cls: type["SkeletonMotion"], dict_repr: OrderedDict, *args, **kwargs + ) -> "SkeletonMotion": + rot = TensorUtils.from_dict(dict_repr["rotation"], *args, **kwargs) + rt = TensorUtils.from_dict(dict_repr["root_translation"], *args, **kwargs) + vel = TensorUtils.from_dict(dict_repr["global_velocity"], *args, **kwargs) + avel = TensorUtils.from_dict(dict_repr["global_angular_velocity"], *args, **kwargs) + return cls( + SkeletonMotion._to_state_vector(rot, rt, vel, avel), + skeleton_tree=SkeletonTree.from_dict(dict_repr["skeleton_tree"], *args, **kwargs), + is_local=dict_repr["is_local"], + fps=dict_repr["fps"], + ) + + def to_dict(self) -> OrderedDict: + return OrderedDict( + [ + ("rotation", tensor_to_dict(self.rotation)), + ("root_translation", tensor_to_dict(self.root_translation)), + ("global_velocity", tensor_to_dict(self.global_velocity)), + ( + "global_angular_velocity", + tensor_to_dict(self.global_angular_velocity), + ), + ("skeleton_tree", self.skeleton_tree.to_dict()), + ("is_local", self.is_local), + ("fps", self.fps), + ] + ) + + # @classmethod + # def from_fbx( + # cls: type["SkeletonMotion"], + # fbx_file_path, + # fbx_configs, + # skeleton_tree=None, + # is_local=True, + # fps=120, + # root_joint="", + # root_trans_index=0, + # *args, + # **kwargs, + # ) -> "SkeletonMotion": + # """ + # Construct a skeleton motion from a fbx file (TODO - generalize this). If the skeleton tree + # is not given, it will use the first frame of the mocap to construct the skeleton tree. + + # :param fbx_file_path: the path of the fbx file + # :type fbx_file_path: string + # :param fbx_configs: the configuration in terms of {"tmp_path": ..., "fbx_py27_path": ...} + # :type fbx_configs: dict + # :param skeleton_tree: the optional skeleton tree that the rotation will be applied to + # :type skeleton_tree: SkeletonTree, optional + # :param is_local: the state vector uses local or global rotation as the representation + # :type is_local: bool, optional, default=True + # :rtype: SkeletonMotion + # """ + # joint_names, joint_parents, transforms, fps = fbx_to_array(fbx_file_path, fbx_configs, root_joint, fps) + # # swap the last two axis to match the convention + # local_transform = euclidean_to_transform(transformation_matrix=torch.from_numpy(np.swapaxes(np.array(transforms), -1, -2),).float()) + # local_rotation = transform_rotation(local_transform) + # root_translation = transform_translation(local_transform)[..., root_trans_index, :] + # joint_parents = torch.from_numpy(np.array(joint_parents)).int() + + # if skeleton_tree is None: + # local_translation = transform_translation(local_transform).reshape(-1, len(joint_parents), 3)[0] + # skeleton_tree = SkeletonTree(joint_names, joint_parents, local_translation) + # skeleton_state = SkeletonState.from_rotation_and_root_translation(skeleton_tree, r=local_rotation, t=root_translation, is_local=True) + # if not is_local: + # skeleton_state = skeleton_state.global_repr() + # return cls.from_skeleton_state(skeleton_state=skeleton_state, fps=fps) + + @staticmethod + def _compute_velocity(p, time_delta, guassian_filter=True): + velocity = np.gradient(p.numpy(), axis=-3) / time_delta + if guassian_filter: + velocity = torch.from_numpy( + filters.gaussian_filter1d(velocity, 2, axis=-3, mode="nearest") + ).to(p) + else: + velocity = torch.from_numpy(velocity).to(p) + + return velocity + + @staticmethod + def _compute_angular_velocity(r, time_delta: float, guassian_filter=True): + # assume the second last dimension is the time axis + diff_quat_data = quat_identity_like(r).to(r) + diff_quat_data[..., :-1, :, :] = quat_mul_norm( + r[..., 1:, :, :], quat_inverse(r[..., :-1, :, :]) + ) + diff_angle, diff_axis = quat_angle_axis(diff_quat_data) + angular_velocity = diff_axis * diff_angle.unsqueeze(-1) / time_delta + if guassian_filter: + angular_velocity = torch.from_numpy( + filters.gaussian_filter1d(angular_velocity.numpy(), 2, axis=-3, mode="nearest"), + ) + return angular_velocity + + def crop(self, start: int, end: int, fps: int | None = None): + """ + Crop the motion along its last axis. This is equivalent to performing a slicing on the + object with [..., start: end: skip_every] where skip_every = old_fps / fps. Note that the + new fps provided must be a factor of the original fps. + + :param start: the beginning frame index + :type start: int + :param end: the ending frame index + :type end: int + :param fps: number of frames per second in the output (if not given the original fps will be used) + :type fps: int, optional + :rtype: SkeletonMotion + """ + if fps is None: + new_fps = int(self.fps) + old_fps = int(self.fps) + else: + new_fps = int(fps) + old_fps = int(self.fps) + assert old_fps % fps == 0, ( + "the resampling doesn't support fps with non-integer division " + f"from the original fps: {old_fps} => {fps}" + ) + skip_every = old_fps // new_fps + s = slice(start, end, skip_every) + z = self[..., s] + + rot = z.local_rotation if z.is_local else z.global_rotation + rt = z.root_translation + vel = z.global_velocity + avel = z.global_angular_velocity + return SkeletonMotion( + SkeletonMotion._to_state_vector(rot, rt, vel, avel), + skeleton_tree=z.skeleton_tree, + is_local=z.is_local, + fps=new_fps, + ) + + def retarget_to( + self, + joint_mapping: dict[str, str], + source_tpose_local_rotation, + source_tpose_root_translation: np.ndarray, + target_skeleton_tree: "SkeletonTree", + target_tpose_local_rotation, + target_tpose_root_translation: np.ndarray, + rotation_to_target_skeleton, + scale_to_target_skeleton: float, + z_up: bool = True, + ) -> "SkeletonMotion": + """ + Same as the one in :class:`SkeletonState`. This method discards all velocity information before + retargeting and re-estimate the velocity after the retargeting. The same fps is used in the + new retargetted motion. + + :param joint_mapping: a dictionary of that maps the joint node from the source skeleton to \ + the target skeleton + :type joint_mapping: dict[str, str] + + :param source_tpose_local_rotation: the local rotation of the source skeleton + :type source_tpose_local_rotation: Tensor + + :param source_tpose_root_translation: the root translation of the source tpose + :type source_tpose_root_translation: np.ndarray + + :param target_skeleton_tree: the target skeleton tree + :type target_skeleton_tree: SkeletonTree + + :param target_tpose_local_rotation: the local rotation of the target skeleton + :type target_tpose_local_rotation: Tensor + + :param target_tpose_root_translation: the root translation of the target tpose + :type target_tpose_root_translation: Tensor + + :param rotation_to_target_skeleton: the rotation that needs to be applied to the source\ + skeleton to align with the target skeleton. Essentially the rotation is t_R_s, where t is\ + the frame of reference of the target skeleton and s is the frame of reference of the source\ + skeleton + :type rotation_to_target_skeleton: Tensor + :param scale_to_target_skeleton: the factor that needs to be multiplied from source\ + skeleton to target skeleton (unit in distance). For example, to go from `cm` to `m`, the \ + factor needs to be 0.01. + :type scale_to_target_skeleton: float + :rtype: SkeletonMotion + """ + return SkeletonMotion.from_skeleton_state( + super().retarget_to( + joint_mapping, + source_tpose_local_rotation, + source_tpose_root_translation, + target_skeleton_tree, + target_tpose_local_rotation, + target_tpose_root_translation, + rotation_to_target_skeleton, + scale_to_target_skeleton, + z_up, + ), + self.fps, + ) + + def retarget_to_by_tpose( + self, + joint_mapping: dict[str, str], + source_tpose: "SkeletonState", + target_tpose: "SkeletonState", + rotation_to_target_skeleton, + scale_to_target_skeleton: float, + z_up: bool = True, + ) -> "SkeletonMotion": + """ + Same as the one in :class:`SkeletonState`. This method discards all velocity information before + retargeting and re-estimate the velocity after the retargeting. The same fps is used in the + new retargetted motion. + + :param joint_mapping: a dictionary of that maps the joint node from the source skeleton to \ + the target skeleton + :type joint_mapping: dict[str, str] + + :param source_tpose: t-pose of the source skeleton + :type source_tpose: SkeletonState + + :param target_tpose: t-pose of the target skeleton + :type target_tpose: SkeletonState + + :param rotation_to_target_skeleton: the rotation that needs to be applied to the source\ + skeleton to align with the target skeleton. Essentially the rotation is t_R_s, where t is\ + the frame of reference of the target skeleton and s is the frame of reference of the source\ + skeleton + :type rotation_to_target_skeleton: Tensor + :param scale_to_target_skeleton: the factor that needs to be multiplied from source\ + skeleton to target skeleton (unit in distance). For example, to go from `cm` to `m`, the \ + factor needs to be 0.01. + :type scale_to_target_skeleton: float + :rtype: SkeletonMotion + """ + return self.retarget_to( + joint_mapping, + source_tpose.local_rotation, + source_tpose.root_translation, + target_tpose.skeleton_tree, + target_tpose.local_rotation, + target_tpose.root_translation, + rotation_to_target_skeleton, + scale_to_target_skeleton, + z_up, + ) diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/torch_humanoid_batch.py b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/torch_humanoid_batch.py new file mode 100644 index 0000000000000000000000000000000000000000..858f7611e17385a6c11ebfc19801eb8f153c9f6c --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/motion_lib/torch_humanoid_batch.py @@ -0,0 +1,885 @@ +#!/usr/bin/env python3 +from collections import OrderedDict, defaultdict +import copy +from io import BytesIO +import os +import os.path as osp +from pathlib import Path +import xml.etree.ElementTree as ETree + +from easydict import EasyDict +import hydra +from loguru import logger +from lxml.etree import XMLParser, parse +import numpy as np +from omegaconf import DictConfig + +# import logging +import open3d as o3d +from rich.progress import track +import scipy.ndimage.filters as filters +from scipy.spatial.transform import Rotation as sRot +import torch + +from gear_sonic.isaac_utils.rotations import ( + axis_angle_to_quaternion, + matrix_to_quaternion, + quat_angle_axis, + quat_identity_like, + quat_inverse, + quat_mul_norm, + quaternion_to_matrix, + slerp, + wxyz_to_xyzw, +) +from gear_sonic.trl.utils.torch_transform import quaternion_to_angle_axis + +# Configure logging +# logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s') + +logger.info("Using Humanoid Batch") + + +# ============================================================================= +# Utility functions for DOF <-> rotation matrix conversion +# ============================================================================= + + +def dof_to_rotation_matrices(dof_angles: torch.Tensor, dof_axis: torch.Tensor) -> torch.Tensor: + """Convert DOF angles [..., N] to rotation matrices [..., N, 3, 3].""" + half_angles = dof_angles / 2 + cos_half, sin_half = torch.cos(half_angles), torch.sin(half_angles) + + axis = dof_axis.to(dof_angles.device) + for _ in range(dof_angles.dim() - 1): + axis = axis.unsqueeze(0) + axis = axis.expand(*dof_angles.shape, 3) + + quaternion = torch.cat([cos_half.unsqueeze(-1), sin_half.unsqueeze(-1) * axis], dim=-1) + return quaternion_to_matrix(quaternion) + + +def rotation_matrices_to_dof( + rotation_matrices: torch.Tensor, dof_axis: torch.Tensor +) -> torch.Tensor: + """Extract DOF angles [..., N] from rotation matrices [..., N, 3, 3].""" + R = rotation_matrices + x_angle = torch.atan2(R[..., 2, 1], R[..., 2, 2]) + y_angle = torch.atan2(R[..., 0, 2], R[..., 0, 0]) + z_angle = torch.atan2(R[..., 1, 0], R[..., 1, 1]) + xyz_angles = torch.stack([x_angle, y_angle, z_angle], dim=-1) + + axis = dof_axis.to(rotation_matrices.device) + for _ in range(xyz_angles.dim() - 2): + axis = axis.unsqueeze(0) + axis = axis.expand(*xyz_angles.shape[:-1], 3) + + return (xyz_angles * axis).sum(dim=-1) + + +def qpos_to_root_and_dof(qpos: torch.Tensor, num_dof: int, root_quat_wxyz: bool = True): + """Parse qpos into (root_trans, root_quat_wxyz, dof_angles). + + Args: + qpos: Joint positions tensor [..., 7 + num_dof]. + num_dof: Number of DOF angles. + root_quat_wxyz: If True, root quaternion is in wxyz order. + + Returns: + root_trans: Root translation [..., 3]. + root_quat: Root quaternion in wxyz order [..., 4]. + dof_angles: DOF angles [..., num_dof]. + """ + root_trans = qpos[..., :3] + root_quat = qpos[..., 3:7] + dof_angles = qpos[..., 7 : 7 + num_dof] + if not root_quat_wxyz: + root_quat = root_quat[..., [3, 0, 1, 2]] + return root_trans, root_quat, dof_angles + + +def root_and_dof_to_qpos(root_trans, root_quat, dof_angles, root_quat_wxyz: bool = True): + """Assemble qpos from (root_trans, root_quat_wxyz, dof_angles). + + Args: + root_trans: Root translation [..., 3]. + root_quat: Root quaternion in wxyz order [..., 4]. + dof_angles: DOF angles [..., num_dof]. + root_quat_wxyz: If True, output root quaternion in wxyz order. + + Returns: + qpos: Joint positions tensor [..., 7 + num_dof]. + """ + if not root_quat_wxyz: + root_quat = root_quat[..., [1, 2, 3, 0]] + return torch.cat([root_trans, root_quat, dof_angles], dim=-1) + + +def _compute_idx_levels(parents): + """Group joint indices by their depth level in the kinematic tree. + + This enables parallel FK by processing all joints at the same level simultaneously. + + Args: + parents: Tensor of parent indices for each joint, where -1 indicates root. + + Returns: + List of tensors, where idx_levels[i] contains indices of all joints at level i. + """ + idx_levels = [] + level_dict = {} + + for i in range(len(parents)): + parent_idx = int(parents[i]) + if parent_idx == -1: + # Root joint is at level -1 (handled separately) + level_dict[i] = -1 + else: + # Child is one level deeper than its parent + parent_level = level_dict[parent_idx] + level = parent_level + 1 + level_dict[i] = level + + # Extend idx_levels if needed + while len(idx_levels) <= level: + idx_levels.append([]) + idx_levels[level].append(i) + + # Convert to tensors + idx_levels = [torch.tensor(indices, dtype=torch.long) for indices in idx_levels] + return idx_levels + + +class Humanoid_Batch: + def __init__(self, cfg, device=torch.device("cpu")): + self.cfg = cfg + self.asset_root = Path(cfg.asset.assetRoot) + + self.asset_file = cfg.asset.assetFileName + self.mjcf_file = self.asset_root / self.asset_file + + parser = XMLParser(remove_blank_text=True) + tree = parse( + BytesIO(open(self.mjcf_file, "rb").read()), + parser=parser, + ) + self.dof_axis = [] + + joints = sorted( + [j.attrib["name"] for j in tree.getroot().find("worldbody").findall(".//joint")] + ) + motors = sorted([m.attrib["name"] for m in tree.getroot().find("actuator").getchildren()]) + + assert len(motors) > 0, "No motors found in the mjcf file" + + self.num_dof = len(motors) + self.num_extend_dof = self.num_dof + + self.mjcf_data = mjcf_data = self.from_mjcf(self.mjcf_file) + self.body_names = copy.deepcopy(mjcf_data["node_names"]) + self._parents = mjcf_data["parent_indices"] + + self.body_names_augment = copy.deepcopy(mjcf_data["node_names"]) + self._offsets = mjcf_data["local_translation"][None,].to(device) + + self._local_rotation = mjcf_data["local_rotation"][None,].to(device) + self.actuated_joints_idx = np.array( + [self.body_names.index(k) for k, v in mjcf_data["body_to_joint"].items()] + ) + + for m in motors: + if m not in joints: + print(m) + + if ( + "type" in tree.getroot().find("worldbody").findall(".//joint")[0].attrib + and tree.getroot().find("worldbody").findall(".//joint")[0].attrib["type"] == "free" + ): + for j in tree.getroot().find("worldbody").findall(".//joint")[1:]: + self.dof_axis.append([int(i) for i in j.attrib["axis"].split(" ")]) + self.has_freejoint = True + elif "type" not in tree.getroot().find("worldbody").findall(".//joint")[0].attrib: + for j in tree.getroot().find("worldbody").findall(".//joint"): + self.dof_axis.append([int(i) for i in j.attrib["axis"].split(" ")]) + self.has_freejoint = True + else: + for j in tree.getroot().find("worldbody").findall(".//joint")[6:]: + self.dof_axis.append([int(i) for i in j.attrib["axis"].split(" ")]) + self.has_freejoint = False + + self.dof_axis = torch.tensor(self.dof_axis) + + for extend_config in cfg.extend_config: + self.body_names_augment += [extend_config.joint_name] + self._parents = torch.cat( + [ + self._parents, + torch.tensor([self.body_names.index(extend_config.parent_name)]).to(device), + ], + dim=0, + ) + self._offsets = torch.cat( + [self._offsets, torch.tensor([[extend_config.pos]]).to(device)], dim=1 + ) + self._local_rotation = torch.cat( + [self._local_rotation, torch.tensor([[extend_config.rot]]).to(device)], dim=1 + ) + self.num_extend_dof += 1 + + self.num_bodies = len(self.body_names) + self.num_bodies_augment = len(self.body_names_augment) + + self.joints_range = mjcf_data["joints_range"].to(device) + self._local_rotation_mat = quaternion_to_matrix(self._local_rotation).float() # w, x, y ,z + + # Pre-compute index levels for parallel FK + self._idx_levels = _compute_idx_levels(self._parents) + + self._device = device + self.load_mesh() + + @property + def device(self): + """Return the device of the humanoid tensors.""" + return self._device + + def to(self, device): + """Move all tensors to the specified device.""" + self._device = device + self._offsets = self._offsets.to(device) + self._local_rotation = self._local_rotation.to(device) + self._local_rotation_mat = self._local_rotation_mat.to(device) + self._parents = self._parents.to(device) + self.dof_axis = self.dof_axis.to(device) + self.joints_range = self.joints_range.to(device) + return self + + def from_mjcf(self, path): + # function from Poselib: + tree = ETree.parse(path) + xml_doc_root = tree.getroot() + xml_world_body = xml_doc_root.find("worldbody") + if xml_world_body is None: + raise ValueError("MJCF parsed incorrectly please verify it.") + # assume this is the root + xml_body_root = xml_world_body.find("body") + if xml_body_root is None: + raise ValueError("MJCF parsed incorrectly please verify it.") + + xml_joint_root = xml_body_root.find("joint") + + node_names = [] + parent_indices = [] + local_translation = [] + local_rotation = [] + joints_range = [] + body_to_joint = OrderedDict() + + # recursively adding all nodes into the skel_tree + def _add_xml_node(xml_node, parent_index, node_index): + node_name = xml_node.attrib.get("name") + # parse the local translation into float list + pos = np.fromstring(xml_node.attrib.get("pos", "0 0 0"), dtype=float, sep=" ") + quat = np.fromstring(xml_node.attrib.get("quat", "1 0 0 0"), dtype=float, sep=" ") + node_names.append(node_name) + parent_indices.append(parent_index) + local_translation.append(pos) + local_rotation.append(quat) + curr_index = node_index + node_index += 1 + all_joints = xml_node.findall("joint") # joints need to remove the first 6 joints + if len(all_joints) == 6: + all_joints = all_joints[6:] + + for joint in all_joints: + if joint.attrib.get("range") is not None: + joints_range.append( + np.fromstring(joint.attrib.get("range"), dtype=float, sep=" ") + ) + else: + if not joint.attrib.get("type") == "free": + joints_range.append([-np.pi, np.pi]) + for joint_node in xml_node.findall("joint"): + body_to_joint[node_name] = joint_node.attrib.get("name") + + for next_node in xml_node.findall("body"): + node_index = _add_xml_node(next_node, curr_index, node_index) + + return node_index + + _add_xml_node(xml_body_root, -1, 0) + assert len(joints_range) == self.num_dof + return { + "node_names": node_names, + "parent_indices": torch.from_numpy(np.array(parent_indices, dtype=np.int32)), + "local_translation": torch.from_numpy(np.array(local_translation, dtype=np.float32)), + "local_rotation": torch.from_numpy(np.array(local_rotation, dtype=np.float32)), + "joints_range": torch.from_numpy(np.array(joints_range)), + "body_to_joint": body_to_joint, + } + + def _lerp(self, a: torch.Tensor, b: torch.Tensor, blend: torch.Tensor) -> torch.Tensor: + """Linear interpolation between two tensors.""" + return a * (1 - blend) + b * blend + + def _slerp(self, a: torch.Tensor, b: torch.Tensor, blend: torch.Tensor) -> torch.Tensor: + """Spherical linear interpolation between two quaternions.""" + slerped_quats = torch.zeros_like(a) + for i in range(a.shape[0]): + slerped_quats[i] = slerp(a[i], b[i], blend[i]) + return slerped_quats + + def _compute_frame_blend( + self, times: torch.Tensor, duration: torch.Tensor, input_frames: int + ) -> torch.Tensor: + """Computes the frame blend for the motion.""" + phase = times / duration + index_0 = (phase * (input_frames - 1)).floor().long() + index_1 = torch.minimum(index_0 + 1, torch.tensor(input_frames - 1).to(times.device)) + blend = phase * (input_frames - 1) - index_0 + return index_0, index_1, blend + + def interploate_pose(self, pose_quat, trans, fps, target_fps): + in_shape = trans.shape if pose_quat is None else pose_quat.shape + device = pose_quat.device if pose_quat is not None else trans.device + assert in_shape[0] == 1, "Only support single sequence for now" + duration = (in_shape[1] - 1) * 1 / fps + + times = torch.arange(0, duration, 1 / target_fps, dtype=torch.float32, device=device) + index_0, index_1, blend = self._compute_frame_blend( + times=times, duration=duration, input_frames=in_shape[1] + ) + if pose_quat is not None: + pose_quat = self._slerp(pose_quat[0, index_0], pose_quat[0, index_1], blend) + pose_quat = pose_quat.unsqueeze(0) + if trans is not None: + trans = self._lerp(trans[0, index_0], trans[0, index_1], blend.unsqueeze(1)) + trans = trans.unsqueeze(0) + return pose_quat, trans + + def fk_batch( + self, + pose, + trans, + return_full=False, + fps=30, + target_fps=50, + interpolate_data=False, + use_parallel_fk: bool = False, + ): + device, dtype = pose.device, pose.dtype + + B, seq_len = pose.shape[:2] + + pose = pose[..., : len(self._parents), :] # G1 fitted joints might have extra joints + + if ( + self.num_bodies_augment > 0 + and pose.shape[2] < self.num_bodies + self.num_bodies_augment + ): + pose = torch.cat( + [ + pose, + torch.zeros( + B, seq_len, self.num_bodies_augment - pose.shape[2], pose.shape[3] + ).to(device), + ], + dim=2, + ) + + pose_quat = axis_angle_to_quaternion(pose.clone()) + if interpolate_data and fps != target_fps: + pose_quat, trans = self.interploate_pose(pose_quat, trans, fps, target_fps) + dt = 1 / target_fps + else: + dt = 1 / fps + pose = quaternion_to_angle_axis(pose_quat) + B, seq_len = pose.shape[:2] + pose_mat = quaternion_to_matrix(pose_quat) + + if pose_mat.shape != 5: + pose_mat = pose_mat.reshape(B, seq_len, -1, 3, 3) + J = pose_mat.shape[2] - 1 # Exclude root + wbody_pos, wbody_mat = self.forward_kinematics_batch( + pose_mat[:, :, 1:], pose_mat[:, :, 0:1], trans, use_parallel_fk=use_parallel_fk + ) + + return_dict = EasyDict() + + wbody_rot = wxyz_to_xyzw(matrix_to_quaternion(wbody_mat)) + if len(self.cfg.extend_config) > 0: + if return_full: + return_dict.global_velocity_extend = self._compute_velocity(wbody_pos, dt) + return_dict.global_angular_velocity_extend = self._compute_angular_velocity( + wbody_rot, dt + ) + + return_dict.global_translation_extend = wbody_pos.clone() + return_dict.global_rotation_mat_extend = wbody_mat.clone() + return_dict.global_rotation_extend = wbody_rot + + wbody_pos = wbody_pos[..., : self.num_bodies, :] + wbody_mat = wbody_mat[..., : self.num_bodies, :, :] + wbody_rot = wbody_rot[..., : self.num_bodies, :] + + return_dict.global_translation = wbody_pos + return_dict.global_rotation_mat = wbody_mat + return_dict.global_rotation = wbody_rot + if return_full: + rigidbody_linear_velocity = self._compute_velocity( + wbody_pos, dt + ) # Isaac gym is [x, y, z, w]. All the previous functions are [w, x, y, z] + rigidbody_angular_velocity = self._compute_angular_velocity(wbody_rot, dt) + return_dict.local_rotation = wxyz_to_xyzw(pose_quat) + return_dict.global_root_velocity = rigidbody_linear_velocity[..., 0, :] + return_dict.global_root_angular_velocity = rigidbody_angular_velocity[..., 0, :] + return_dict.global_angular_velocity = rigidbody_angular_velocity + return_dict.global_velocity = rigidbody_linear_velocity + + if len(self.cfg.extend_config) > 0: + return_dict.dof_pos = pose.sum(dim=-1)[ + ..., 1 : self.num_bodies + ] # you can sum it up since unitree's each joint has 1 dof. Last two are for hands. doesn't really matter. + else: + if not len(self.actuated_joints_idx) == len(self.body_names): + return_dict.dof_pos = pose.sum(dim=-1)[..., self.actuated_joints_idx] + else: + return_dict.dof_pos = pose.sum(dim=-1)[..., 1:] + + dof_vel = (return_dict.dof_pos[:, 1:] - return_dict.dof_pos[:, :-1]) / dt + return_dict.dof_vels = torch.cat([dof_vel, dof_vel[:, -2:-1]], dim=1) + return_dict.fps = int(1 / dt) + + return return_dict + + def forward_kinematics_batch( + self, rotations, root_rotations, root_positions, use_parallel_fk: bool = False + ): + """ + Perform forward kinematics using the given trajectory and local rotations. + + Arguments (where B = batch size, T = sequence length, J = number of joints): + rotations: (B, T, J-1, 3, 3) tensor of rotation matrices for non-root joints. + root_rotations: (B, T, 1, 3, 3) tensor of root rotation matrix. + root_positions: (B, T, 3) tensor describing the root joint positions. + use_parallel_fk: If True, use level-wise parallel FK (faster for large batches). + If False, use sequential FK (original implementation). + + Output: + positions_world: (B, T, J, 3) world positions of all joints + rotations_world: (B, T, J, 3, 3) world rotation matrices of all joints + """ + device, dtype = root_rotations.device, root_rotations.dtype + B, seq_len = rotations.size()[0:2] + J = self._offsets.shape[1] + + expanded_offsets = self._offsets[:, None].expand(B, seq_len, J, 3).to(device).type(dtype) + + if use_parallel_fk: + # Initialize transforms + eye = ( + torch.eye(3, device=device, dtype=dtype) + .view(1, 1, 1, 3, 3) + .expand(B, seq_len, 1, 3, 3) + ) + local_rot_mat = self._local_rotation_mat.to(device, dtype).unsqueeze( + 1 + ) # [1, 1, J, 3, 3] + local_transforms = torch.matmul( + local_rot_mat, torch.cat([eye, rotations], dim=2) + ) # [B, T, J, 3, 3] + + positions_world = torch.zeros(B, seq_len, J, 3, device=device, dtype=dtype) + rotations_world = torch.zeros(B, seq_len, J, 3, 3, device=device, dtype=dtype) + root_idx = (self._parents == -1).nonzero(as_tuple=True)[0].item() + positions_world[:, :, root_idx] = root_positions + rotations_world[:, :, root_idx] = root_rotations[:, :, 0] + + # Process level by level (all joints at same depth in parallel) + for level_indices in self._idx_levels: + if len(level_indices) == 0: + continue + + level_indices = level_indices.to(device) + parent_indices = self._parents[level_indices].long().to(device) + + parent_pos = positions_world[:, :, parent_indices] # [B, T, L, 3] + parent_rot = rotations_world[:, :, parent_indices] # [B, T, L, 3, 3] + + local_rot = local_transforms[:, :, level_indices] # [B, T, L, 3, 3] + offsets = expanded_offsets[:, :, level_indices] # [B, T, L, 3] + world_pos = parent_pos + torch.matmul(parent_rot, offsets.unsqueeze(-1)).squeeze(-1) + world_rot = torch.matmul(parent_rot, local_rot) + + positions_world[:, :, level_indices] = world_pos + rotations_world[:, :, level_indices] = world_rot + + else: + # for loop version, should be deprecated but kept here for compatibility + positions_world = [] + rotations_world = [] + + for i in range( + J + ): # Tingwu: this will be super slow; should do parallel forward kinematics instead + if self._parents[i] == -1: + positions_world.append(root_positions) + rotations_world.append(root_rotations) + else: + try: + jpos = ( + torch.matmul( + rotations_world[self._parents[i]][:, :, 0], + expanded_offsets[:, :, i, :, None], + ).squeeze(-1) + + positions_world[self._parents[i]] + ) + rot_mat = torch.matmul( + rotations_world[self._parents[i]], + torch.matmul( + self._local_rotation_mat[:, (i) : (i + 1)], + rotations[:, :, (i - 1) : i, :], + ), + ) + except Exception as e: + logger.error(f"Error at joint index {i}") + logger.error(f"Parent index: {self._parents[i]}") + logger.error(f"Error details: {str(e)}") + + positions_world.append(jpos) + rotations_world.append(rot_mat) + + positions_world = torch.stack(positions_world, dim=2) + rotations_world = torch.cat(rotations_world, dim=2) + + return positions_world, rotations_world + + def global_to_local_rotations(self, global_rotations: torch.Tensor) -> torch.Tensor: + """Convert global rotations to local rotations. + + Args: + global_rotations: Global rotation matrices [..., J, 3, 3]. + + Returns: + local_rotations: Local rotation matrices [..., J, 3, 3]. + """ + parents = self._parents[: global_rotations.shape[-3]] + root_mask = parents == -1 # [J] + parent_indices = parents.clone() + parent_indices[root_mask] = 0 # placeholder index for gathering (will be overwritten) + + # Gather parent rotations for all joints: [..., J, 3, 3] + parent_rot = global_rotations[..., parent_indices.long(), :, :] + + # local = parent^T @ global + local_rotations = torch.matmul(parent_rot.transpose(-1, -2), global_rotations) + + # Root joints: local = global (overwrite) + if root_mask.any(): + local_rotations[..., root_mask, :, :] = global_rotations[..., root_mask, :, :] + + return local_rotations + + def qpos_to_global_transforms( + self, + qpos: torch.Tensor, + root_quat_wxyz: bool = True, + include_extended: bool = False, + use_parallel_fk: bool = True, + ): + """Convert qpos to global positions and rotations. + + Args: + qpos: Joint positions tensor [..., D]. + Format: [root_trans(3), root_quat(4), dof_angles(N)]. + Supports arbitrary leading batch dimensions. + root_quat_wxyz: If True, root quaternion is in wxyz order. + include_extended: If True, include extended bodies in output. + use_parallel_fk: If True, use level-wise parallel FK. + + Returns: + global_pos: Global positions [..., J, 3]. + global_rot: Global rotation matrices [..., J, 3, 3]. + """ + # Flatten arbitrary leading dims into [flat_B, 1, D] for FK + orig_shape = qpos.shape[:-1] # e.g., (B,), (B, T), (B, T, N), ... + D = qpos.shape[-1] + qpos = qpos.reshape(-1, 1, D) # [flat_B, 1, D] + + flat_B = qpos.shape[0] + root_trans, root_quat, dof_angles = qpos_to_root_and_dof(qpos, self.num_dof, root_quat_wxyz) + + root_rot_mat = quaternion_to_matrix(root_quat).unsqueeze(2) + joint_rot_mat = dof_to_rotation_matrices(dof_angles, self.dof_axis) + + # Pad with identity matrices for extended bodies + num_extended = self.num_bodies_augment - self.num_bodies + if num_extended > 0: + eye = torch.eye(3, device=joint_rot_mat.device, dtype=joint_rot_mat.dtype) + eye = eye.view(1, 1, 1, 3, 3).expand(flat_B, 1, num_extended, 3, 3) + joint_rot_mat = torch.cat([joint_rot_mat, eye], dim=2) + + global_pos, global_rot = self.forward_kinematics_batch( + joint_rot_mat, root_rot_mat, root_trans, use_parallel_fk=use_parallel_fk + ) + + if not include_extended: + global_pos = global_pos[..., : self.num_bodies, :] + global_rot = global_rot[..., : self.num_bodies, :, :] + + # Reshape back: [flat_B, 1, J, ...] -> [*orig_shape, J, ...] + J_pos = global_pos.shape[-2] + global_pos = global_pos.squeeze(1).reshape(*orig_shape, J_pos, 3) + global_rot = global_rot.squeeze(1).reshape(*orig_shape, J_pos, 3, 3) + return global_pos, global_rot + + def global_transforms_to_qpos( + self, + global_rotations: torch.Tensor, + global_positions: torch.Tensor, + root_quat_wxyz: bool = True, + ) -> torch.Tensor: + """Convert global transforms back to qpos. + + Args: + global_rotations: Global rotation matrices [B, T, J, 3, 3] or [B, J, 3, 3]. + global_positions: Global positions [B, T, J, 3] or [B, J, 3]. + root_quat_wxyz: If True, output root quaternion in wxyz order. + + Returns: + qpos: Joint positions tensor [B, T, D] or [B, D]. + """ + squeeze_time = global_rotations.dim() == 4 + if squeeze_time: + global_rotations = global_rotations.unsqueeze(1) + global_positions = global_positions.unsqueeze(1) + + root_trans = global_positions[..., 0, :] + local_rotations = self.global_to_local_rotations(global_rotations) + + root_quat = matrix_to_quaternion(local_rotations[..., 0, :, :]) + local_rot_mat = self._local_rotation_mat.to(local_rotations.device) + joint_rot_mat = torch.matmul( + local_rot_mat[:, 1 : self.num_bodies].transpose(-1, -2), + local_rotations[..., 1 : self.num_bodies, :, :], + ) + dof_angles = rotation_matrices_to_dof(joint_rot_mat, self.dof_axis) + + qpos = root_and_dof_to_qpos(root_trans, root_quat, dof_angles, root_quat_wxyz) + + return qpos.squeeze(1) if squeeze_time else qpos + + def append_extended_transforms( + self, + base_positions: torch.Tensor, + base_rotations: torch.Tensor, + ): + """Append extended body transforms to base body transforms (global frame). + + Computes world-frame transforms for extended bodies from their parent's + global transforms and concatenates them with base transforms. + All extended joints have base-body parents, so processed in one batched pass. + + Args: + base_positions: Base body positions in global frame [..., num_bodies, 3]. + base_rotations: Base body rotation matrices in global frame [..., num_bodies, 3, 3]. + + Returns: + full_positions: All body positions [..., num_bodies_augment, 3]. + full_rotations: All body rotation matrices [..., num_bodies_augment, 3, 3]. + """ + if ( + self.num_bodies_augment - self.num_bodies == 0 + or base_positions.shape[-2] == self.num_bodies_augment + ): + return base_positions, base_rotations + assert ( + base_positions.shape[-2] == self.num_bodies + ), "Must provide the non-extended base body positions" + + device = base_positions.device + dtype = base_positions.dtype + + offsets = self._offsets[0].to(device, dtype) # [num_bodies_augment, 3] + local_rot = self._local_rotation_mat[0].to(device, dtype) # [num_bodies_augment, 3, 3] + + # All extended joints have base-body parents, so process in one pass + ext_indices = torch.arange(self.num_bodies, self.num_bodies_augment, device=device) + parent_indices = self._parents[ext_indices].long().to(device) # [K] + + # Gather parent transforms from base bodies: [..., K, 3] and [..., K, 3, 3] + parent_pos = base_positions[..., parent_indices, :] + parent_rot = base_rotations[..., parent_indices, :, :] + + ext_offsets = offsets[ext_indices] # [K, 3] + ext_local_rot = local_rot[ext_indices] # [K, 3, 3] + + # child_pos = parent_pos + parent_rot @ offset + # child_rot = parent_rot @ local_rot + ext_positions = parent_pos + torch.matmul(parent_rot, ext_offsets.unsqueeze(-1)).squeeze(-1) + ext_rotations = torch.matmul(parent_rot, ext_local_rot) + + full_positions = torch.cat([base_positions, ext_positions], dim=-2) + full_rotations = torch.cat([base_rotations, ext_rotations], dim=-3) + + return full_positions, full_rotations + + @staticmethod + def _compute_velocity(p, time_delta, guassian_filter=True): + velocity = np.gradient(p.numpy(), axis=-3) / time_delta + if guassian_filter: + velocity = torch.from_numpy( + filters.gaussian_filter1d(velocity, 2, axis=-3, mode="nearest") + ).to(p) + else: + velocity = torch.from_numpy(velocity).to(p) + + return velocity + + @staticmethod + def _compute_angular_velocity(r, time_delta: float, guassian_filter=True): + # assume the second last dimension is the time axis + diff_quat_data = quat_identity_like(r).to(r) + diff_quat_data[..., :-1, :, :] = quat_mul_norm( + r[..., 1:, :, :], quat_inverse(r[..., :-1, :, :], w_last=True), w_last=True + ) + diff_angle, diff_axis = quat_angle_axis(diff_quat_data, w_last=True) + angular_velocity = diff_axis * diff_angle.unsqueeze(-1) / time_delta + if guassian_filter: + angular_velocity = torch.from_numpy( + filters.gaussian_filter1d(angular_velocity.numpy(), 2, axis=-3, mode="nearest"), + ) + return angular_velocity + + def load_mesh(self): + xml_base = os.path.dirname(self.mjcf_file) + # Read the compiler tag from the g1.xml file to find if there is a meshdir defined + tree = ETree.parse(self.mjcf_file) + xml_doc_root = tree.getroot() + compiler_tag = xml_doc_root.find("compiler") + + if compiler_tag is not None and "meshdir" in compiler_tag.attrib: + mesh_base = os.path.join(xml_base, compiler_tag.attrib["meshdir"]) + else: + mesh_base = xml_base + + self.tree = tree = ETree.parse(self.mjcf_file) + xml_doc_root = tree.getroot() + xml_world_body = xml_doc_root.find("worldbody") + + xml_assets = xml_doc_root.find("asset") + all_mesh = xml_assets.findall(".//mesh") + + geoms = xml_world_body.findall(".//geom") + + all_joints = xml_world_body.findall(".//joint") + all_motors = tree.findall(".//motor") + all_bodies = xml_world_body.findall(".//body") + + def find_parent(root, child): + for parent in root.iter(): + for elem in parent: + if elem == child: + return parent + return None + + mesh_dict = {} + mesh_parent_dict = {} + + for mesh_file_node in track(all_mesh, description="Loading Meshes ..."): + mesh_name = mesh_file_node.attrib["name"] + mesh_file = mesh_file_node.attrib["file"] + mesh_full_file = osp.join(mesh_base, mesh_file) + mesh_obj = o3d.io.read_triangle_mesh(mesh_full_file) + mesh_dict[mesh_name] = mesh_obj + + geom_transform = {} + + body_to_mesh = defaultdict(set) + mesh_to_body = {} + for geom_node in track(geoms, description="Loading Geoms..."): + if "mesh" in geom_node.attrib: + parent = find_parent(xml_doc_root, geom_node) + body_to_mesh[parent.attrib["name"]].add(geom_node.attrib["mesh"]) + mesh_to_body[geom_node] = parent + if "pos" in geom_node.attrib or "quat" in geom_node.attrib: + geom_transform[parent.attrib["name"]] = {} + geom_transform[parent.attrib["name"]]["pos"] = np.array([0.0, 0.0, 0.0]) + geom_transform[parent.attrib["name"]]["quat"] = np.array([1.0, 0.0, 0.0, 0.0]) + if "pos" in geom_node.attrib: + geom_transform[parent.attrib["name"]]["pos"] = np.array( + [float(f) for f in geom_node.attrib["pos"].split(" ")] + ) + if "quat" in geom_node.attrib: + geom_transform[parent.attrib["name"]]["quat"] = np.array( + [float(f) for f in geom_node.attrib["quat"].split(" ")] + ) + + else: + pass + + self.geom_transform = geom_transform + self.mesh_dict = mesh_dict + self.body_to_mesh = body_to_mesh + self.mesh_to_body = mesh_to_body + + def mesh_fk(self, pose=None, trans=None): + """ + Load the mesh from the XML file and merge them into the humanoid based on the current pose. + """ + if pose is None: + fk_res = self.fk_batch( + torch.zeros(1, 1, len(self.body_names_augment), 3), torch.zeros(1, 1, 3) + ) + else: + fk_res = self.fk_batch(pose, trans) + + g_trans = fk_res.global_translation.squeeze() + g_rot = fk_res.global_rotation_mat.squeeze() + geoms = self.tree.find("worldbody").findall(".//geom") + joined_mesh_obj = [] + for geom in geoms: + if "mesh" not in geom.attrib: + continue + parent_name = geom.attrib["mesh"] + + k = self.mesh_to_body[geom].attrib["name"] + mesh_names = self.body_to_mesh[k] + body_idx = self.body_names.index(k) + + body_trans = g_trans[body_idx].numpy().copy() + body_rot = g_rot[body_idx].numpy().copy() + for mesh_name in mesh_names: + mesh_obj = copy.deepcopy(self.mesh_dict[mesh_name]) + if k in self.geom_transform: + pos = self.geom_transform[k]["pos"] + quat = self.geom_transform[k]["quat"] + body_trans = body_trans + body_rot @ pos + global_rot = (body_rot @ sRot.from_quat(quat[[1, 2, 3, 0]]).as_matrix()).T + else: + global_rot = body_rot.T + mesh_obj.rotate(global_rot.T, center=(0, 0, 0)) + mesh_obj.translate(body_trans) + joined_mesh_obj.append(mesh_obj) + + # Merge all meshes into a single mesh + merged_mesh = joined_mesh_obj[0] + for mesh in joined_mesh_obj[1:]: + merged_mesh += mesh + + # Save the merged mesh to a file + + # merged_mesh.compute_vertex_normals() # Debugging + # o3d.io.write_triangle_mesh(f"data/combined_{self.cfg.humanoid_type}.stl", merged_mesh) + return merged_mesh + + +@hydra.main(version_base=None, config_path="../../config", config_name="base") +def main(config: DictConfig): + device = torch.device("cpu") + humanoid_fk = Humanoid_Batch(config.robot.motion, device) + humanoid_fk.mesh_fk() + + +if __name__ == "__main__": + main() diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/__init__.py b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/base_sim.py b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/base_sim.py new file mode 100644 index 0000000000000000000000000000000000000000..fe8f4a0852818922d2fc01da61dffc671fa2ef66 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/base_sim.py @@ -0,0 +1,659 @@ +"""MuJoCo simulation environment and loop for the G1 (and H1) humanoid robots. + +DefaultEnv owns the MuJoCo model/data, computes PD torques from Unitree SDK +commands, steps physics, and publishes observations back via the SDK bridge. +BaseSimulator wraps DefaultEnv with rate-limiting and viewer/image update loops. +""" + +import os +import pathlib +from pathlib import Path +import pickle +import tempfile +from threading import Lock, Thread +import time +from typing import Dict +import xml.etree.ElementTree as ET + +import mujoco +import mujoco.viewer +import numpy as np +from scipy.spatial.transform import Rotation +from unitree_sdk2py.core.channel import ChannelFactoryInitialize + +from gear_sonic.utils.mujoco_sim.metric_utils import check_contact, check_height +from gear_sonic.utils.mujoco_sim.sim_utils import get_subtree_body_names +from gear_sonic.utils.mujoco_sim.unitree_sdk2py_bridge import ElasticBand, UnitreeSdk2Bridge +from gear_sonic.utils.mujoco_sim.robot import Robot + +GEAR_SONIC_ROOT = Path(__file__).resolve().parent.parent.parent.parent + + +class DefaultEnv: + """Base environment class that handles simulation environment setup and step""" + + def __init__( + self, + config: Dict[str, any], + env_name: str = "default", + camera_configs: Dict[str, any] = {}, + onscreen: bool = False, + offscreen: bool = False, + enable_image_publish: bool = False, + ): + self.config = config + self.env_name = env_name + self.robot = Robot(self.config) + self.num_body_dof = self.robot.NUM_JOINTS + self.num_hand_dof = self.robot.NUM_HAND_JOINTS + self.sim_dt = self.config["SIMULATE_DT"] + self.obs = None + self.torques = np.zeros(self.num_body_dof + self.num_hand_dof * 2) + self.torque_limit = np.array(self.robot.MOTOR_EFFORT_LIMIT_LIST) + self.camera_configs = camera_configs + + if not camera_configs and offscreen and enable_image_publish: + self.camera_configs = { + "ego_view": {"height": 480, "width": 640, "mjcf_name": "head_camera"}, + } + + self.reward_lock = Lock() + self.unitree_bridge = None + self.onscreen = onscreen + + self.init_scene() + self.last_reward = 0 + + self.offscreen = offscreen + if self.offscreen: + self.init_renderers() + self.image_dt = self.config.get("IMAGE_DT", 0.033333) + self.image_publish_process = None + + def start_image_publish_subprocess(self, start_method: str = "spawn", camera_port: int = 5555): + from gear_sonic.utils.mujoco_sim.image_publish_utils import ImagePublishProcess + + if len(self.camera_configs) == 0: + print( + "Warning: No camera configs provided, image publishing subprocess will not be started" + ) + return + start_method = self.config.get("MP_START_METHOD", "spawn") + self.image_publish_process = ImagePublishProcess( + camera_configs=self.camera_configs, + image_dt=self.image_dt, + zmq_port=camera_port, + start_method=start_method, + verbose=self.config.get("verbose", False), + ) + self.image_publish_process.start_process() + + def _get_dof_indices_by_class(self): + with tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".xml") as f: + mujoco.mj_saveLastXML(f.name, self.mj_model) + temp_xml_path = f.name + + try: + tree = ET.parse(temp_xml_path) + root = tree.getroot() + + joint_class_map = {} + for joint_element in root.findall(".//joint[@class]"): + joint_name = joint_element.get("name") + joint_class = joint_element.get("class") + if joint_name and joint_class: + joint_id = mujoco.mj_name2id( + self.mj_model, mujoco.mjtObj.mjOBJ_JOINT, joint_name + ) + if joint_id != -1: + dof_adr = self.mj_model.jnt_dofadr[joint_id] + if joint_class not in joint_class_map: + joint_class_map[joint_class] = [] + joint_class_map[joint_class].append(dof_adr) + finally: + os.remove(temp_xml_path) + + return joint_class_map + + def _get_default_dof_properties(self): + with tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".xml") as f: + mujoco.mj_saveLastXML(f.name, self.mj_model) + temp_xml_path = f.name + + try: + tree = ET.parse(temp_xml_path) + root = tree.getroot() + + default_dof_properties = {} + for default_element in root.findall(".//default/default[@class]"): + class_name = default_element.get("class") + joint_element = default_element.find("joint") + if class_name and joint_element is not None: + properties = {} + if "damping" in joint_element.attrib: + properties["damping"] = float(joint_element.get("damping")) + if "armature" in joint_element.attrib: + properties["armature"] = float(joint_element.get("armature")) + if "frictionloss" in joint_element.attrib: + properties["frictionloss"] = float(joint_element.get("frictionloss")) + + if properties: + default_dof_properties[class_name] = properties + finally: + os.remove(temp_xml_path) + + return default_dof_properties + + def init_scene(self): + """Initialize the default robot scene""" + xml_path = str(pathlib.Path(GEAR_SONIC_ROOT) / self.config["ROBOT_SCENE"]) + self.mj_model = mujoco.MjModel.from_xml_path(xml_path) + self.mj_data = mujoco.MjData(self.mj_model) + self.mj_model.opt.timestep = self.sim_dt + self.torso_index = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, "torso_link") + self.root_body = "pelvis" + self.root_body_id = self.mj_model.body(self.root_body).id + + self.joint_class_map = self._get_dof_indices_by_class() + + self.perform_sysid_search = self.config.get("perform_sysid_search", False) + + # Check for static root link (fixed base) + self.use_floating_root_link = "floating_base_joint" in [ + self.mj_model.joint(i).name for i in range(self.mj_model.njnt) + ] + self.use_constrained_root_link = "constrained_base_joint" in [ + self.mj_model.joint(i).name for i in range(self.mj_model.njnt) + ] + + # MuJoCo qpos/qvel arrays start with root DOFs before joint DOFs: + # floating base has 7 qpos (pos + quat) and 6 qvel (lin + ang velocity) + if self.use_floating_root_link: + self.qpos_offset = 7 + self.qvel_offset = 6 + else: + if self.use_constrained_root_link: + self.qpos_offset = 1 + self.qvel_offset = 1 + else: + raise ValueError( + "No root link found --" + "The absolute static root will make the simulation unstable." + ) + + # Enable the elastic band + if self.config["ENABLE_ELASTIC_BAND"] and self.use_floating_root_link: + self.elastic_band = ElasticBand() + if "g1" in self.config["ROBOT_TYPE"]: + if self.config["enable_waist"]: + self.band_attached_link = self.mj_model.body("pelvis").id + else: + self.band_attached_link = self.mj_model.body("torso_link").id + elif "h1" in self.config["ROBOT_TYPE"]: + self.band_attached_link = self.mj_model.body("torso_link").id + else: + self.band_attached_link = self.mj_model.body("base_link").id + + if self.onscreen: + self.viewer = mujoco.viewer.launch_passive( + self.mj_model, + self.mj_data, + key_callback=self.elastic_band.MujuocoKeyCallback, + show_left_ui=False, + show_right_ui=False, + ) + else: + mujoco.mj_forward(self.mj_model, self.mj_data) + self.viewer = None + else: + if self.onscreen: + self.viewer = mujoco.viewer.launch_passive( + self.mj_model, self.mj_data, show_left_ui=False, show_right_ui=False + ) + else: + mujoco.mj_forward(self.mj_model, self.mj_data) + self.viewer = None + + if self.viewer: + self.viewer.cam.azimuth = 120 + self.viewer.cam.elevation = -30 + self.viewer.cam.distance = 2.0 + self.viewer.cam.lookat = np.array([0, 0, 0.5]) + self.viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING + self.viewer.cam.trackbodyid = self.mj_model.body("pelvis").id + + self.body_joint_index = [] + self.left_hand_index = [] + self.right_hand_index = [] + for i in range(self.mj_model.njnt): + name = self.mj_model.joint(i).name + if any( + [ + part_name in name + for part_name in ["hip", "knee", "ankle", "waist", "shoulder", "elbow", "wrist"] + ] + ): + self.body_joint_index.append(i) + elif "left_hand" in name: + self.left_hand_index.append(i) + elif "right_hand" in name: + self.right_hand_index.append(i) + + assert len(self.body_joint_index) == self.robot.NUM_JOINTS + assert len(self.left_hand_index) == self.robot.NUM_HAND_JOINTS + assert len(self.right_hand_index) == self.robot.NUM_HAND_JOINTS + + self.body_joint_index = np.array(self.body_joint_index) + self.left_hand_index = np.array(self.left_hand_index) + self.right_hand_index = np.array(self.right_hand_index) + + def init_renderers(self): + self.renderers = {} + for camera_name, camera_config in self.camera_configs.items(): + renderer = mujoco.Renderer( + self.mj_model, height=camera_config["height"], width=camera_config["width"] + ) + self.renderers[camera_name] = renderer + + def compute_body_torques(self) -> np.ndarray: + # PD control: tau = tau_ff + kp * (q_des - q) + kd * (dq_des - dq) + body_torques = np.zeros(self.num_body_dof) + if self.unitree_bridge is not None and self.unitree_bridge.low_cmd: + for i in range(self.unitree_bridge.num_body_motor): + if self.unitree_bridge.use_sensor: + body_torques[i] = ( + self.unitree_bridge.low_cmd.motor_cmd[i].tau + + self.unitree_bridge.low_cmd.motor_cmd[i].kp + * (self.unitree_bridge.low_cmd.motor_cmd[i].q - self.mj_data.sensordata[i]) + + self.unitree_bridge.low_cmd.motor_cmd[i].kd + * ( + self.unitree_bridge.low_cmd.motor_cmd[i].dq + - self.mj_data.sensordata[i + self.unitree_bridge.num_body_motor] + ) + ) + else: + body_torques[i] = ( + self.unitree_bridge.low_cmd.motor_cmd[i].tau + + self.unitree_bridge.low_cmd.motor_cmd[i].kp + * ( + self.unitree_bridge.low_cmd.motor_cmd[i].q + - self.mj_data.qpos[self.body_joint_index[i] + self.qpos_offset - 1] + ) + + self.unitree_bridge.low_cmd.motor_cmd[i].kd + * ( + self.unitree_bridge.low_cmd.motor_cmd[i].dq + - self.mj_data.qvel[self.body_joint_index[i] + self.qvel_offset - 1] + ) + ) + return body_torques + + def get_head_pose(self) -> np.ndarray: + root_pos = self.mj_data.body("torso_link").xpos.copy() + # Reorder quaternion from MuJoCo [w,x,y,z] to scipy [x,y,z,w] + root_quat = self.mj_data.body("torso_link").xquat.copy()[[1, 2, 3, 0]] + head_pos = root_pos + Rotation.from_quat(root_quat).apply(np.array([0.0, 0.0, -0.044])) + return np.concatenate((head_pos, root_quat)) + + def get_root_vel(self) -> np.ndarray: + return self.mj_data.qvel[:6] + + def compute_hand_torques(self) -> np.ndarray: + left_hand_torques = np.zeros(self.num_hand_dof) + right_hand_torques = np.zeros(self.num_hand_dof) + if self.unitree_bridge is not None and self.unitree_bridge.low_cmd: + for i in range(self.unitree_bridge.num_hand_motor): + left_hand_torques[i] = ( + self.unitree_bridge.left_hand_cmd.motor_cmd[i].tau + + self.unitree_bridge.left_hand_cmd.motor_cmd[i].kp + * ( + self.unitree_bridge.left_hand_cmd.motor_cmd[i].q + - self.mj_data.qpos[self.left_hand_index[i] + self.qpos_offset - 1] + ) + + self.unitree_bridge.left_hand_cmd.motor_cmd[i].kd + * ( + self.unitree_bridge.left_hand_cmd.motor_cmd[i].dq + - self.mj_data.qvel[self.left_hand_index[i] + self.qvel_offset - 1] + ) + ) + right_hand_torques[i] = ( + self.unitree_bridge.right_hand_cmd.motor_cmd[i].tau + + self.unitree_bridge.right_hand_cmd.motor_cmd[i].kp + * ( + self.unitree_bridge.right_hand_cmd.motor_cmd[i].q + - self.mj_data.qpos[self.right_hand_index[i] + self.qpos_offset - 1] + ) + + self.unitree_bridge.right_hand_cmd.motor_cmd[i].kd + * ( + self.unitree_bridge.right_hand_cmd.motor_cmd[i].dq + - self.mj_data.qvel[self.right_hand_index[i] + self.qvel_offset - 1] + ) + ) + return np.concatenate((left_hand_torques, right_hand_torques)) + + def compute_body_qpos(self) -> np.ndarray: + body_qpos = np.zeros(self.num_body_dof) + if self.unitree_bridge is not None and self.unitree_bridge.low_cmd: + for i in range(self.unitree_bridge.num_body_motor): + body_qpos[i] = self.unitree_bridge.low_cmd.motor_cmd[i].q + return body_qpos + + def compute_hand_qpos(self) -> np.ndarray: + hand_qpos = np.zeros(self.num_hand_dof * 2) + if self.unitree_bridge is not None and self.unitree_bridge.low_cmd: + for i in range(self.unitree_bridge.num_hand_motor): + hand_qpos[i] = self.unitree_bridge.left_hand_cmd.motor_cmd[i].q + hand_qpos[i + self.num_hand_dof] = self.unitree_bridge.right_hand_cmd.motor_cmd[i].q + return hand_qpos + + def prepare_obs(self) -> Dict[str, any]: + obs = {} + if self.use_floating_root_link: + obs["floating_base_pose"] = self.mj_data.qpos[:7] + obs["floating_base_vel"] = self.mj_data.qvel[:6] + obs["floating_base_acc"] = self.mj_data.qacc[:6] + else: + obs["floating_base_pose"] = np.zeros(7) + obs["floating_base_vel"] = np.zeros(6) + obs["floating_base_acc"] = np.zeros(6) + + obs["secondary_imu_quat"] = self.mj_data.xquat[self.torso_index] + + pose = np.zeros(13) + torso_link = self.mj_model.body("torso_link").id + # mj_objectVelocity returns [ang_vel, lin_vel]; swap to [lin_vel, ang_vel] + mujoco.mj_objectVelocity( + self.mj_model, self.mj_data, mujoco.mjtObj.mjOBJ_BODY, torso_link, pose[7:13], 1 + ) + pose[7:10], pose[10:13] = ( + pose[10:13], + pose[7:10].copy(), + ) + obs["secondary_imu_vel"] = pose[7:13] + + obs["body_q"] = self.mj_data.qpos[self.body_joint_index + 7 - 1] + obs["body_dq"] = self.mj_data.qvel[self.body_joint_index + 6 - 1] + obs["body_ddq"] = self.mj_data.qacc[self.body_joint_index + 6 - 1] + obs["body_tau_est"] = self.mj_data.actuator_force[self.body_joint_index - 1] + if self.num_hand_dof > 0: + obs["left_hand_q"] = self.mj_data.qpos[self.left_hand_index + self.qpos_offset - 1] + obs["left_hand_dq"] = self.mj_data.qvel[self.left_hand_index + self.qvel_offset - 1] + obs["left_hand_ddq"] = self.mj_data.qacc[self.left_hand_index + self.qvel_offset - 1] + obs["left_hand_tau_est"] = self.mj_data.actuator_force[self.left_hand_index - 1] + obs["right_hand_q"] = self.mj_data.qpos[self.right_hand_index + self.qpos_offset - 1] + obs["right_hand_dq"] = self.mj_data.qvel[self.right_hand_index + self.qvel_offset - 1] + obs["right_hand_ddq"] = self.mj_data.qacc[self.right_hand_index + self.qvel_offset - 1] + obs["right_hand_tau_est"] = self.mj_data.actuator_force[self.right_hand_index - 1] + obs["time"] = self.mj_data.time + return obs + + def sim_step(self): + self.obs = self.prepare_obs() + self.unitree_bridge.PublishLowState(self.obs) + if self.unitree_bridge.joystick: + self.unitree_bridge.PublishWirelessController() + if self.elastic_band: + if self.elastic_band.enable and self.use_floating_root_link: + pose = np.concatenate( + [ + self.mj_data.xpos[self.band_attached_link], + self.mj_data.xquat[self.band_attached_link], + np.zeros(6), + ] + ) + mujoco.mj_objectVelocity( + self.mj_model, + self.mj_data, + mujoco.mjtObj.mjOBJ_BODY, + self.band_attached_link, + pose[7:13], + 0, + ) + pose[7:10], pose[10:13] = pose[10:13], pose[7:10].copy() + self.mj_data.xfrc_applied[self.band_attached_link] = self.elastic_band.Advance(pose) + else: + self.mj_data.xfrc_applied[self.band_attached_link] = np.zeros(6) + body_torques = self.compute_body_torques() + hand_torques = self.compute_hand_torques() + # -1: actuator array is 0-based while joint indices from the model are 1-based + self.torques[self.body_joint_index - 1] = body_torques + if self.num_hand_dof > 0: + self.torques[self.left_hand_index - 1] = hand_torques[: self.num_hand_dof] + self.torques[self.right_hand_index - 1] = hand_torques[self.num_hand_dof :] + + self.torques = np.clip(self.torques, -self.torque_limit, self.torque_limit) + + if self.config["FREE_BASE"]: + # Prepend 6 zeros for the floating-base root DOF actuators + self.mj_data.ctrl = np.concatenate((np.zeros(6), self.torques)) + else: + self.mj_data.ctrl = self.torques + mujoco.mj_step(self.mj_model, self.mj_data) + + self.check_fall() + + def apply_perturbation(self, key): + perturbation_x_body = 0.0 + perturbation_y_body = 0.0 + if key == "up": + perturbation_x_body = 1.0 + elif key == "down": + perturbation_x_body = -1.0 + elif key == "left": + perturbation_y_body = 1.0 + elif key == "right": + perturbation_y_body = -1.0 + + vel_body = np.array([perturbation_x_body, perturbation_y_body, 0.0]) + vel_world = np.zeros(3) + base_quat = self.mj_data.qpos[3:7] + mujoco.mju_rotVecQuat(vel_world, vel_body, base_quat) + + self.mj_data.qvel[0] += vel_world[0] + self.mj_data.qvel[1] += vel_world[1] + mujoco.mj_forward(self.mj_model, self.mj_data) + + def update_viewer(self): + if self.viewer is not None: + self.viewer.sync() + + def update_viewer_camera(self): + if self.viewer is not None: + if self.viewer.cam.type == mujoco.mjtCamera.mjCAMERA_TRACKING: + self.viewer.cam.type = mujoco.mjtCamera.mjCAMERA_FREE + else: + self.viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING + + def update_reward(self): + with self.reward_lock: + self.last_reward = 0 + + def get_reward(self): + with self.reward_lock: + return self.last_reward + + def set_unitree_bridge(self, unitree_bridge): + self.unitree_bridge = unitree_bridge + + def get_privileged_obs(self): + return {} + + def update_render_caches(self): + render_caches = {} + for camera_name, camera_config in self.camera_configs.items(): + renderer = self.renderers[camera_name] + if "params" in camera_config: + renderer.update_scene(self.mj_data, camera=camera_config["params"]) + elif "mjcf_name" in camera_config: + renderer.update_scene(self.mj_data, camera=camera_config["mjcf_name"]) + else: + renderer.update_scene(self.mj_data, camera=camera_name) + render_caches[camera_name + "_image"] = renderer.render() + + if self.image_publish_process is not None: + self.image_publish_process.update_shared_memory(render_caches) + + return render_caches + + def handle_keyboard_button(self, key): + if self.elastic_band: + self.elastic_band.handle_keyboard_button(key) + + if key == "backspace": + self.reset() + if key == "v": + self.update_viewer_camera() + if key in ["up", "down", "left", "right"]: + self.apply_perturbation(key) + + def check_fall(self): + self.fall = False + if self.mj_data.qpos[2] < 0.2: + self.fall = True + print(f"Warning: Robot has fallen, height: {self.mj_data.qpos[2]:.3f} m") + + if self.fall: + self.reset() + + def check_self_collision(self): + robot_bodies = get_subtree_body_names(self.mj_model, self.mj_model.body(self.root_body).id) + self_collision, contact_bodies = check_contact( + self.mj_model, self.mj_data, robot_bodies, robot_bodies, return_all_contact_bodies=True + ) + if self_collision: + print(f"Warning: Self-collision detected: {contact_bodies}") + return self_collision + + def reset(self): + mujoco.mj_resetData(self.mj_model, self.mj_data) + + +class BaseSimulator: + """Base simulator class that handles initialization and running of simulations""" + + def __init__( + self, config: Dict[str, any], env_name: str = "default", redis_client=None, **kwargs + ): + self.config = config + self.env_name = env_name + self.redis_client = redis_client + if self.redis_client is not None: + self.redis_client.set("push_left_hand", "false") + self.redis_client.set("push_right_hand", "false") + self.redis_client.set("push_torso", "false") + + # Create rate objects + self.sim_dt = self.config["SIMULATE_DT"] + self.reward_dt = self.config.get("REWARD_DT", 0.02) + self.image_dt = self.config.get("IMAGE_DT", 0.033333) + self.viewer_dt = self.config.get("VIEWER_DT", 0.02) + self._running = True + + self.robot = Robot(self.config) + + # Create the environment + if env_name == "default": + self.sim_env = DefaultEnv(config, env_name, **kwargs) + else: + raise ValueError( + f"Invalid environment name: {env_name}. " + f"Only 'default' is supported in this minimal build." + ) + + try: + if self.config.get("INTERFACE", None): + ChannelFactoryInitialize(self.config["DOMAIN_ID"], self.config["INTERFACE"]) + else: + ChannelFactoryInitialize(self.config["DOMAIN_ID"]) + except Exception as e: + print(f"Note: Channel factory initialization attempt: {e}") + + self.init_unitree_bridge() + self.sim_env.set_unitree_bridge(self.unitree_bridge) + + self.init_subscriber() + self.init_publisher() + + self.sim_thread = None + + def start_as_thread(self): + self.sim_thread = Thread(target=self.start) + self.sim_thread.start() + + def start_image_publish_subprocess(self, start_method: str = "spawn", camera_port: int = 5555): + self.sim_env.start_image_publish_subprocess(start_method, camera_port) + + def init_subscriber(self): + pass + + def init_publisher(self): + pass + + def init_unitree_bridge(self): + self.unitree_bridge = UnitreeSdk2Bridge(self.config) + if self.config["USE_JOYSTICK"]: + self.unitree_bridge.SetupJoystick( + device_id=self.config["JOYSTICK_DEVICE"], js_type=self.config["JOYSTICK_TYPE"] + ) + + def start(self): + """Main simulation loop""" + sim_cnt = 0 + ts = time.time() + + try: + while self._running and ( + (self.sim_env.viewer and self.sim_env.viewer.is_running()) + or (self.sim_env.viewer is None) + ): + step_start = time.monotonic() + + self.sim_env.sim_step() + now = time.time() + if now - ts > 1 / 10.0 and self.redis_client is not None: + head_pose = self.sim_env.get_head_pose() + self.redis_client.set("head_pos", pickle.dumps(head_pose[:3])) + self.redis_client.set("head_quat", pickle.dumps(head_pose[3:])) + ts = now + + if sim_cnt % int(self.viewer_dt / self.sim_dt) == 0: + self.sim_env.update_viewer() + + if sim_cnt % int(self.reward_dt / self.sim_dt) == 0: + self.sim_env.update_reward() + + if sim_cnt % int(self.image_dt / self.sim_dt) == 0: + self.sim_env.update_render_caches() + + # Simple rate limiter (replaces ROS rate) + elapsed = time.monotonic() - step_start + sleep_time = self.sim_dt - elapsed + if sleep_time > 0: + time.sleep(sleep_time) + + sim_cnt += 1 + except KeyboardInterrupt: + print("Simulator interrupted by user.") + finally: + self.close() + + def __del__(self): + self.close() + + def reset(self): + self.sim_env.reset() + + def close(self): + self._running = False + try: + if self.sim_env.image_publish_process is not None: + self.sim_env.image_publish_process.stop() + if self.sim_env.viewer is not None: + self.sim_env.viewer.close() + except Exception as e: + print(f"Warning during close: {e}") + + def get_privileged_obs(self): + return self.sim_env.get_privileged_obs() + + def handle_keyboard_button(self, key): + self.sim_env.handle_keyboard_button(key) diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/configs.py b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/configs.py new file mode 100644 index 0000000000000000000000000000000000000000..5ec03facf8a5853ad6da1e2fa37bc999d51483b1 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/configs.py @@ -0,0 +1,346 @@ +"""Dataclass-based configuration for MuJoCo simulation, WBC deployment, and teleop. + +BaseConfig holds all tuneable knobs (interface, frequency, safety limits, etc.) +and can load/override values from a WBC YAML file. SimLoopConfig extends it +with multiprocessing and image-publish settings for the sim loop entry point. +""" + +from dataclasses import asdict, dataclass +import os +from pathlib import Path +import subprocess +from typing import Any, Literal, Optional + +import yaml + +from gear_sonic.utils.network.network_utils import resolve_interface + +WBC_VERSIONS = ["sonic_model12"] + +@dataclass +class ArgsConfigTemplate: + """Args Config for running the data collection loop.""" + + def update( + self, + config_dict: dict, + strict: bool = False, + skip_keys: list[str] = [], + allowed_keys: list[str] | None = None, + ): + for k, v in config_dict.items(): + if k in skip_keys: + continue + if allowed_keys is not None and k not in allowed_keys: + continue + if strict and not hasattr(self, k): + raise ValueError(f"Config {k} not found in {self.__class__.__name__}") + if not strict and not hasattr(self, k): + continue + setattr(self, k, v) + + @classmethod + def from_dict( + cls, + config_dict: dict, + strict: bool = False, + skip_keys: list[str] = [], + allowed_keys: list[str] | None = None, + ): + instance = cls() + instance.update( + config_dict=config_dict, strict=strict, skip_keys=skip_keys, allowed_keys=allowed_keys + ) + return instance + + def to_dict(self): + return asdict(self) + + def get(self, key: str, default: Any = None): + return getattr(self, key) if hasattr(self, key) else default + + +def override_wbc_config( + wbc_config: dict, config: "BaseConfig", missed_keys_only: bool = False +) -> dict: + """Override WBC YAML values with dataclass values.""" + key_to_value = { + "INTERFACE": config.interface, + "ENV_TYPE": config.env_type, + "VERSION": config.wbc_version, + "SIMULATOR": config.simulator, + "SIMULATE_DT": 1 / float(config.sim_frequency), + "ENABLE_OFFSCREEN": config.enable_offscreen, + "ENABLE_ONSCREEN": config.enable_onscreen, + "model_path": config.wbc_model_path, + "enable_waist": config.enable_waist, + "with_hands": config.with_hands, + "verbose": config.verbose, + "verbose_timing": config.verbose_timing, + "upper_body_max_joint_speed": config.upper_body_joint_speed, + "keyboard_dispatcher_type": config.keyboard_dispatcher_type, + "enable_gravity_compensation": config.enable_gravity_compensation, + "gravity_compensation_joints": config.gravity_compensation_joints, + "high_elbow_pose": config.high_elbow_pose, + "joint_safety_mode": config.joint_safety_mode, + "arm_velocity_limit": config.arm_velocity_limit, + "hand_velocity_limit": config.hand_velocity_limit, + "lower_body_velocity_limit": config.lower_body_velocity_limit, + "waist_pitch_limit": config.waist_pitch_limit, + "hand_torque_limit": config.hand_torque_limit, + "enable_natural_walk": config.enable_natural_walk, + } + + if missed_keys_only: + for key in key_to_value: + if key not in wbc_config: + wbc_config[key] = key_to_value[key] + else: + for key in key_to_value: + wbc_config[key] = key_to_value[key] + + # Sim-to-real KD gap: waist pitch (index 14) is over-damped in sim; + # reduce KD by 10 on the real robot to avoid sluggish response + if config.env_type == "real": + wbc_config["MOTOR_KD"][14] = wbc_config["MOTOR_KD"][14] - 10 + + return wbc_config + + +@dataclass +class BaseConfig(ArgsConfigTemplate): + """Base config inherited by all G1 control loops""" + + dataset_version: str = "sonic_model12" + + # WBC Configuration + wbc_version: Literal[tuple(WBC_VERSIONS)] = "sonic_model12" + """Version of the whole body controller.""" + + wbc_model_path: str = "policy/stand.onnx,policy/walk.onnx" + """Path to WBC model file (relative to GearCheckpoints/wbc)""" + + wbc_policy_class: str = "G1DecoupledWholeBodyPolicy" + """Whole body policy class.""" + + # System Configuration + interface: str = "sim" + """Interface to use for the control loop. [sim, real, lo, enxe8ea6a9c4e09]""" + + simulator: str = "mujoco" + """Simulator to use.""" + + sim_sync_mode: bool = False + """Whether to run the control loop in sync mode.""" + + control_frequency: int = 50 + """Frequency of the control loop.""" + + sim_frequency: int = 200 + """Frequency of the simulation loop.""" + + # Robot Configuration + enable_waist: bool = True + """Whether to include waist joints in IK.""" + + with_hands: bool = True + """Enable hand functionality.""" + + high_elbow_pose: bool = False + """Enable high elbow pose configuration.""" + + verbose: bool = True + """Whether to print verbose output.""" + + enable_offscreen: bool = False + """Whether to enable offscreen rendering.""" + + enable_onscreen: bool = True + """Whether to enable onscreen rendering.""" + + enable_teleop_evaluator: bool = False + """Whether to enable teleop evaluator.""" + + upper_body_joint_speed: float = 1000 + """Upper body joint speed.""" + + env_name: str = "default" + """Environment name.""" + + ik_indicator: bool = False + """Whether to draw IK indicators.""" + + verbose_timing: bool = False + """Enable verbose timing output every iteration.""" + + keyboard_dispatcher_type: str = "raw" + """Keyboard dispatcher to use. [raw, ros]""" + + # Gravity Compensation Configuration + enable_gravity_compensation: bool = False + """Enable gravity compensation using pinocchio dynamics.""" + + gravity_compensation_joints: Optional[list[str]] = None + """Joint groups to apply gravity compensation to.""" + + # Joint Safety Configuration + joint_safety_mode: Literal["kill", "freeze"] = "kill" + """Joint safety violation mode.""" + + arm_velocity_limit: float = 25.0 + """Arm joint velocity limit in rad/s.""" + + hand_velocity_limit: float = 1000.0 + """Hand/finger joint velocity limit in rad/s.""" + + lower_body_velocity_limit: float = 20.0 + """Lower body joint velocity limit in rad/s.""" + + waist_pitch_limit: float = 15.0 + """Waist pitch position limit in degrees.""" + + hand_torque_limit: float = 0.1 + """Hand torque limit in radians.""" + + enable_natural_walk: bool = False + """Enable natural walk mode.""" + + # Teleop/Device Configuration + body_control_device: str = "dummy" + """Device to use for body control.""" + + hand_control_device: Optional[str] = "dummy" + """Device to use for hand control.""" + + body_streamer_ip: str = "10.112.210.229" + """IP address for body streamer (vive only).""" + + body_streamer_keyword: str = "knee" + """Body streamer keyword (vive only).""" + + enable_visualization: bool = False + """Whether to enable visualization.""" + + enable_real_device: bool = True + """Whether to enable real device.""" + + teleop_frequency: int = 20 + """Teleoperation frequency (Hz).""" + + teleop_replay_path: Optional[str] = None + """Path to teleop replay data.""" + + # Deployment/Camera Configuration + robot_ip: str = "192.168.123.164" + """Robot IP address""" + + data_collection: bool = True + """Enable data collection""" + + data_collection_frequency: int = 20 + """Data collection frequency (Hz)""" + + root_output_dir: str = "outputs" + """Root output directory""" + + offline_dc: bool = False + """Offline data collection.""" + + enable_upper_body_operation: bool = True + """Enable upper body operation""" + + upper_body_operation_mode: Literal["teleop", "inference"] = "teleop" + """Upper body operation mode""" + + inference_host: str = "localhost" + """Inference server host""" + + inference_port: int = 5558 + """Inference server port (default 5558 to avoid conflict with camera_port 5555)""" + + inference_remote: bool = False + """Whether to run inference on a remote server.""" + + inference_prompt: str = "Pick up apple from table to plate" + """Inference prompt""" + + inference_action_horizon: int = 16 + """Inference action horizon""" + + inference_control_freq: int = 20 + """Inference control frequency (Hz)""" + + inference_rate: float = 2.5 + """Inference rate (Hz)""" + + inference_plot_rerun: bool = False + """Enable inference plot rerun""" + + inference_push_evals: bool = True + """Whether to push evals.""" + + inference_publish_single_action: bool = False + """Whether to publish only a single action.""" + + commit_id: str = "" + """Commit ID for the current codebase""" + + enable_mode_switch: bool = False + """Enable operation mode switching via ROS topic.""" + + initial_mode: str = "idle" + """Initial operation mode.""" + + def __post_init__(self): + # Resolve interface + self.interface, self.env_type = resolve_interface(self.interface) + + # Set default gravity compensation joints if not specified + if self.gravity_compensation_joints is None: + self.gravity_compensation_joints = ["arms"] + + try: + self.commit_id = ( + subprocess.check_output(["git", "rev-parse", "HEAD"]).decode("utf-8").strip() + ) + except Exception: + self.commit_id = "" + + def load_wbc_yaml(self) -> dict: + """Load and merge wbc yaml with dataclass overrides""" + import gear_sonic + + gear_sonic_path = Path(os.path.dirname(gear_sonic.__file__)) + configs_dir = gear_sonic_path / "utils" / "mujoco_sim" / "wbc_configs" + + if self.wbc_version == "sonic_model12": + config_path = str(configs_dir / "g1_29dof_sonic_model12.yaml") + else: + raise ValueError( + f"Invalid wbc_version: {self.wbc_version}, please use one of: " + f"sonic_model12" + ) + + with open(config_path) as file: + wbc_config = yaml.load(file, Loader=yaml.FullLoader) + + wbc_config = override_wbc_config(wbc_config, self) + + return wbc_config + + +@dataclass +class SimLoopConfig(BaseConfig): + """Config for running the simulation loop.""" + + mp_start_method: str = "spawn" + """Multiprocessing start method""" + + enable_image_publish: bool = False + """Enable image publishing in simulation""" + + camera_port: int = 5555 + """Camera port for image publishing""" + + verbose: bool = False + """Verbose output, override the base config verbose""" diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/image_publish_utils.py b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/image_publish_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2f559775111925fccbe87e0743785d93e0c9badf --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/image_publish_utils.py @@ -0,0 +1,201 @@ +"""Offloads rendered MuJoCo camera images to a subprocess via shared memory + ZMQ.""" + +import multiprocessing as mp +from multiprocessing import shared_memory +import time +from typing import Any, Dict + +import numpy as np + +from gear_sonic.utils.mujoco_sim.sensor_server import ImageMessageSchema, SensorServer + + +def get_multiprocessing_info(verbose: bool = True): + """Get information about multiprocessing start methods""" + + if verbose: + print(f"Available start methods: {mp.get_all_start_methods()}") + return mp.get_start_method() + + +class ImagePublishProcess: + """Subprocess for publishing images using shared memory and ZMQ""" + + def __init__( + self, + camera_configs: Dict[str, Any], + image_dt: float, + zmq_port: int = 5555, + start_method: str = "spawn", + verbose: bool = False, + ): + self.camera_configs = camera_configs + self.image_dt = image_dt + self.zmq_port = zmq_port + self.verbose = verbose + self.shared_memory_blocks = {} + self.shared_memory_info = {} + self.process = None + + self.mp_context = mp.get_context(start_method) + if self.verbose: + print(f"Using multiprocessing context: {start_method}") + + self.stop_event = self.mp_context.Event() + self.data_ready_event = self.mp_context.Event() + + self.stop_event.clear() + self.data_ready_event.clear() + + for camera_name, camera_config in camera_configs.items(): + height = camera_config["height"] + width = camera_config["width"] + size = height * width * 3 + + shm = shared_memory.SharedMemory(create=True, size=size) + self.shared_memory_blocks[camera_name] = shm + self.shared_memory_info[camera_name] = { + "name": shm.name, + "size": size, + "shape": (height, width, 3), + "dtype": np.uint8, + } + + def start_process(self): + """Start the image publishing subprocess""" + self.process = self.mp_context.Process( + target=self._image_publish_worker, + args=( + self.shared_memory_info, + self.image_dt, + self.zmq_port, + self.stop_event, + self.data_ready_event, + self.verbose, + ), + ) + self.process.start() + + def update_shared_memory(self, render_caches: Dict[str, np.ndarray]): + """Update shared memory with new rendered images""" + images_updated = 0 + for camera_name in self.camera_configs.keys(): + image_key = f"{camera_name}_image" + if image_key in render_caches: + image = render_caches[image_key] + + if image.dtype != np.uint8: + image = (image * 255).astype(np.uint8) + + shm = self.shared_memory_blocks[camera_name] + shared_array = np.ndarray( + self.shared_memory_info[camera_name]["shape"], + dtype=self.shared_memory_info[camera_name]["dtype"], + buffer=shm.buf, + ) + + np.copyto(shared_array, image) + images_updated += 1 + + if images_updated > 0: + self.data_ready_event.set() + + def stop(self): + """Stop the image publishing subprocess""" + self.stop_event.set() + + if self.process and self.process.is_alive(): + self.process.join(timeout=5) + if self.process.is_alive(): + self.process.terminate() + self.process.join(timeout=2) + if self.process.is_alive(): + self.process.kill() + self.process.join() + + for camera_name, shm in self.shared_memory_blocks.items(): + try: + shm.close() + shm.unlink() + except Exception as e: + print(f"Warning: Failed to cleanup shared memory for {camera_name}: {e}") + + self.shared_memory_blocks.clear() + + @staticmethod + def _image_publish_worker( + shared_memory_info, image_dt, zmq_port, stop_event, data_ready_event, verbose + ): + """Worker function that runs in the subprocess""" + try: + sensor_server = SensorServer() + sensor_server.start_server(port=zmq_port) + + shared_arrays = {} + shm_blocks = {} + for camera_name, info in shared_memory_info.items(): + shm = shared_memory.SharedMemory(name=info["name"]) + shm_blocks[camera_name] = shm + shared_arrays[camera_name] = np.ndarray( + info["shape"], dtype=info["dtype"], buffer=shm.buf + ) + + print( + f"Image publishing subprocess started with {len(shared_arrays)} cameras " + f"on ZMQ port {zmq_port}" + ) + + loop_count = 0 + last_data_time = time.time() + + while not stop_event.is_set(): + loop_count += 1 + + timeout = min(image_dt, 0.1) + data_available = data_ready_event.wait(timeout=timeout) + + current_time = time.time() + + if data_available: + data_ready_event.clear() + if loop_count % 50 == 0: + print("Image publish frequency: ", 1 / (current_time - last_data_time)) + last_data_time = current_time + + try: + from gear_sonic.utils.mujoco_sim.sensor_server import ImageUtils + + image_copies = {name: arr.copy() for name, arr in shared_arrays.items()} + + message_dict = { + "images": image_copies, + "timestamps": {name: current_time for name in image_copies.keys()}, + } + + image_msg = ImageMessageSchema( + timestamps=message_dict.get("timestamps"), + images=message_dict.get("images", None), + ) + + serialized_data = image_msg.serialize() + + for camera_name, image_copy in image_copies.items(): + serialized_data[f"{camera_name}"] = ImageUtils.encode_image(image_copy) + + sensor_server.send_message(serialized_data) + + except Exception as e: + print(f"Error publishing images: {e}") + + if not data_available: + time.sleep(0.001) + + except KeyboardInterrupt: + print("Image publisher interrupted by user") + finally: + try: + for shm in shm_blocks.values(): + shm.close() + sensor_server.stop_server() + except Exception as e: + print(f"Error during subprocess cleanup: {e}") diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/metric_utils.py b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/metric_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..4f688ac352cf9cebdf58b0e41011476dbd780a24 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/metric_utils.py @@ -0,0 +1,73 @@ +"""Contact detection and geom height checks for MuJoCo simulations.""" + +from typing import List, Tuple + +import mujoco + +from gear_sonic.utils.mujoco_sim.sim_utils import get_body_geom_ids + + +def check_contact( + mj_model: mujoco.MjModel, + mj_data: mujoco.MjData, + bodies_1: List[str] | str, + bodies_2: List[str] | str, + return_all_contact_bodies: bool = False, +) -> Tuple[bool, List[Tuple[str, str]]] | bool: + """ + Finds contact between two body groups. Any geom in the body is considered to be in contact. + Args: + mj_model (MujocoModel): Current simulation object + mj_data (MjData): Current simulation data + bodies_1 (str or list of int): an individual body name or list of body names. + bodies_2 (str or list of int): another individual body name or list of body names. + Returns: + bool: True if any body in @bodies_1 is in contact with any body in @bodies_2. + """ + if isinstance(bodies_1, str): + bodies_1 = [bodies_1] + if isinstance(bodies_2, str): + bodies_2 = [bodies_2] + + geoms_1 = [get_body_geom_ids(mj_model, mj_model.body(g).id) for g in bodies_1] + geoms_1 = [g for geom_list in geoms_1 for g in geom_list] + geoms_2 = [get_body_geom_ids(mj_model, mj_model.body(g).id) for g in bodies_2] + geoms_2 = [g for geom_list in geoms_2 for g in geom_list] + contact_bodies = [] + for i in range(mj_data.ncon): + contact = mj_data.contact[i] + # check contact geom in geoms + c1_in_g1 = contact.geom1 in geoms_1 + c2_in_g2 = contact.geom2 in geoms_2 if geoms_2 is not None else True + # check contact geom in geoms (flipped) + c2_in_g1 = contact.geom2 in geoms_1 + c1_in_g2 = contact.geom1 in geoms_2 if geoms_2 is not None else True + if (c1_in_g1 and c2_in_g2) or (c1_in_g2 and c2_in_g1): + contact_bodies.append( + ( + mj_model.body(mj_model.geom(contact.geom1).bodyid).name, + mj_model.body(mj_model.geom(contact.geom2).bodyid).name, + ) + ) + if not return_all_contact_bodies: + break + if return_all_contact_bodies: + return len(contact_bodies) > 0, set(contact_bodies) + else: + return len(contact_bodies) > 0 + + +def check_height( + mj_model: mujoco.MjModel, + mj_data: mujoco.MjData, + geom_name: str, + lower_bound: float = -float("inf"), + upper_bound: float = float("inf"), +): + """ + Checks if the height of a geom is greater than a given height. + """ + geom_id = mj_model.geom(geom_name).id + return ( + mj_data.geom_xpos[geom_id][2] < upper_bound and mj_data.geom_xpos[geom_id][2] > lower_bound + ) diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/robot.py b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/robot.py new file mode 100644 index 0000000000000000000000000000000000000000..1e93258a98621afe93dedcf5c68be722d4ae4dc0 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/robot.py @@ -0,0 +1,26 @@ +"""Legacy robot config container (deprecated -- prefer robot_model package).""" + + +# TODO: This class and associated configs should be deleted and instead use the robot_model package +class Robot: + def __init__(self, config): + self.ROBOT_TYPE = config["ROBOT_TYPE"] + self.MOTOR2JOINT = config["MOTOR2JOINT"] + self.JOINT2MOTOR = config["JOINT2MOTOR"] + self.UNITREE_LEGGED_CONST = config["UNITREE_LEGGED_CONST"] + self.MOTOR_KP = config["MOTOR_KP"] + self.MOTOR_KD = config["MOTOR_KD"] + if "HAND_MOTOR_KP" in config: + self.HAND_MOTOR_KP = config["HAND_MOTOR_KP"] + if "HAND_MOTOR_KD" in config: + self.HAND_MOTOR_KD = config["HAND_MOTOR_KD"] + + self.WeakMotorJointIndex = config["WeakMotorJointIndex"] + self.NUM_MOTORS = config["NUM_MOTORS"] + self.NUM_JOINTS = config["NUM_JOINTS"] + self.NUM_HAND_MOTORS = config.get("NUM_HAND_MOTORS", 0) # only 43dof has hand + self.NUM_HAND_JOINTS = config.get("NUM_HAND_JOINTS", 0) + self.DEFAULT_DOF_ANGLES = config["DEFAULT_DOF_ANGLES"] + self.DEFAULT_MOTOR_ANGLES = config["DEFAULT_MOTOR_ANGLES"] + self.USE_SENSOR = config["USE_SENSOR"] + self.MOTOR_EFFORT_LIMIT_LIST = config["motor_effort_limit_list"] diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/sensor_server.py b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/sensor_server.py new file mode 100644 index 0000000000000000000000000000000000000000..21c310980582a78d3cf1803598aae6620c50eaba --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/sensor_server.py @@ -0,0 +1,83 @@ +"""ZMQ PUB server for streaming JPEG-encoded camera images as msgpack payloads.""" + +import base64 +from dataclasses import dataclass, field +from typing import Any, Dict + +import cv2 +import msgpack +import msgpack_numpy as m +import numpy as np +import zmq + + +@dataclass +class ImageMessageSchema: + """ + Standardized message schema for image data. + """ + + timestamps: Dict[str, float] + images: Dict[str, np.ndarray] + + def serialize(self) -> Dict[str, Any]: + serialized_msg = {"timestamps": self.timestamps, "images": {}} + for key, image in self.images.items(): + serialized_msg["images"][key] = ImageUtils.encode_image(image) + return serialized_msg + + @staticmethod + def deserialize(data: Dict[str, Any]) -> "ImageMessageSchema": + timestamps = data.get("timestamps", {}) + images = {} + for key, value in data.get("images", {}).items(): + if isinstance(value, str): + images[key] = ImageUtils.decode_image(value) + else: + images[key] = value + return ImageMessageSchema(timestamps=timestamps, images=images) + + +class SensorServer: + def start_server(self, port: int): + self.context = zmq.Context() + self.socket = self.context.socket(zmq.PUB) + self.socket.setsockopt(zmq.SNDHWM, 20) + self.socket.setsockopt(zmq.LINGER, 0) + self.socket.bind(f"tcp://*:{port}") + print(f"Sensor server running at tcp://*:{port}") + + self.message_sent = 0 + self.message_dropped = 0 + + def stop_server(self): + self.socket.close() + self.context.term() + + def send_message(self, data: Dict[str, Any]): + try: + packed = msgpack.packb(data, use_bin_type=True) + self.socket.send(packed, flags=zmq.NOBLOCK) + except zmq.Again: + self.message_dropped += 1 + print(f"[Warning] message dropped: {self.message_dropped}") + self.message_sent += 1 + + if self.message_sent % 100 == 0: + print( + f"[Sensor server] Message sent: {self.message_sent}, " + f"message dropped: {self.message_dropped}" + ) + + +class ImageUtils: + @staticmethod + def encode_image(image: np.ndarray) -> str: + _, color_buffer = cv2.imencode(".jpg", image, [int(cv2.IMWRITE_JPEG_QUALITY), 80]) + return base64.b64encode(color_buffer).decode("utf-8") + + @staticmethod + def decode_image(image: str) -> np.ndarray: + color_data = base64.b64decode(image) + color_array = np.frombuffer(color_data, dtype=np.uint8) + return cv2.imdecode(color_array, cv2.IMREAD_COLOR) diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/sim_utils.py b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/sim_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e032eb6d801469cb543dd10fa6cef11a90f643a8 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/sim_utils.py @@ -0,0 +1,96 @@ +""" +Utility functions for working with Mujoco models. +copied from https://github.com/kevinzakka/mink/blob/main/mink/utils.py +""" + +from typing import List + +import mujoco + + +def get_body_body_ids(model: mujoco.MjModel, body_id: int) -> List[int]: + """Get immediate children bodies belonging to a given body. + + Args: + model: Mujoco model. + body_id: ID of body. + + Returns: + A List containing all child body ids. + """ + return [ + i + for i in range(model.nbody) + if model.body_parentid[i] == body_id and body_id != i # Exclude the body itself. + ] + + +def get_subtree_body_ids(model: mujoco.MjModel, body_id: int) -> List[int]: + """Get all bodies belonging to subtree starting at a given body. + + Args: + model: Mujoco model. + body_id: ID of body where subtree starts. + + Returns: + A List containing all subtree body ids. + """ + body_ids: List[int] = [] + stack = [body_id] + while stack: + body_id = stack.pop() + body_ids.append(body_id) + stack += get_body_body_ids(model, body_id) + return body_ids + + +def get_subtree_body_names(model: mujoco.MjModel, body_id: int) -> List[str]: + """Get all bodies belonging to subtree starting at a given body. + Args: + model: Mujoco model. + body_id: ID of body where subtree starts. + + Returns: + A List containing all subtree body names. + """ + return [model.body(i).name for i in get_subtree_body_ids(model, body_id)] + + +def get_body_geom_ids(model: mujoco.MjModel, body_id: int) -> List[int]: + """Get immediate geoms belonging to a given body. + + Here, immediate geoms are those directly attached to the body and not its + descendants. + + Args: + model: Mujoco model. + body_id: ID of body. + + Returns: + A list containing all body geom ids. + """ + geom_start = model.body_geomadr[body_id] + geom_end = geom_start + model.body_geomnum[body_id] + return list(range(geom_start, geom_end)) + + +def get_subtree_geom_ids(model: mujoco.MjModel, body_id: int) -> List[int]: + """Get all geoms belonging to subtree starting at a given body. + + Here, a subtree is defined as the kinematic tree starting at the body and including + all its descendants. + + Args: + model: Mujoco model. + body_id: ID of body where subtree starts. + + Returns: + A list containing all subtree geom ids. + """ + geom_ids: List[int] = [] + stack = [body_id] + while stack: + body_id = stack.pop() + geom_ids.extend(get_body_geom_ids(model, body_id)) + stack += get_body_body_ids(model, body_id) + return geom_ids diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/simulator_factory.py b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/simulator_factory.py new file mode 100644 index 0000000000000000000000000000000000000000..008ae87eaa8b71a5e11c4a0a64bb410a51c3fee9 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/simulator_factory.py @@ -0,0 +1,98 @@ +"""Factory for creating and launching MuJoCo simulators with Unitree SDK channel setup.""" + +import time +from typing import Any, Dict + +from unitree_sdk2py.core.channel import ChannelFactoryInitialize + +from gear_sonic.utils.mujoco_sim.base_sim import BaseSimulator + + +def init_channel(config: Dict[str, Any]) -> None: + """ + Initialize the communication channel for simulator/robot communication. + + Args: + config: Configuration dictionary containing DOMAIN_ID and optionally INTERFACE + """ + if config.get("INTERFACE", None): + ChannelFactoryInitialize(config["DOMAIN_ID"], config["INTERFACE"]) + else: + ChannelFactoryInitialize(config["DOMAIN_ID"]) + + +class SimulatorFactory: + """Factory class for creating different types of simulators.""" + + @staticmethod + def create_simulator(config: Dict[str, Any], env_name: str = "default", **kwargs): + """ + Create a simulator based on the configuration. + + Args: + config: Configuration dictionary containing SIMULATOR type + env_name: Environment name + **kwargs: Additional keyword arguments for specific simulators + """ + simulator_type = config.get("SIMULATOR", "mujoco") + if simulator_type == "mujoco": + return SimulatorFactory._create_mujoco_simulator(config, env_name, **kwargs) + else: + print( + f"Warning: Invalid simulator type: {simulator_type}. " + "If you are using run_sim_loop, please ignore this warning." + ) + return None + + @staticmethod + def _create_mujoco_simulator(config: Dict[str, Any], env_name: str = "default", **kwargs): + """Create a MuJoCo simulator instance.""" + return BaseSimulator( + onscreen=kwargs.pop("onscreen", True), + offscreen=kwargs.pop("offscreen", False), + enable_image_publish=kwargs.get("enable_image_publish", False), + config=config, + env_name=env_name, + redis_client=kwargs.get("redis_client", None), + ) + + @staticmethod + def start_simulator( + simulator, + as_thread: bool = True, + enable_image_publish: bool = False, + mp_start_method: str = "spawn", + camera_port: int = 5555, + ): + """ + Start the simulator either as a thread or as a separate process. + + Args: + simulator: The simulator instance to start + as_thread: If True, start as thread; if False, start as subprocess + enable_image_publish: If True and not as_thread, start image publishing + mp_start_method: Multiprocessing start method + camera_port: Camera port for image publishing + """ + + if as_thread: + simulator.start_as_thread() + else: + try: + if enable_image_publish: + simulator.start_image_publish_subprocess( + start_method=mp_start_method, + camera_port=camera_port, + ) + time.sleep(1) + simulator.start() + except KeyboardInterrupt: + print("+++++Simulator interrupted by user.") + except Exception as e: + print(f"++++error in simulator: {e} ++++") + finally: + print("++++closing simulator ++++") + simulator.close() + + # Allow simulator to initialize + time.sleep(1) diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/unitree_sdk2py_bridge.py b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/unitree_sdk2py_bridge.py new file mode 100644 index 0000000000000000000000000000000000000000..fac16cad1f385329fcdaafb07875586d9f42115c --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/unitree_sdk2py_bridge.py @@ -0,0 +1,410 @@ +"""Bridge between Unitree SDK2 DDS topics and the MuJoCo simulation. + +Subscribes to low-level motor commands (body + hands) and publishes +simulated sensor state (joint pos/vel, IMU, odometry) back over DDS, +so the WBC policy sees the sim as a real robot. +""" + +import sys +import threading +from typing import Dict, Tuple + +import numpy as np +import scipy.spatial.transform +from unitree_sdk2py.core.channel import ChannelPublisher, ChannelSubscriber +from unitree_sdk2py.idl.default import ( + unitree_go_msg_dds__WirelessController_, + unitree_hg_msg_dds__HandCmd_ as HandCmd_default, + unitree_hg_msg_dds__HandState_ as HandState_default, +) +from unitree_sdk2py.idl.unitree_go.msg.dds_ import WirelessController_ +from unitree_sdk2py.idl.unitree_hg.msg.dds_ import HandCmd_, HandState_, OdoState_ + + +class UnitreeSdk2Bridge: + """ + This class is responsible for bridging the Unitree SDK2 with the Groot environment. + It is responsible for sending and receiving messages to and from the Unitree SDK2. + Both the body and hand are supported. + """ + + def __init__(self, config): + # Note that we do not give the mjdata and mjmodel to the UnitreeSdk2Bridge. + # It is unsafe and would be unflexible if we use a hand-plugged robot model + + robot_type = config["ROBOT_TYPE"] + if "g1" in robot_type or "h1-2" in robot_type: + from unitree_sdk2py.idl.default import ( + unitree_hg_msg_dds__IMUState_ as IMUState_default, + unitree_hg_msg_dds__LowCmd_, + unitree_hg_msg_dds__LowState_ as LowState_default, + unitree_hg_msg_dds__OdoState_ as OdoState_default, + ) + from unitree_sdk2py.idl.unitree_hg.msg.dds_ import IMUState_, LowCmd_, LowState_ + + self.low_cmd = unitree_hg_msg_dds__LowCmd_() + elif "h1" == robot_type or "go2" == robot_type: + from unitree_sdk2py.idl.default import ( + unitree_go_msg_dds__LowCmd_, + unitree_go_msg_dds__LowState_ as LowState_default, + unitree_hg_msg_dds__IMUState_ as IMUState_default, + ) + from unitree_sdk2py.idl.unitree_go.msg.dds_ import IMUState_, LowCmd_, LowState_ + + self.low_cmd = unitree_go_msg_dds__LowCmd_() + else: + raise ValueError(f"Invalid robot type '{robot_type}'. Expected 'g1', 'h1', or 'go2'.") + + self.num_body_motor = config["NUM_MOTORS"] + self.num_hand_motor = config.get("NUM_HAND_MOTORS", 0) + self.use_sensor = config["USE_SENSOR"] + + self.have_imu_ = False + self.have_frame_sensor_ = False + + # Unitree sdk2 message + self.low_state = LowState_default() + self.low_state_puber = ChannelPublisher("rt/lowstate", LowState_) + self.low_state_puber.Init() + + # Only create odo_state for supported robot types + if "g1" in robot_type or "h1-2" in robot_type: + self.odo_state = OdoState_default() + self.odo_state_puber = ChannelPublisher("rt/odostate", OdoState_) + self.odo_state_puber.Init() + else: + self.odo_state = None + self.odo_state_puber = None + self.torso_imu_state = IMUState_default() + self.torso_imu_puber = ChannelPublisher("rt/secondary_imu", IMUState_) + self.torso_imu_puber.Init() + + self.left_hand_state = HandState_default() + self.left_hand_state_puber = ChannelPublisher("rt/dex3/left/state", HandState_) + self.left_hand_state_puber.Init() + self.right_hand_state = HandState_default() + self.right_hand_state_puber = ChannelPublisher("rt/dex3/right/state", HandState_) + self.right_hand_state_puber.Init() + + self.low_cmd_suber = ChannelSubscriber("rt/lowcmd", LowCmd_) + self.low_cmd_suber.Init(self.LowCmdHandler, 1) + + self.left_hand_cmd = HandCmd_default() + self.left_hand_cmd_suber = ChannelSubscriber("rt/dex3/left/cmd", HandCmd_) + self.left_hand_cmd_suber.Init(self.LeftHandCmdHandler, 1) + self.right_hand_cmd = HandCmd_default() + self.right_hand_cmd_suber = ChannelSubscriber("rt/dex3/right/cmd", HandCmd_) + self.right_hand_cmd_suber.Init(self.RightHandCmdHandler, 1) + + self.low_cmd_lock = threading.Lock() + self.left_hand_cmd_lock = threading.Lock() + self.right_hand_cmd_lock = threading.Lock() + + self.wireless_controller = unitree_go_msg_dds__WirelessController_() + self.wireless_controller_puber = ChannelPublisher( + "rt/wirelesscontroller", WirelessController_ + ) + self.wireless_controller_puber.Init() + + # joystick + self.key_map = { + "R1": 0, + "L1": 1, + "start": 2, + "select": 3, + "R2": 4, + "L2": 5, + "F1": 6, + "F2": 7, + "A": 8, + "B": 9, + "X": 10, + "Y": 11, + "up": 12, + "right": 13, + "down": 14, + "left": 15, + } + self.joystick = None + + self.reset() + + def reset(self): + with self.low_cmd_lock: + self.low_cmd_received = False + self.new_low_cmd = False + with self.left_hand_cmd_lock: + self.left_hand_cmd_received = False + self.new_left_hand_cmd = False + with self.right_hand_cmd_lock: + self.right_hand_cmd_received = False + self.new_right_hand_cmd = False + + def LowCmdHandler(self, msg): + with self.low_cmd_lock: + self.low_cmd = msg + self.low_cmd_received = True + self.new_low_cmd = True + + def LeftHandCmdHandler(self, msg): + with self.left_hand_cmd_lock: + self.left_hand_cmd = msg + self.left_hand_cmd_received = True + self.new_left_hand_cmd = True + + def RightHandCmdHandler(self, msg): + with self.right_hand_cmd_lock: + self.right_hand_cmd = msg + self.right_hand_cmd_received = True + self.new_right_hand_cmd = True + + def cmd_received(self): + with self.low_cmd_lock: + low_cmd_received = self.low_cmd_received + with self.left_hand_cmd_lock: + left_hand_cmd_received = self.left_hand_cmd_received + with self.right_hand_cmd_lock: + right_hand_cmd_received = self.right_hand_cmd_received + return low_cmd_received or left_hand_cmd_received or right_hand_cmd_received + + def PublishLowState(self, obs: Dict[str, any]): + # publish body state + if self.use_sensor: + raise NotImplementedError("Sensor data is not implemented yet.") + else: + for i in range(self.num_body_motor): + self.low_state.motor_state[i].q = obs["body_q"][i] + self.low_state.motor_state[i].dq = obs["body_dq"][i] + self.low_state.motor_state[i].ddq = obs["body_ddq"][i] + self.low_state.motor_state[i].tau_est = obs["body_tau_est"][i] + + if self.use_sensor and self.have_frame_sensor_: + raise NotImplementedError("Frame sensor data is not implemented yet.") + else: + # Get data from ground truth + self.odo_state.position[:] = obs["floating_base_pose"][:3] + self.odo_state.linear_velocity[:] = obs["floating_base_vel"][:3] + self.odo_state.orientation[:] = obs["floating_base_pose"][3:7] + self.odo_state.angular_velocity[:] = obs["floating_base_vel"][3:6] + # quaternion: w, x, y, z + self.low_state.imu_state.quaternion[:] = obs["floating_base_pose"][3:7] + # angular velocity + self.low_state.imu_state.gyroscope[:] = obs["floating_base_vel"][3:6] + # linear acceleration + self.low_state.imu_state.accelerometer[:] = obs["floating_base_acc"][:3] + + self.torso_imu_state.quaternion[:] = obs["secondary_imu_quat"] + self.torso_imu_state.gyroscope[:] = obs["secondary_imu_vel"][3:6] + + # acceleration: x, y, z (only available when frame sensor is enabled) + if self.have_frame_sensor_: + raise NotImplementedError("Frame sensor data is not implemented yet.") + self.low_state.tick = int(obs["time"] * 1e3) + self.low_state_puber.Write(self.low_state) + + self.odo_state.tick = int(obs["time"] * 1e3) + self.odo_state_puber.Write(self.odo_state) + + self.torso_imu_puber.Write(self.torso_imu_state) + + # publish hand state + for i in range(self.num_hand_motor): + self.left_hand_state.motor_state[i].q = obs["left_hand_q"][i] + self.left_hand_state.motor_state[i].dq = obs["left_hand_dq"][i] + self.left_hand_state_puber.Write(self.left_hand_state) + + for i in range(self.num_hand_motor): + self.right_hand_state.motor_state[i].q = obs["right_hand_q"][i] + self.right_hand_state.motor_state[i].dq = obs["right_hand_dq"][i] + self.right_hand_state_puber.Write(self.right_hand_state) + + def GetAction(self) -> Tuple[np.ndarray, bool, bool]: + with self.low_cmd_lock: + body_q = [self.low_cmd.motor_cmd[i].q for i in range(self.num_body_motor)] + with self.left_hand_cmd_lock: + left_hand_q = [self.left_hand_cmd.motor_cmd[i].q for i in range(self.num_hand_motor)] + with self.right_hand_cmd_lock: + right_hand_q = [self.right_hand_cmd.motor_cmd[i].q for i in range(self.num_hand_motor)] + with self.low_cmd_lock and self.left_hand_cmd_lock and self.right_hand_cmd_lock: + is_new_action = self.new_low_cmd and self.new_left_hand_cmd and self.new_right_hand_cmd + if is_new_action: + self.new_low_cmd = False + self.new_left_hand_cmd = False + self.new_right_hand_cmd = False + + return ( + np.concatenate([body_q[:-7], left_hand_q, body_q[-7:], right_hand_q]), + self.cmd_received(), + is_new_action, + ) + + def PublishWirelessController(self): + import pygame + + if self.joystick is not None: + pygame.event.get() + key_state = [0] * 16 + key_state[self.key_map["R1"]] = self.joystick.get_button(self.button_id["RB"]) + key_state[self.key_map["L1"]] = self.joystick.get_button(self.button_id["LB"]) + key_state[self.key_map["start"]] = self.joystick.get_button(self.button_id["START"]) + key_state[self.key_map["select"]] = self.joystick.get_button(self.button_id["SELECT"]) + key_state[self.key_map["R2"]] = self.joystick.get_axis(self.axis_id["RT"]) > 0 + key_state[self.key_map["L2"]] = self.joystick.get_axis(self.axis_id["LT"]) > 0 + key_state[self.key_map["F1"]] = 0 + key_state[self.key_map["F2"]] = 0 + key_state[self.key_map["A"]] = self.joystick.get_button(self.button_id["A"]) + key_state[self.key_map["B"]] = self.joystick.get_button(self.button_id["B"]) + key_state[self.key_map["X"]] = self.joystick.get_button(self.button_id["X"]) + key_state[self.key_map["Y"]] = self.joystick.get_button(self.button_id["Y"]) + key_state[self.key_map["up"]] = self.joystick.get_hat(0)[1] > 0 + key_state[self.key_map["right"]] = self.joystick.get_hat(0)[0] > 0 + key_state[self.key_map["down"]] = self.joystick.get_hat(0)[1] < 0 + key_state[self.key_map["left"]] = self.joystick.get_hat(0)[0] < 0 + + # Pack 16 button states into a single integer via bit-shifting + key_value = 0 + for i in range(16): + key_value += key_state[i] << i + + self.wireless_controller.keys = key_value + self.wireless_controller.lx = self.joystick.get_axis(self.axis_id["LX"]) + self.wireless_controller.ly = -self.joystick.get_axis(self.axis_id["LY"]) + self.wireless_controller.rx = self.joystick.get_axis(self.axis_id["RX"]) + self.wireless_controller.ry = -self.joystick.get_axis(self.axis_id["RY"]) + + self.wireless_controller_puber.Write(self.wireless_controller) + + def SetupJoystick(self, device_id=0, js_type="xbox"): + import pygame + + pygame.init() + pygame.joystick.init() + joystick_count = pygame.joystick.get_count() + if joystick_count > 0: + self.joystick = pygame.joystick.Joystick(device_id) + self.joystick.init() + else: + print("No gamepad detected.") + sys.exit() + + if js_type == "xbox": + if sys.platform.startswith("linux"): + self.axis_id = { + "LX": 0, "LY": 1, "RX": 3, "RY": 4, + "LT": 2, "RT": 5, "DX": 6, "DY": 7, + } + self.button_id = { + "X": 2, "Y": 3, "B": 1, "A": 0, + "LB": 4, "RB": 5, "SELECT": 6, "START": 7, + "XBOX": 8, "LSB": 9, "RSB": 10, + } + elif sys.platform == "darwin": + self.axis_id = { + "LX": 0, "LY": 1, "RX": 2, "RY": 3, + "LT": 4, "RT": 5, + } + self.button_id = { + "X": 2, "Y": 3, "B": 1, "A": 0, + "LB": 9, "RB": 10, "SELECT": 4, "START": 6, + "XBOX": 5, "LSB": 7, "RSB": 8, + "DYU": 11, "DYD": 12, "DXL": 13, "DXR": 14, + } + else: + print("Unsupported OS. ") + + elif js_type == "switch": + self.axis_id = { + "LX": 0, "LY": 1, "RX": 2, "RY": 3, + "LT": 5, "RT": 4, "DX": 6, "DY": 7, + } + self.button_id = { + "X": 3, "Y": 4, "B": 1, "A": 0, + "LB": 6, "RB": 7, "SELECT": 10, "START": 11, + } + else: + print("Unsupported gamepad. ") + + def PrintSceneInformation(self): + import mujoco + from loguru import logger + from termcolor import colored + + print(" ") + logger.info(colored("<<------------- Link ------------->>", "green")) + for i in range(self.mj_model.nbody): + name = mujoco.mj_id2name(self.mj_model, mujoco._enums.mjtObj.mjOBJ_BODY, i) + if name: + logger.info(f"link_index: {i}, name: {name}") + print(" ") + + logger.info(colored("<<------------- Joint ------------->>", "green")) + for i in range(self.mj_model.njnt): + name = mujoco.mj_id2name(self.mj_model, mujoco._enums.mjtObj.mjOBJ_JOINT, i) + if name: + logger.info(f"joint_index: {i}, name: {name}") + print(" ") + + logger.info(colored("<<------------- Actuator ------------->>", "green")) + for i in range(self.mj_model.nu): + name = mujoco.mj_id2name(self.mj_model, mujoco._enums.mjtObj.mjOBJ_ACTUATOR, i) + if name: + logger.info(f"actuator_index: {i}, name: {name}") + print(" ") + + logger.info(colored("<<------------- Sensor ------------->>", "green")) + index = 0 + for i in range(self.mj_model.nsensor): + name = mujoco.mj_id2name(self.mj_model, mujoco._enums.mjtObj.mjOBJ_SENSOR, i) + if name: + logger.info( + f"sensor_index: {index}, name: {name}, dim: {self.mj_model.sensor_dim[i]}" + ) + index = index + self.mj_model.sensor_dim[i] + print(" ") + + +class ElasticBand: + """ + ref: https://github.com/unitreerobotics/unitree_mujoco + """ + + def __init__(self): + self.kp_pos = 10000 + self.kd_pos = 1000 + self.kp_ang = 1000 + self.kd_ang = 10 + self.point = np.array([0, 0, 1]) + self.length = 0 + self.enable = True + + def Advance(self, pose): + pos = pose[0:3] + quat = pose[3:7] + lin_vel = pose[7:10] + ang_vel = pose[10:13] + + δx = self.point - pos + f = self.kp_pos * (δx + np.array([0, 0, self.length])) + self.kd_pos * (0 - lin_vel) + + # Convert quaternion from MuJoCo [w,x,y,z] to scipy [x,y,z,w] + quat = np.array([quat[1], quat[2], quat[3], quat[0]]) + rot = scipy.spatial.transform.Rotation.from_quat(quat) + rotvec = rot.as_rotvec() + torque = -self.kp_ang * rotvec - self.kd_ang * ang_vel + + return np.concatenate([f, torque]) + + def MujuocoKeyCallback(self, key): + import glfw + + if key == glfw.KEY_7: + self.length -= 0.1 + if key == glfw.KEY_8: + self.length += 0.1 + if key == glfw.KEY_9: + self.enable = not self.enable + + def handle_keyboard_button(self, key): + if key == "9": + self.enable = not self.enable + print(f"ElasticBand enable: {self.enable}") diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/wbc_configs/g1_29dof_sonic_model12.yaml b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/wbc_configs/g1_29dof_sonic_model12.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7eaddda1105674ecb44399e74f987f8dec18b653 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/wbc_configs/g1_29dof_sonic_model12.yaml @@ -0,0 +1,419 @@ +# copy from g1_43dof_hist.yaml +ROBOT_TYPE: 'g1_29dof' # Robot name, "go2", "b2", "b2w", "h1", "go2w", "g1" +ROBOT_SCENE: "gear_sonic/data/robot_model/model_data/g1/scene_43dof.xml" # Robot scene, for Sim2Sim +# ROBOT_SCENE: "gear_sonic/data/robot_model/model_data/g1/scene_29dof_activated3dex.xml" + +DOMAIN_ID: 0 # Domain id +# Network Interface, "lo" for simulation and the one with "192.168.123.222" for real robot +# INTERFACE: "enxe8ea6a9c4e09" +# INTERFACE: "enxc8a3623c9cb7" +INTERFACE: "lo" +SIMULATOR: "mujoco" # "robocasa" + +USE_JOYSTICK: 0 # Simulate Unitree WirelessController using a gamepad (0: disable, 1: enable) +JOYSTICK_TYPE: "xbox" # support "xbox" and "switch" gamepad layout +JOYSTICK_DEVICE: 0 # Joystick number + +FREE_BASE: False + +PRINT_SCENE_INFORMATION: True # Print link, joint and sensors information of robot +ENABLE_ELASTIC_BAND: True # Virtual spring band, used for lifting h1 + +SIMULATE_DT: 0.005 # Need to be larger than the runtime of viewer.sync() +VIEWER_DT: 0.02 # Viewer update time +REWARD_DT: 0.02 +USE_SENSOR: False +USE_HISTORY: True +USE_HISTORY_LOCO: True +USE_HISTORY_MIMIC: True + +GAIT_PERIOD: 0.9 # 1.25 + +MOTOR2JOINT: [0, 1, 2, 3, 4, 5, + 6, 7, 8, 9, 10, 11, + 12, 13, 14, + 15, 16, 17, 18, 19, 20, 21, + 22, 23, 24, 25, 26, 27, 28] + +JOINT2MOTOR: [0, 1, 2, 3, 4, 5, + 6, 7, 8, 9, 10, 11, + 12, 13, 14, + 15, 16, 17, 18, 19, 20, 21, + 22, 23, 24, 25, 26, 27, 28] + + +UNITREE_LEGGED_CONST: + HIGHLEVEL: 0xEE + LOWLEVEL: 0xFF + TRIGERLEVEL: 0xF0 + PosStopF: 2146000000.0 + VelStopF: 16000.0 + MODE_MACHINE: 5 + MODE_PR: 0 + +JOINT_KP: [ + 100, 100, 100, 200, 20, 20, + 100, 100, 100, 200, 20, 20, + 400, 400, 400, + 90, 60, 20, 60, 4, 4, 4, + 90, 60, 20, 60, 4, 4, 4 +] + + +JOINT_KD: [ + 2.5, 2.5, 2.5, 5, 0.2, 0.1, + 2.5, 2.5, 2.5, 5, 0.2, 0.1, + 5.0, 5.0, 5.0, + 2.0, 1.0, 0.4, 1.0, 0.2, 0.2, 0.2, + 2.0, 1.0, 0.4, 1.0, 0.2, 0.2, 0.2 +] + +# arm kp +# soft kp, safe, test it first +# 50, 50, 20, 20, 10, 10, 10 +# hard kp, use only if policy is safe +# 200, 200, 80, 80, 50, 50, 50, + +# MOTOR_KP: [ +# 100, 100, 100, 200, 20, 20, +# 100, 100, 100, 200, 20, 20, +# 400, 400, 400, +# 50, 50, 20, 20, 10, 10, 10, +# 50, 50, 20, 20, 10, 10, 10 +# ] + +MOTOR_KP: [ + 150, 150, 150, 200, 40, 40, + 150, 150, 150, 200, 40, 40, + 250, 250, 250, + 100, 100, 40, 40, 20, 20, 20, + 100, 100, 40, 40, 20, 20, 20 +] + +MOTOR_KD: [ + 2, 2, 2, 4, 2, 2, + 2, 2, 2, 4, 2, 2, + 5, 5, 5, + 5, 5, 2, 2, 2, 2, 2, + 5, 5, 2, 2, 2, 2, 2 +] + +# MOTOR_KP: [ +# 100, 100, 100, 200, 20, 20, +# 100, 100, 100, 200, 20, 20, +# 400, 400, 400, +# 90, 60, 20, 60, 4, 4, 4, +# 90, 60, 20, 60, 4, 4, 4 +# ] + + +# MOTOR_KD: [ +# 2.5, 2.5, 2.5, 5, 0.2, 0.1, +# 2.5, 2.5, 2.5, 5, 0.2, 0.1, +# 5.0, 5.0, 5.0, +# 2.0, 1.0, 0.4, 1.0, 0.2, 0.2, 0.2, +# 2.0, 1.0, 0.4, 1.0, 0.2, 0.2, 0.2 +# ] + + +WeakMotorJointIndex: + left_hip_yaw_joint: 0 + left_hip_roll_joint: 1 + left_hip_pitch_joint: 2 + left_knee_joint: 3 + left_ankle_pitch_joint: 4 + left_ankle_roll_joint: 5 + right_hip_yaw_joint: 6 + right_hip_roll_joint: 7 + right_hip_pitch_joint: 8 + right_knee_joint: 9 + right_ankle_pitch_joint: 10 + right_ankle_roll_joint: 11 + waist_yaw_joint : 12 + waist_roll_joint : 13 + waist_pitch_joint : 14 + left_shoulder_pitch_joint: 15 + left_shoulder_roll_joint: 16 + left_shoulder_yaw_joint: 17 + left_elbow_joint: 18 + left_wrist_roll_joint: 19 + left_wrist_pitch_joint: 20 + left_wrist_yaw_joint: 21 + right_shoulder_pitch_joint: 22 + right_shoulder_roll_joint: 23 + right_shoulder_yaw_joint: 24 + right_elbow_joint: 25 + right_wrist_roll_joint: 26 + right_wrist_pitch_joint: 27 + right_wrist_yaw_joint: 28 + +NUM_MOTORS: 29 +NUM_JOINTS: 29 +NUM_HAND_MOTORS: 7 +NUM_HAND_JOINTS: 7 +NUM_UPPER_BODY_JOINTS: 17 + +DEFAULT_DOF_ANGLES: [ + -0.1, # left_hip_pitch_joint + 0.0, # left_hip_roll_joint + 0.0, # left_hip_yaw_joint + 0.3, # left_knee_joint + -0.2, # left_ankle_pitch_joint + 0.0, # left_ankle_roll_joint + -0.1, # right_hip_pitch_joint + 0.0, # right_hip_roll_joint + 0.0, # right_hip_yaw_joint + 0.3, # right_knee_joint + -0.2, # right_ankle_pitch_joint + 0.0, # right_ankle_roll_joint + 0.0, # waist_yaw_joint + 0.0, # waist_roll_joint + 0.0, # waist_pitch_joint + 0.0, # left_shoulder_pitch_joint + 0.0, # left_shoulder_roll_joint + 0.0, # left_shoulder_yaw_joint + 0.0, # left_elbow_joint + 0.0, # left_wrist_roll_joint + 0.0, # left_wrist_pitch_joint + 0.0, # left_wrist_yaw_joint + 0.0, # right_shoulder_pitch_joint + 0.0, # right_shoulder_roll_joint + 0.0, # right_shoulder_yaw_joint + 0.0, # right_elbow_joint + 0.0, # right_wrist_roll_joint + 0.0, # right_wrist_pitch_joint + 0.0 # right_wrist_yaw_joint +] + +DEFAULT_MOTOR_ANGLES: [ + -0.1, # left_hip_pitch_joint + 0.0, # left_hip_roll_joint + 0.0, # left_hip_yaw_joint + 0.3, # left_knee_joint + -0.2, # left_ankle_pitch_joint + 0.0, # left_ankle_roll_joint + -0.1, # right_hip_pitch_joint + 0.0, # right_hip_roll_joint + 0.0, # right_hip_yaw_joint + 0.3, # right_knee_joint + -0.2, # right_ankle_pitch_joint + 0.0, # right_ankle_roll_joint + 0.0, # waist_yaw_joint + 0.0, # waist_roll_joint + 0.0, # waist_pitch_joint + 0.0, # left_shoulder_pitch_joint + 0.0, # left_shoulder_roll_joint + 0.0, # left_shoulder_yaw_joint + 0.0, # left_elbow_joint + 0.0, # left_wrist_roll_joint + 0.0, # left_wrist_pitch_joint + 0.0, # left_wrist_yaw_joint + 0.0, # right_shoulder_pitch_joint + 0.0, # right_shoulder_roll_joint + 0.0, # right_shoulder_yaw_joint + 0.0, # right_elbow_joint + 0.0, # right_wrist_roll_joint + 0.0, # right_wrist_pitch_joint + 0.0 # right_wrist_yaw_joint +] + +motor_pos_lower_limit_list: [-2.5307, -0.5236, -2.7576, -0.087267, -0.87267, -0.2618, + -2.5307, -2.9671, -2.7576, -0.087267, -0.87267, -0.2618, + -2.618, -0.52, -0.52, + -3.0892, -1.5882, -2.618, -1.0472, + -1.972222054, -1.61443, -1.61443, + -3.0892, -2.2515, -2.618, -1.0472, + -1.972222054, -1.61443, -1.61443] +motor_pos_upper_limit_list: [2.8798, 2.9671, 2.7576, 2.8798, 0.5236, 0.2618, + 2.8798, 0.5236, 2.7576, 2.8798, 0.5236, 0.2618, + 2.618, 0.52, 0.52, + 2.6704, 2.2515, 2.618, 2.0944, + 1.972222054, 1.61443, 1.61443, + 2.6704, 1.5882, 2.618, 2.0944, + 1.972222054, 1.61443, 1.61443] +motor_vel_limit_list: [32.0, 32.0, 32.0, 20.0, 37.0, 37.0, + 32.0, 32.0, 32.0, 20.0, 37.0, 37.0, + 32.0, 37.0, 37.0, + 37.0, 37.0, 37.0, 37.0, + 37.0, 22.0, 22.0, + 37.0, 37.0, 37.0, 37.0, + 37.0, 22.0, 22.0] +motor_effort_limit_list: [88.0, 88.0, 88.0, 139.0, 50.0, 50.0, + 88.0, 88.0, 88.0, 139.0, 50.0, 50.0, + 88.0, 50.0, 50.0, + 25.0, 25.0, 25.0, 25.0, + 25.0, 5.0, 5.0, + 2.45, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7, + 25.0, 25.0, 25.0, 25.0, + 25.0, 5.0, 5.0, + 2.45, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7] +history_config: { + base_ang_vel: 4, + projected_gravity: 4, + command_lin_vel: 4, + command_ang_vel: 4, + command_base_height: 4, + command_stand: 4, + ref_upper_dof_pos: 4, + dof_pos: 4, + dof_vel: 4, + actions: 4, + # phase_time: 4, + ref_motion_phase: 4, + sin_phase: 4, + cos_phase: 4 + } +history_loco_config: { + base_ang_vel: 4, + projected_gravity: 4, + command_lin_vel: 4, + command_ang_vel: 4, + # command_base_height: 4, + command_stand: 4, + ref_upper_dof_pos: 4, + dof_pos: 4, + dof_vel: 4, + actions: 4, + # phase_time: 4, + sin_phase: 4, + cos_phase: 4 + } +history_loco_height_config: { + base_ang_vel: 4, + projected_gravity: 4, + command_lin_vel: 4, + command_ang_vel: 4, + command_base_height: 4, + command_stand: 4, + ref_upper_dof_pos: 4, + dof_pos: 4, + dof_vel: 4, + actions: 4, + # phase_time: 4, + sin_phase: 4, + cos_phase: 4 + } +history_mimic_config: { + base_ang_vel: 4, + projected_gravity: 4, + dof_pos: 4, + dof_vel: 4, + actions: 4, + ref_motion_phase: 4, + } +obs_dims: { + base_lin_vel: 3, + base_ang_vel: 3, + projected_gravity: 3, + command_lin_vel: 2, + command_ang_vel: 1, + command_stand: 1, + command_base_height: 1, + ref_upper_dof_pos: 17, # upper body actions + dof_pos: 29, + dof_vel: 29, + # actions: 12, # lower body actions + actions: 29, # full body actions + phase_time: 1, + ref_motion_phase: 1, # mimic motion phase + sin_phase: 1, + cos_phase: 1, + } +obs_loco_dims: { + base_lin_vel: 3, + base_ang_vel: 3, + projected_gravity: 3, + command_lin_vel: 2, + command_ang_vel: 1, + command_stand: 1, + command_base_height: 1, + ref_upper_dof_pos: 17, # upper body actions + dof_pos: 29, + dof_vel: 29, + actions: 12, # lower body actions + phase_time: 1, + sin_phase: 1, + cos_phase: 1, + } +obs_mimic_dims: { + base_lin_vel: 3, + base_ang_vel: 3, + projected_gravity: 3, + dof_pos: 29, + dof_vel: 29, + actions: 29, # full body actions + ref_motion_phase: 1, # mimic motion phase + } +obs_scales: { + base_lin_vel: 2.0, + base_ang_vel: 0.25, + projected_gravity: 1.0, + command_lin_vel: 1, + command_ang_vel: 1, + command_stand: 1, + command_base_height: 2, # Yuanhang: it's 2, not 1! + ref_upper_dof_pos: 1.0, + dof_pos: 1.0, + dof_vel: 0.05, + history: 1.0, + history_loco: 1.0, + history_mimic: 1.0, + actions: 1.0, + phase_time: 1.0, + ref_motion_phase: 1.0, + sin_phase: 1.0, + cos_phase: 1.0 + } + +loco_upper_body_dof_pos: [ + 0.0, 0.0, 0.0, # waist + 0.0, 0.3, 0.0, 1.0, # left shoulder and elbow + 0.0, 0.0, 0.0, # left wrist + 0.0, -0.3, 0.0, 1.0, # right shoulder and elbow + 0.0, 0.0, 0.0 # right wrist +] + +robot_dofs: { + "g1_29dof": [1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, + 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1], + "g1_29dof_anneal_23dof": [1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, + 1, 1, 1, + 1, 1, 1, 1, 0, 0, 0, + 1, 1, 1, 1, 0, 0, 0], +} + +mimic_robot_types: { + + "APT_level1": "g1_29dof_anneal_23dof", +} + + + + +# 01281657 +mimic_models: { + "APT_level1": "20250116_225127-TairanTestbed_G129dofANNEAL23dof_dm_APT_video_APT_level1_MinimalFriction-0.3_RfiTrue_Far0.325_RESUME_LARGENOISE-motion_tracking-g1_29dof_anneal_23dof/exported/model_176500.onnx", + +} + + + +start_upper_body_dof_pos: { + + "APT_level1": + [0.19964170455932617, 0.07710712403059006, -0.2882401943206787, + 0.21672365069389343, 0.15629297494888306, -0.5167576670646667, 0.5782126784324646, + 0.0, 0.0, 0.0, + 0.25740593671798706, -0.2504104673862457, 0.22500675916671753, 0.5127624273300171, + 0.0, 0.0, 0.0], + +} + +motion_length_s: { + "APT_level1": 7.66, + +} diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/network/network_utils.py b/GR00T-WholeBodyControl/gear_sonic/utils/network/network_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..105a5cd2e7958e3732cf3ed135576b5d90b81e30 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/network/network_utils.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python3 +""" +Simple network interface utilities +""" + +import platform +import re +import subprocess + + +def get_network_interfaces(): + """Get network interfaces with their IP addresses""" + try: + result = subprocess.run( + ["/sbin/ip", "addr", "show"], capture_output=True, text=True, check=True + ) + return _parse_ip_output(result.stdout) + except (subprocess.CalledProcessError, FileNotFoundError): + try: + result = subprocess.run(["ifconfig"], capture_output=True, text=True, check=True) + return _parse_ifconfig_output(result.stdout) + except (subprocess.CalledProcessError, FileNotFoundError): + return {} + + +def _parse_ip_output(output): + """Parse 'ip addr' command output""" + interfaces = {} + current_interface = None + + for line in output.split("\n"): + interface_match = re.match(r"^\d+:\s+(\w+):", line) + if interface_match: + current_interface = interface_match.group(1) + interfaces[current_interface] = [] + + ip_match = re.search(r"inet\s+(\d+\.\d+\.\d+\.\d+)", line) + if ip_match and current_interface: + interfaces[current_interface].append(ip_match.group(1)) + + return interfaces + + +def _parse_ifconfig_output(output): + """Parse 'ifconfig' command output""" + interfaces = {} + current_interface = None + + for line in output.split("\n"): + interface_match = re.match(r"^(\w+):", line) + if interface_match: + current_interface = interface_match.group(1) + interfaces[current_interface] = [] + + ip_match = re.search(r"inet\s+(\d+\.\d+\.\d+\.\d+)", line) + if ip_match and current_interface: + interfaces[current_interface].append(ip_match.group(1)) + + return interfaces + + +def find_interface_by_ip(target_ip): + """Find interface name for given IP address""" + interfaces = get_network_interfaces() + for interface, ip_list in interfaces.items(): + if target_ip in ip_list: + return interface + return None + + +def resolve_interface(interface: str) -> tuple[str, str]: + """ + Resolve interface parameter to actual network interface name and environment type + + Args: + interface: "sim", "real", or direct interface name or IP address + + Returns: + tuple: (interface_name, env_type) where env_type is "sim" or "real" + """ + # Check if interface is an IP address + if re.match(r"^\d+\.\d+\.\d+\.\d+$", interface): + if interface == "127.0.0.1": + return interface, "sim" + else: + return interface, "real" + + if interface == "sim": + lo_interface = find_interface_by_ip("127.0.0.1") + if lo_interface: + # macOS uses lo0 instead of lo + if platform.system() == "Darwin" and lo_interface == "lo": + return "lo0", "sim" + return lo_interface, "sim" + return ("lo0" if platform.system() == "Darwin" else "lo"), "sim" + + elif interface == "real": + interfaces = get_network_interfaces() + for iface, ip_list in interfaces.items(): + for ip in ip_list: + if ip.startswith("192.168.123."): + return iface, "real" + return interface, "real" # fallback + + else: + # Direct interface name - check if it has 127.0.0.1 to determine env_type + interfaces = get_network_interfaces() + if interface in interfaces: + for ip in interfaces[interface]: + if ip == "127.0.0.1": + return interface, "sim" + + # macOS lo interface handling + if platform.system() == "Darwin" and interface == "lo": + return "lo0", "sim" + + # Default to real for unknown interfaces + return interface, "real" + + +if __name__ == "__main__": + interfaces = get_network_interfaces() + + if not interfaces: + print("No network interfaces found") + exit(1) + + # Show all interfaces + print("Network interfaces:") + for interface, ip_list in interfaces.items(): + print(f" {interface}: {', '.join(ip_list)}") + + # Test resolve_interface function + print("\nTesting resolve_interface:") + for test_interface in ["sim", "real", "lo", "127.0.0.1"]: + interface_name, env_type = resolve_interface(test_interface) + print(f" {test_interface} -> {interface_name} ({env_type})") diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/obs_utils.py b/GR00T-WholeBodyControl/gear_sonic/utils/obs_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6e63220c5bd195289232174f58e3e3e8a53538d2 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/obs_utils.py @@ -0,0 +1,92 @@ +""" +Utility functions for observation processing and indexing. +""" + +import numpy as np + + +def get_obs_index_map(observation_manager): + """ + Compute a dictionary that maps each observation term of each group to the corresponding + start and end indices in the observation tensor. + + Args: + group_obs_term_dim (dict): Dictionary with group names as keys and lists of dimension tuples as values + e.g., {'policy': [(15,), (20,), ...], 'critic': [(58,), (3,), ...]} + group_obs_term_names (dict): Dictionary with group names as keys and lists of observation term names as values + e.g., {'policy': ['root_pos_multi_future', 'root_quat_multi_future', ...], + 'critic': ['command', 'motion_anchor_pos_b', ...]} + + Returns: + dict: Nested dictionary mapping group -> obs_term -> (start_idx, end_idx) + e.g., {'policy': {'root_pos_multi_future': (0, 15), 'root_quat_multi_future': (15, 35), ...}, + 'critic': {'command': (0, 58), 'motion_anchor_pos_b': (58, 61), ...}} + """ + obs_index_map = {} + group_obs_term_dim = observation_manager._group_obs_term_dim + group_obs_term_names = observation_manager._group_obs_term_names + + for group_name in group_obs_term_dim.keys(): + obs_index_map[group_name] = {} + + # Get dimensions and names for this group + dims = group_obs_term_dim[group_name] + names = group_obs_term_names[group_name] + + # Ensure dimensions and names lists have the same length + assert len(dims) == len( + names + ), f"Mismatch in group '{group_name}': {len(dims)} dims vs {len(names)} names" + + # Compute cumulative indices + current_idx = 0 + for i, (dim_tuple, obs_name) in enumerate(zip(dims, names)): + # Extract the actual dimension from the tuple (assuming single dimension per tuple) + dim = ( + dim_tuple[0] if isinstance(dim_tuple, tuple) and len(dim_tuple) == 1 else dim_tuple + ) + + start_idx = current_idx + end_idx = current_idx + dim + + obs_index_map[group_name][obs_name] = (start_idx, end_idx) + current_idx = end_idx + + return obs_index_map + + +def get_group_obs_shape(observation_manager, group_name): + group_obs_term_dim = observation_manager.group_obs_term_dim[group_name] + total_dim = sum([dim[-1] for dim in group_obs_term_dim]) + group_obs_first_shape = group_obs_term_dim[0] + group_obs_shape = tuple(group_obs_first_shape[:-1]) + (total_dim,) + return group_obs_shape + + +def get_group_term_obs_shape(example_obs, group_name): + """Get observation shapes for a group. + + Handles both cases: + - Dict observations (concatenate_terms: False) - returns individual term dims/names + - Tensor observations (concatenate_terms: True) - returns total dim only + """ + obs_data = example_obs[group_name] + + # Handle case where observation is already concatenated to a tensor + # (when concatenate_terms: True in observation group config) + if not isinstance(obs_data, dict): + # obs_data is a tensor, not a dict + group_obs_total_dim = int(np.prod(obs_data.shape[1:]).item()) + # Return single entry with the group name as key + group_obs_dims = {group_name: tuple(obs_data.shape[1:])} + group_obs_names = [group_name] + return group_obs_dims, group_obs_names, group_obs_total_dim + + # Original behavior for dict observations + group_obs_dims = {} + group_obs_names = list(obs_data.keys()) + group_obs_total_dim = 0 + for key, value in obs_data.items(): + group_obs_dims[key] = tuple(value.shape[1:]) + group_obs_total_dim += np.prod(group_obs_dims[key]).item() + return group_obs_dims, group_obs_names, group_obs_total_dim diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/running_mean_std.py b/GR00T-WholeBodyControl/gear_sonic/utils/running_mean_std.py new file mode 100644 index 0000000000000000000000000000000000000000..6039aed7d7a1a1e1ccf49187f732cb63a3408e0b --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/running_mean_std.py @@ -0,0 +1,293 @@ +import torch +import torch.nn as nn + +""" +updates statistic from a full data + +Memory optimization note: If you encounter CUDA out of memory errors, consider setting: +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +This can help reduce memory fragmentation. +""" + + +class RunningMeanStd(nn.Module): + + def __init__(self, insize, epsilon=1e-05, per_channel=False, norm_only=False): + super().__init__() + print("RunningMeanStd: ", insize) + self.insize = insize + self.mean_size = insize[0] + self.epsilon = epsilon + + self.norm_only = norm_only + self.per_channel = per_channel + if per_channel: + if len(self.insize) == 3: + self.axis = [0, 2, 3] + if len(self.insize) == 2: + self.axis = [0, 2] + if len(self.insize) == 1: + self.axis = [0] + in_size = self.insize[0] + else: + self.axis = [0] + in_size = insize + + self.register_buffer("running_mean", torch.zeros(in_size, dtype=torch.float32)) + self.register_buffer("running_var", torch.ones(in_size, dtype=torch.float32)) + self.register_buffer("count", torch.ones((), dtype=torch.float32)) + + self.frozen = False + self.frozen_partial = False + + def freeze(self): + self.frozen = True + + def unfreeze(self): + self.frozen = False + + def freeze_partial(self, diff): + self.frozen_partial = True + self.diff = diff + + def sync_across_gpus(self, accelerator): + if accelerator.num_processes <= 1: + return + # ZL: this is the right formulation but a crude sync. The correct way is to sync + # batch stats. + flat_stats = torch.cat( + [ + self.running_mean.flatten(), # (D,) + self.running_var.flatten(), # (D,) + ] + ) # shape = (2D + C,) + + # Gather stats from all processes + gathered = accelerator.gather(flat_stats[None]) # shape = (world_size, 2D + C) + + # Reshape + world_size = gathered.shape[0] + D = self.running_mean.numel() + + means_ = gathered[:, :D].reshape(world_size, *self.running_mean.shape).mean(dim=0) + vars_ = gathered[:, D : 2 * D].reshape(world_size, *self.running_var.shape).mean(dim=0) + + # Update local stats + self.running_mean.copy_(means_) + self.running_var.copy_(vars_) + + def _update_mean_var_count_from_moments( + self, mean, var, count, batch_mean, batch_var, batch_count + ): + delta = batch_mean - mean + tot_count = count + batch_count + + new_mean = mean + delta * batch_count / tot_count + m_a = var * count + m_b = batch_var * batch_count + M2 = m_a + m_b + delta**2 * count * batch_count / tot_count + new_var = M2 / tot_count + new_count = tot_count + return new_mean, new_var, new_count + + def forward(self, input, unnorm=False): + # change shape + input_shape = input.shape + if len(input.shape) == 3: + input = input.reshape(-1, input_shape[-1]) + + if self.per_channel: + if len(self.insize) == 3: + # Use broadcasting instead of expand_as to avoid memory issues + current_mean = self.running_mean.view([1, self.insize[0], 1, 1]) + current_var = self.running_var.view([1, self.insize[0], 1, 1]) + elif len(self.insize) == 2: + current_mean = self.running_mean.view([1, self.insize[0], 1]) + current_var = self.running_var.view([1, self.insize[0], 1]) + elif len(self.insize) == 1: + current_mean = self.running_mean.view([1, self.insize[0]]) + current_var = self.running_var.view([1, self.insize[0]]) + + else: + current_mean = self.running_mean + current_var = self.running_var + # get output + + if unnorm: + y = torch.clamp(input, min=-5.0, max=5.0) + y = torch.sqrt(current_var + self.epsilon) * y + current_mean + else: + if self.norm_only: + y = input / torch.sqrt(current_var + self.epsilon) + else: + # Use in-place operations where possible to reduce memory usage + y = input - current_mean + y = y / torch.sqrt(current_var + self.epsilon) + y = torch.clamp(y, min=-5.0, max=5.0) + + # update After normalization, so that the values used for training and testing are the same. + if self.training and not self.frozen: + mean = input.mean(self.axis) # along channel axis + var = input.var(self.axis) + + new_mean, new_var, new_count = self._update_mean_var_count_from_moments( + self.running_mean, self.running_var, self.count, mean, var, input.size()[0] + ) + if self.frozen_partial: + # Only update the last bit (futures) + self.running_mean[-self.diff :], self.running_var[-self.diff :], self.count = ( + new_mean[-self.diff :], + new_var[-self.diff :], + new_count, + ) + else: + self.running_mean, self.running_var, self.count = new_mean, new_var, new_count + + if len(input_shape) == 3: + y = y.view(input_shape) + + return y + + +class RunningMeanStdObs(nn.Module): + + def __init__(self, insize, epsilon=1e-05, per_channel=False, norm_only=False): + assert isinstance(insize, dict) + super().__init__() + self.running_mean_std = nn.ModuleDict( + {k: RunningMeanStd(v, epsilon, per_channel, norm_only) for k, v in insize.items()} + ) + + def forward(self, input, unnorm=False): + res = {k: self.running_mean_std[k](v, unnorm) for k, v in input.items()} + return res + + +from collections.abc import Sequence +from copy import deepcopy + + +class VecNorm(nn.Module): + """Simple running normalization for observations. + + Keeps track of running mean and variance to normalize observations on-the-fly. + + Args: + obs_keys: List of observation keys to normalize + decay: Decay rate for running statistics (default: 0.99) + eps: Small constant for numerical stability (default: 1e-4) + device: Device to store statistics on + """ + + def __init__( + self, + obs_keys: Sequence[str] | None, + decay: float = 0.9999, + eps: float = 1e-4, + ): + super().__init__() + self.obs_keys = obs_keys + self.decay = decay + self.eps = eps + + # Running statistics as nn.Parameters + self.sum: dict[str, nn.Parameter] = nn.ParameterDict({}) + self.ssq: dict[str, nn.Parameter] = nn.ParameterDict({}) + self.cnt: dict[str, nn.Parameter] = nn.ParameterDict({}) + self.mean: dict[str, nn.Parameter] = nn.ParameterDict({}) + self.var: dict[str, nn.Parameter] = nn.ParameterDict({}) + + self.initialized = False + self.frozen = False + + def init_stats(self, obs_dict: dict[str, torch.Tensor]): + """Initialize running statistics based on observation shapes.""" + if self.obs_keys is None: + self.obs_keys = list(obs_dict.keys()) + for key in self.obs_keys: + if key in obs_dict: + v = obs_dict[key] + self.sum[key] = nn.Parameter(torch.zeros_like(v[0]), requires_grad=False) + self.ssq[key] = nn.Parameter(torch.zeros_like(v[0]), requires_grad=False) + self.cnt[key] = nn.Parameter(torch.zeros_like(v[0]), requires_grad=False) + self.mean[key] = nn.Parameter(torch.zeros_like(v[0]), requires_grad=False) + self.var[key] = nn.Parameter(torch.ones_like(v[0]), requires_grad=False) + self.initialized = True + + def update(self, obs_dict: dict[str, torch.Tensor]): + """Update running statistics with new observations.""" + if not self.initialized: + self.init_stats(obs_dict) + + if self.frozen: + return + + for key in self.obs_keys: + if key not in obs_dict: + continue + + x = obs_dict[key] + sum_x = x.sum(dim=0) + ssq_x = (x**2).sum(dim=0) + cnt_x = x.shape[0] + + self.sum[key] = self.sum[key] * self.decay + sum_x + self.ssq[key] = self.ssq[key] * self.decay + ssq_x + self.cnt[key] = self.cnt[key] * self.decay + cnt_x + + self.mean[key] = self.sum[key] / self.cnt[key] + self.var[key] = (self.ssq[key] / self.cnt[key] - self.mean[key] ** 2).clamp( + min=self.eps + ) + + def normalize(self, obs_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Normalize observations using running statistics.""" + if not self.initialized: + self.init_stats(obs_dict) + return obs_dict.copy() + + normalized = obs_dict.copy() + for key in self.obs_keys: + normalized[key] = (normalized[key] - self.mean[key]) / self.var[key].sqrt() + + return normalized + + def denormalize(self, obs_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Denormalize observations using running statistics.""" + if not self.initialized: + self.init_stats(obs_dict) + return obs_dict.copy() + + denormalized = obs_dict.copy() + for key in self.obs_keys: + denormalized[key] = denormalized[key] * self.var[key].sqrt() + self.mean[key] + + return denormalized + + def freeze(self): + """Freeze running statistics updates.""" + self.frozen = True + + def unfreeze(self): + """Unfreeze running statistics updates.""" + self.frozen = False + + def get_stats(self): + """Get current running statistics.""" + return { + "sum": {key: param.data.clone() for key, param in self.sum.items()}, + "ssq": {key: param.data.clone() for key, param in self.ssq.items()}, + "cnt": {key: param.data.clone() for key, param in self.cnt.items()}, + "mean": {key: param.data.clone() for key, param in self.mean.items()}, + "var": {key: param.data.clone() for key, param in self.var.items()}, + } + + def load_stats(self, stats): + """Load running statistics.""" + for key in stats["sum"]: + self.sum[key] = nn.Parameter(deepcopy(stats["sum"][key]), requires_grad=False) + self.ssq[key] = nn.Parameter(deepcopy(stats["ssq"][key]), requires_grad=False) + self.cnt[key] = nn.Parameter(deepcopy(stats["cnt"][key]), requires_grad=False) + self.mean[key] = nn.Parameter(deepcopy(stats["mean"][key]), requires_grad=False) + self.var[key] = nn.Parameter(deepcopy(stats["var"][key]), requires_grad=False) + self.initialized = True diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/teleop/input_readers.py b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/input_readers.py new file mode 100644 index 0000000000000000000000000000000000000000..47a6dc43728c7ff314abc721fc68827dee32b6b1 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/input_readers.py @@ -0,0 +1,493 @@ +"""Input source readers for body tracking data. + +PicoReader -- pulls data from XRoboToolkit SDK (Pico headset). +IsaacTeleopReader -- in-process IsaacTeleop / CloudXR DeviceIO session. +""" + +import logging +import threading +import time +from typing import Any + +import numpy as np + +logger = logging.getLogger(__name__) + +try: + import xrobotoolkit_sdk as xrt +except ImportError: + xrt = None + +try: + from gear_sonic.utils.teleop.isaac_teleop_client import IsaacTeleopClient +except ImportError: + IsaacTeleopClient = None + + +class PicoReader: + """Background reader that pulls Pico/XRT data and computes dt/FPS.""" + + STALE_TIMEOUT = 5.0 + + def __init__(self, max_queue_size: int = 15): + del max_queue_size + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, daemon=True) + self._fps_ema = 0.0 + self._last_stamp_ns = None + self._latest = None + self._lock = threading.Lock() + self._last_new_data_time = time.monotonic() + self._disconnected = threading.Event() + + def start(self): + if not self._thread.is_alive(): + self._thread.start() + + def stop(self): + self._stop.set() + if self._thread.is_alive(): + self._thread.join(timeout=1.0) + + def get_latest(self): + with self._lock: + return self._latest + + @property + def disconnected(self) -> bool: + return self._disconnected.is_set() + + def clear_disconnect(self): + self._disconnected.clear() + self._last_new_data_time = time.monotonic() + self._last_stamp_ns = None + self._fps_ema = 0.0 + + def get_timestamp_ns(self) -> int: + if xrt is None: + return 0 + return int(xrt.get_time_stamp_ns()) + + def _run(self): + last_report = time.time() + while not self._stop.is_set(): + if xrt is None or not xrt.is_body_data_available(): + if ( + time.monotonic() - self._last_new_data_time > self.STALE_TIMEOUT + and not self._disconnected.is_set() + ): + logger.warning( + "[PicoReader] No new data for %.1fs, flagging disconnect", + self.STALE_TIMEOUT, + ) + self._disconnected.set() + time.sleep(0.001) + continue + + stamp_ns = xrt.get_time_stamp_ns() + prev_stamp_ns = self._last_stamp_ns + if prev_stamp_ns is not None and stamp_ns == prev_stamp_ns: + if ( + time.monotonic() - self._last_new_data_time > self.STALE_TIMEOUT + and not self._disconnected.is_set() + ): + logger.warning( + "[PicoReader] Timestamps stale for %.1fs, flagging disconnect", + self.STALE_TIMEOUT, + ) + self._disconnected.set() + time.sleep(0.000001) + continue + + self._last_new_data_time = time.monotonic() + if self._disconnected.is_set(): + logger.info("[PicoReader] Fresh data received, connection restored") + self._disconnected.clear() + + device_dt = ((stamp_ns - prev_stamp_ns) * 1e-9) if prev_stamp_ns is not None else 0.0 + if device_dt > 0.0: + inst = 1.0 / device_dt + self._fps_ema = inst if self._fps_ema == 0.0 else (0.9 * self._fps_ema + 0.1 * inst) + self._last_stamp_ns = stamp_ns + + try: + body_poses = xrt.get_body_joints_pose() + sample = { + "body_poses_np": np.array(body_poses), + "timestamp_realtime": time.time(), + "timestamp_monotonic": time.monotonic(), + "timestamp_ns": stamp_ns, + "dt": device_dt, + "fps": self._fps_ema, + } + with self._lock: + self._latest = sample + + now = time.time() + if now - last_report >= 5.0: + logger.info( + "[PicoReader] dt_ts: %.2f ms, fps: %.2f", + device_dt * 1000.0, + self._fps_ema, + ) + last_report = now + except Exception: + logger.exception("[PicoReader] read error") + + +def _attr_or_item(obj: Any, name: str, default: Any = None) -> Any: + """Return ``obj.`` if present, else ``obj[]`` if dict-like, else ``default``.""" + if obj is None: + return default + sentinel = object() + val = getattr(obj, name, sentinel) + if val is not sentinel: + return val + if hasattr(obj, "get"): + try: + return obj.get(name, default) + except Exception: + return default + return default + + +def _vec3(point: Any) -> tuple[float, float, float] | None: + """Extract (x, y, z) from a point-like (.x/.y/.z attrs or 3-sequence).""" + if point is None: + return None + x = _attr_or_item(point, "x") + y = _attr_or_item(point, "y") + z = _attr_or_item(point, "z") + if x is not None and y is not None and z is not None: + return float(x), float(y), float(z) + try: + return float(point[0]), float(point[1]), float(point[2]) + except Exception: + return None + + +def _quat_xyzw(orientation: Any) -> tuple[float, float, float, float] | None: + """Extract (qx, qy, qz, qw) from an orientation-like.""" + if orientation is None: + return None + qx = _attr_or_item(orientation, "x") + qy = _attr_or_item(orientation, "y") + qz = _attr_or_item(orientation, "z") + qw = _attr_or_item(orientation, "w") + if all(v is not None for v in (qx, qy, qz, qw)): + return float(qx), float(qy), float(qz), float(qw) + try: + return ( + float(orientation[0]), + float(orientation[1]), + float(orientation[2]), + float(orientation[3]), + ) + except Exception: + return None + + +# Number of joints in the IsaacTeleop FullBodyPosePicoT (XR_BD_body_tracking). +# Mirrors core.BodyJointPico.NUM_JOINTS in IsaacTeleop's schema bindings. +_NUM_BODY_JOINTS = 24 + +_UNRECOGNISED_SCHEMA_LOGGED: set[str] = set() + + +def _log_unrecognised_schema_once(body_data: Any) -> None: + """One-shot diagnostic if ``body_data`` doesn't look like either schema we + expect. Logs the type once per process so it doesn't flood the streamer. + """ + type_name = type(body_data).__name__ + if type_name in _UNRECOGNISED_SCHEMA_LOGGED: + return + _UNRECOGNISED_SCHEMA_LOGGED.add(type_name) + attrs = sorted(a for a in dir(body_data) if not a.startswith("_"))[:25] + logger.warning( + "[IsaacTeleopReader] Unrecognised body_data schema: type=%s attrs=%s. " + "Update _body_data_to_24x7() to handle this layout.", + type_name, + attrs, + ) + + +def _body_data_to_24x7(body_data: Any) -> np.ndarray | None: + """Convert ``FullBodyTrackerPico.get_body_pose().data`` to a (24, 7) array. + + Returns ``None`` while no joint is valid (typical when the headset isn't + connected yet — every ``BodyJointPose.is_valid`` is False, the streamer + keeps polling and the C++ deploy doesn't see fake zero pose). + + Two accepted schemas: + + Schema A — IsaacTeleop ``FullBodyPosePicoT`` (DeviceIO direct). + Defined in IsaacTeleop's ``schema/full_body.fbs`` / + ``schema/python/full_body_bindings.h``:: + + FullBodyPosePicoT.joints → BodyJointsPico (attr) + BodyJointsPico.joints(index) → BodyJointPose (METHOD; index 0..23) + BodyJointPose.is_valid → bool + BodyJointPose.pose.position → Point (.x .y .z) + BodyJointPose.pose.orientation → Quaternion (.x .y .z .w) + + Schema B — msgpack wire format published by ``teleop_ros2_ref`` (kept for + compatibility with ROS2 bridges; consumed when ``body_data`` already + looks like a dict with ``joint_positions`` / ``joint_orientations``). + """ + if body_data is None: + return None + + # Schema B: msgpack wire format (teleop_ros2_ref-compatible). + positions = _attr_or_item(body_data, "joint_positions") + orientations = _attr_or_item(body_data, "joint_orientations") + if positions is not None and orientations is not None: + n = min(len(positions), len(orientations), _NUM_BODY_JOINTS) + if n == 0: + return None + body_poses = np.zeros((_NUM_BODY_JOINTS, 7), dtype=np.float32) + for i in range(n): + pos = _vec3(positions[i]) + quat = _quat_xyzw(orientations[i]) + if pos is None or quat is None: + continue + body_poses[i, :3] = pos + body_poses[i, 3:] = quat + return body_poses + + # Schema A: native FullBodyPosePicoT — joints exposed via + # BodyJointsPico.joints(index) method (one BodyJointPose per call). + joints_container = getattr(body_data, "joints", None) + if joints_container is None: + _log_unrecognised_schema_once(body_data) + return None + get_joint = getattr(joints_container, "joints", None) + if not callable(get_joint): + _log_unrecognised_schema_once(body_data) + return None + + body_poses = np.zeros((_NUM_BODY_JOINTS, 7), dtype=np.float32) + any_valid = False + for i in range(_NUM_BODY_JOINTS): + try: + joint = get_joint(i) + except Exception: + continue + if joint is None: + continue + # Older builds may omit is_valid — default to True so we don't drop + # samples on schema drift; per-field validity falls out below. + if not getattr(joint, "is_valid", True): + continue + pose = getattr(joint, "pose", None) + if pose is None: + continue + pos = _vec3(getattr(pose, "position", None)) + quat = _quat_xyzw(getattr(pose, "orientation", None)) + if pos is None or quat is None: + continue + body_poses[i, :3] = pos + body_poses[i, 3:] = quat + any_valid = True + + return body_poses if any_valid else None + + +def _controller_inputs_to_dict_side(snapshot: Any) -> dict[str, Any] | None: + """Project one ControllerSnapshot.inputs into the dict shape consumed by helpers.""" + if snapshot is None: + return None + inputs = _attr_or_item(snapshot, "inputs") + if inputs is None: + return None + return { + "trigger_value": float(_attr_or_item(inputs, "trigger_value", 0.0) or 0.0), + "squeeze_value": float(_attr_or_item(inputs, "squeeze_value", 0.0) or 0.0), + "thumbstick_x": float(_attr_or_item(inputs, "thumbstick_x", 0.0) or 0.0), + "thumbstick_y": float(_attr_or_item(inputs, "thumbstick_y", 0.0) or 0.0), + "thumbstick_click": float(_attr_or_item(inputs, "thumbstick_click", 0.0) or 0.0), + "primary_click": float(_attr_or_item(inputs, "primary_click", 0.0) or 0.0), + "secondary_click": float(_attr_or_item(inputs, "secondary_click", 0.0) or 0.0), + } + + +def _build_controller_dict(raw: dict[str, Any] | None) -> dict[str, Any] | None: + """Convert ``IsaacTeleopClient._get_tracker_data()`` into the controller dict + schema that ``pico_manager_thread_server`` consumes (left/right trigger, + squeeze, thumbstick, click, primary/secondary click).""" + if raw is None: + return None + + left = _controller_inputs_to_dict_side(raw.get("left_controller")) + right = _controller_inputs_to_dict_side(raw.get("right_controller")) + if left is None and right is None: + return None + + out: dict[str, Any] = {} + if left is not None: + out["left_trigger_value"] = left["trigger_value"] + out["left_squeeze_value"] = left["squeeze_value"] + out["left_thumbstick"] = [left["thumbstick_x"], left["thumbstick_y"]] + out["left_thumbstick_click"] = left["thumbstick_click"] + out["left_primary_click"] = left["primary_click"] + out["left_secondary_click"] = left["secondary_click"] + if right is not None: + out["right_trigger_value"] = right["trigger_value"] + out["right_squeeze_value"] = right["squeeze_value"] + out["right_thumbstick"] = [right["thumbstick_x"], right["thumbstick_y"]] + out["right_thumbstick_click"] = right["thumbstick_click"] + out["right_primary_click"] = right["primary_click"] + out["right_secondary_click"] = right["secondary_click"] + return out + + +class IsaacTeleopReader: + """Background reader using the in-process IsaacTeleop / CloudXR DeviceIO session. + + Drop-in alternative to ``PicoReader`` — same ``get_latest()`` / + ``get_controller_data()`` contract. Hosts the CloudXR runtime in-process + via :class:`IsaacTeleopClient` (no separate publisher container, no host + ``~/.cloudxr`` sharing required). + """ + + STALE_TIMEOUT = 5.0 + + def __init__( + self, + max_queue_size: int = 15, + use_adb: bool = False, + poll_hz: float = 90.0, + ): + del max_queue_size + + if IsaacTeleopClient is None: + raise RuntimeError( + "isaacteleop is required for --input-source isaac-teleop but was not " + "found. Install via install_scripts/install_pico.sh, which runs:\n" + " uv pip install 'isaacteleop[cloudxr]~=1.3.0' --prerelease=allow " + "--extra-index-url https://pypi.nvidia.com" + ) + + self._client = IsaacTeleopClient(use_adb=use_adb) + self._period = 1.0 / max(1.0, float(poll_hz)) + + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, daemon=True) + self._lock = threading.Lock() + self._ctrl_lock = threading.Lock() + self._latest: dict[str, Any] | None = None + self._latest_controller: dict[str, Any] | None = None + self._fps_ema = 0.0 + self._last_stamp_ns: int | None = None + self._last_new_data_time = time.monotonic() + self._disconnected = threading.Event() + self._unrecognised_logged = False + + def start(self) -> None: + self._client.start_streaming() + if not self._thread.is_alive(): + self._thread.start() + + def stop(self) -> None: + self._stop.set() + if self._thread.is_alive(): + self._thread.join(timeout=1.0) + try: + self._client.close() + except Exception: + logger.exception("Failed to close IsaacTeleopClient cleanly") + + def get_latest(self) -> dict[str, Any] | None: + with self._lock: + return self._latest + + def get_controller_data(self) -> dict[str, Any] | None: + with self._ctrl_lock: + return self._latest_controller + + @property + def disconnected(self) -> bool: + return self._disconnected.is_set() + + def clear_disconnect(self) -> None: + self._disconnected.clear() + self._last_new_data_time = time.monotonic() + self._last_stamp_ns = None + self._fps_ema = 0.0 + + def get_timestamp_ns(self) -> int: + with self._lock: + sample = self._latest + return int(sample["timestamp_ns"]) if sample else 0 + + def _run(self) -> None: + last_report = time.time() + while not self._stop.is_set(): + try: + raw = self._client._get_tracker_data() # noqa: SLF001 — internal API by design + except Exception: + logger.exception("[IsaacTeleopReader] DeviceIO update failed") + time.sleep(self._period) + continue + + if raw is None: + if ( + time.monotonic() - self._last_new_data_time > self.STALE_TIMEOUT + and not self._disconnected.is_set() + ): + logger.warning( + "[IsaacTeleopReader] No DeviceIO data for %.1fs, flagging disconnect", + self.STALE_TIMEOUT, + ) + self._disconnected.set() + time.sleep(self._period) + continue + + controller = _build_controller_dict(raw) + if controller is not None: + with self._ctrl_lock: + self._latest_controller = controller + + body_poses = _body_data_to_24x7(raw.get("full_body")) + if body_poses is None: + if not self._unrecognised_logged and not _attr_or_item( + raw.get("full_body"), "joint_positions" + ): + self._unrecognised_logged = True + time.sleep(self._period) + continue + + stamp_ns = int(self._client.get_timestamp_ns()) + prev_stamp_ns = self._last_stamp_ns + device_dt = ((stamp_ns - prev_stamp_ns) * 1e-9) if prev_stamp_ns is not None else 0.0 + if device_dt > 0.0: + inst = 1.0 / device_dt + self._fps_ema = inst if self._fps_ema == 0.0 else (0.9 * self._fps_ema + 0.1 * inst) + self._last_stamp_ns = stamp_ns + self._last_new_data_time = time.monotonic() + if self._disconnected.is_set(): + logger.info("[IsaacTeleopReader] Fresh data received, connection restored") + self._disconnected.clear() + + sample = { + "body_poses_np": body_poses, + "timestamp_realtime": time.time(), + "timestamp_monotonic": time.monotonic(), + "timestamp_ns": stamp_ns, + "dt": device_dt, + "fps": self._fps_ema, + } + with self._lock: + self._latest = sample + + now = time.time() + if now - last_report >= 5.0: + logger.info( + "[IsaacTeleopReader] dt: %.2f ms, fps: %.2f", + device_dt * 1000.0, + self._fps_ema, + ) + last_report = now + + time.sleep(self._period) + + diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/teleop/isaac_teleop_client.py b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/isaac_teleop_client.py new file mode 100644 index 0000000000000000000000000000000000000000..7601c0b3282e5a78f738fce93603f727c48f666d --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/isaac_teleop_client.py @@ -0,0 +1,423 @@ +"""Isaac Teleop client wrapping CloudXR + DeviceIO + OpenXR for the host. + +Launches the CloudXR runtime in-process (via ``CloudXRLauncher``), opens an +OpenXR session, and starts the DeviceIO trackers (head, hands, controllers, +full-body Pico). Provides synchronous getters that the gear_sonic teleop +readers poll on a background thread. + +This replaces the legacy multi-container path (``run_cloudxr_via_docker.sh`` +plus the ROS2 ``teleop_ros2_ref`` publisher) — see +``docs/source/tutorials/isaac_teleop_publisher_setup.md`` for the in-process +setup. Requires ``isaacteleop[cloudxr]`` from ``pypi.nvidia.com`` (installed +by ``install_scripts/install_pico.sh``). +""" + +from __future__ import annotations + +import time +from contextlib import ExitStack +from pathlib import Path +from typing import Any + +import numpy as np + +import isaacteleop.deviceio as deviceio +import isaacteleop.oxr as oxr +from isaacteleop.cloudxr import CloudXRLauncher + + +def _default_pose_vec() -> np.ndarray: + return np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0], dtype=np.float64) + + +def _controller_pose_vec(controller_data: Any) -> np.ndarray: + """``[x, y, z, qx, qy, qz, qw]`` from ``ControllerSnapshot``; aim preferred, then grip.""" + if controller_data is None: + return _default_pose_vec() + controller_pose = None + aim_pose = controller_data.aim_pose + if aim_pose is not None and aim_pose.is_valid: + controller_pose = aim_pose + else: + grip_pose = controller_data.grip_pose + if grip_pose is not None and grip_pose.is_valid: + controller_pose = grip_pose + if controller_pose is None: + return _default_pose_vec() + p = controller_pose.pose.position + o = controller_pose.pose.orientation + pos = np.array([p.x, p.y, p.z], dtype=np.float64) + quat = np.array([o.x, o.y, o.z, o.w], dtype=np.float64) + return np.concatenate([pos, quat]) + + +def _pose_vec_from_head_data(head_pose: Any) -> np.ndarray: + if head_pose is None or not head_pose.is_valid: + return _default_pose_vec() + p = head_pose.pose.position + o = head_pose.pose.orientation + return np.array([p.x, p.y, p.z, o.x, o.y, o.z, o.w], dtype=np.float64) + + +def _controller_inputs(snapshot: Any) -> Any | None: + if snapshot is None: + return None + return snapshot.inputs + + +class IsaacTeleopClient: + """Single-process CloudXR + DeviceIO + OpenXR session. + + Args: + app_name: OpenXR application name (shows up in CloudXR runtime logs). + use_adb: If True, route signalling/WebRTC over a USB ADB tunnel + (``setup_oob`` + ``usb_local`` on the launcher) so the headset + reaches the host on loopback without needing shared Wi-Fi. + Requires ``adb`` and ``coturn`` on PATH. + cloudxr_install_dir: Where to install/find the CloudXR runtime. + Defaults to ``~/.cloudxr``. + cloudxr_env_config: Path to ``cloudxr.env`` (selects ``NV_DEVICE_PROFILE``). + Defaults to ``~/cloudxr.env``; created by ``install_pico.sh``. + """ + + def __init__( + self, + app_name: str = "GearSonicIsaacTeleopClient", + use_adb: bool = False, + cloudxr_install_dir: str | Path | None = None, + cloudxr_env_config: str | Path | None = None, + ) -> None: + self._app_name = app_name + self._use_adb = bool(use_adb) + self._cloudxr_install_dir = str( + cloudxr_install_dir if cloudxr_install_dir is not None else Path.home() / ".cloudxr" + ) + self._cloudxr_env_config = str( + cloudxr_env_config if cloudxr_env_config is not None else Path.home() / "cloudxr.env" + ) + + self._exit_stack: ExitStack | None = None + self._deviceio_session: Any = None + self._head_tracker: Any = None + self._hand_tracker: Any = None + self._controller_tracker: Any = None + self._body_tracker: Any = None + self._cloudxr_launcher: CloudXRLauncher | None = None + + def _clear_trackers_and_session_ref(self) -> None: + """Clear held refs after ``ExitStack.close()`` or a failed connect.""" + if self._cloudxr_launcher is not None: + try: + self._cloudxr_launcher.stop() + except Exception: + pass + self._cloudxr_launcher = None + self._deviceio_session = None + self._head_tracker = None + self._hand_tracker = None + self._controller_tracker = None + self._body_tracker = None + + def start_streaming(self) -> None: + """Launch CloudXR + open OpenXR session + start DeviceIO trackers.""" + stack = ExitStack() + try: + # OOB hub + USB-local (when use_adb=True): route signalling and + # WebRTC media over the USB cable via `adb reverse`, so the headset + # reaches the host on loopback without needing shared Wi-Fi. Needs + # `coturn` and `adb` on PATH. + self._cloudxr_launcher = CloudXRLauncher( + install_dir=self._cloudxr_install_dir, + env_config=self._cloudxr_env_config, + accept_eula=True, + setup_oob=self._use_adb, + usb_local=self._use_adb, + ) + + self._head_tracker = deviceio.HeadTracker() + self._hand_tracker = deviceio.HandTracker() + self._controller_tracker = deviceio.ControllerTracker() + self._body_tracker = deviceio.FullBodyTrackerPico() + trackers = [ + self._head_tracker, + self._hand_tracker, + self._controller_tracker, + self._body_tracker, + ] + required_extensions = deviceio.DeviceIOSession.get_required_extensions(trackers) + oxr_session = stack.enter_context( + oxr.OpenXRSession(self._app_name, required_extensions) + ) + handles = oxr_session.get_handles() + self._deviceio_session = stack.enter_context( + deviceio.DeviceIOSession.run(trackers, handles) + ) + self._exit_stack = stack + print("Isaac Teleop session initialized.") + + except RuntimeError as e: + stack.close() + self._exit_stack = None + self._clear_trackers_and_session_ref() + if "Failed to get OpenXR system" in str(e) or "OpenXR" in str(e): + print(f"IsaacTeleopClient: no XR session yet ({e}).") + else: + raise + except Exception as e: + stack.close() + self._exit_stack = None + self._clear_trackers_and_session_ref() + print(f"IsaacTeleopClient: failed to start sessions ({e}).") + + def _get_tracker_data(self) -> dict[str, Any] | None: + """Poll current tracking data and return it as a dictionary. + + Returns: + Dict with keys ``left_controller``, ``right_controller``, ``head``, + ``left_hand``, ``right_hand``, ``full_body``. Each value is the + corresponding tracker's ``.data`` payload (raw DeviceIO type). + """ + if self._deviceio_session is None: + return None + try: + self._deviceio_session.update() + except RuntimeError as e: + print(f"IsaacTeleopClient: DeviceIO update failed ({e}); closing.") + self.close() + return None + + session = self._deviceio_session + return { + "left_controller": self._controller_tracker.get_left_controller(session).data, + "right_controller": self._controller_tracker.get_right_controller(session).data, + "head": self._head_tracker.get_head(session).data, + "left_hand": self._hand_tracker.get_left_hand(session).data, + "right_hand": self._hand_tracker.get_right_hand(session).data, + "full_body": self._body_tracker.get_body_pose(session).data, + } + + def get_pose_by_name(self, name: str) -> np.ndarray: + """Return ``[x, y, z, qx, qy, qz, qw]`` for ``name`` ∈ {left_controller, right_controller, headset}.""" + raw = self._get_tracker_data() + if raw is None: + return _default_pose_vec() + + if name == "left_controller": + return _controller_pose_vec(raw.get("left_controller")) + if name == "right_controller": + return _controller_pose_vec(raw.get("right_controller")) + if name == "headset": + return _pose_vec_from_head_data(raw.get("head")) + raise ValueError( + f"Invalid name: {name}. Valid names: 'left_controller', 'right_controller', 'headset'." + ) + + def _snapshot_side(self, raw: dict[str, Any] | None, side: str) -> Any: + if raw is None: + return None + key = "left_controller" if side == "left" else "right_controller" + return raw.get(key) + + def get_key_value_by_name(self, name: str) -> float: + """Return trigger/grip value for ``name`` ∈ {left,right}_{trigger,grip}.""" + raw = self._get_tracker_data() + if name == "left_trigger": + inp = _controller_inputs(self._snapshot_side(raw, "left")) + return float(inp.trigger_value) if inp is not None else 0.0 + if name == "right_trigger": + inp = _controller_inputs(self._snapshot_side(raw, "right")) + return float(inp.trigger_value) if inp is not None else 0.0 + if name == "left_grip": + inp = _controller_inputs(self._snapshot_side(raw, "left")) + return float(inp.squeeze_value) if inp is not None else 0.0 + if name == "right_grip": + inp = _controller_inputs(self._snapshot_side(raw, "right")) + return float(inp.squeeze_value) if inp is not None else 0.0 + raise ValueError( + f"Invalid name: {name}. Valid names: " + "'left_trigger', 'right_trigger', 'left_grip', 'right_grip'." + ) + + def get_button_state_by_name(self, name: str) -> bool: + """Return True/False for face buttons and stick clicks. + + Valid names: ``A``, ``B``, ``X``, ``Y``, + ``left_menu_button``, ``right_menu_button``, + ``left_axis_click``, ``right_axis_click``. + """ + raw = self._get_tracker_data() + left = _controller_inputs(self._snapshot_side(raw, "left")) + right = _controller_inputs(self._snapshot_side(raw, "right")) + + if name == "A": + return right is not None and float(right.primary_click) > 0.5 + if name == "B": + return right is not None and float(right.secondary_click) > 0.5 + if name == "X": + return left is not None and float(left.primary_click) > 0.5 + if name == "Y": + return left is not None and float(left.secondary_click) > 0.5 + if name in ("left_menu_button", "right_menu_button"): + # Pico-specific menu button is not exposed via standard OpenXR + # input bindings — DeviceIO doesn't surface it for the supported + # controllers, so report not-pressed. + return False + if name == "left_axis_click": + return left is not None and float(left.thumbstick_click) > 0.5 + if name == "right_axis_click": + return right is not None and float(right.thumbstick_click) > 0.5 + raise ValueError( + f"Invalid name: {name}. Valid names: 'A', 'B', 'X', 'Y', " + "'left_menu_button', 'right_menu_button', 'left_axis_click', 'right_axis_click'." + ) + + def get_joystick_state(self, controller: str) -> list[float]: + """Return ``[x, y]`` joystick state for ``controller`` ∈ {left, right}.""" + side = controller.lower() + if side not in ("left", "right"): + raise ValueError( + f"Invalid controller: {controller}. Valid controllers: 'left', 'right'." + ) + + raw = self._get_tracker_data() + inp = _controller_inputs(self._snapshot_side(raw, side)) + if inp is None: + return [0.0, 0.0] + return [float(inp.thumbstick_x), float(inp.thumbstick_y)] + + def get_full_body_data(self) -> Any | None: + """Return the raw DeviceIO ``FullBodyTrackerPico`` data payload (or None). + + Body-joint extraction (24×7 pose array) lives in + ``input_readers.IsaacTeleopReader``, which knows the gear_sonic schema. + """ + raw = self._get_tracker_data() + if raw is None: + return None + return raw.get("full_body") + + def get_timestamp_ns(self) -> int: + """Return the host monotonic timestamp in nanoseconds.""" + return int(time.monotonic_ns()) + + def close(self) -> None: + """Close OpenXR session, stop DeviceIO + CloudXR runtime.""" + if self._exit_stack is not None: + try: + self._exit_stack.close() + except Exception: + pass + self._exit_stack = None + self._clear_trackers_and_session_ref() + + +def main() -> None: + """Poll ``IsaacTeleopClient`` getters periodically and print (Ctrl+C to stop). + + Usage:: + + source .venv_teleop/bin/activate + + # Setup verification: bring CloudXR up, populate ~/.cloudxr/, exit clean + python -m gear_sonic.utils.teleop.isaac_teleop_client --init-only + + # Live print (default): poll getters and print until Ctrl+C + python -m gear_sonic.utils.teleop.isaac_teleop_client --hz 5 + """ + import argparse + import sys + from pathlib import Path + + parser = argparse.ArgumentParser( + description="Print IsaacTeleopClient getter outputs at a fixed rate." + ) + parser.add_argument("--hz", type=float, default=5.0, help="Print rate in Hz (default: 5)") + parser.add_argument( + "--use-adb", action="store_true", help="Route CloudXR over USB ADB (OOB / usb-local)." + ) + parser.add_argument( + "--init-only", + action="store_true", + help=( + "Bring CloudXR up to populate ~/.cloudxr/ ownership + env file, " + "then exit. Use as a setup verification step." + ), + ) + args = parser.parse_args() + + client = IsaacTeleopClient(use_adb=args.use_adb) + client.start_streaming() + + if args.init_only: + # Confirm the runtime actually populated the install dir before we tear down. + run_env = Path.home() / ".cloudxr" / "run" / "cloudxr.env" + if run_env.exists(): + print(f"[OK] CloudXR runtime initialized; {run_env} written.") + client.close() + sys.exit(0) + print( + f"[ERROR] CloudXR runtime did not populate {run_env} — check the " + "earlier log for IsaacTeleopClient errors.", + file=sys.stderr, + ) + client.close() + sys.exit(1) + + period = 1.0 / max(0.1, float(args.hz)) + pose_names = ("left_controller", "right_controller", "headset") + key_names = ("left_trigger", "right_trigger", "left_grip", "right_grip") + button_names = ( + "A", + "B", + "X", + "Y", + "left_menu_button", + "right_menu_button", + "left_axis_click", + "right_axis_click", + ) + joy_sides = ("left", "right") + + try: + while True: + t0 = time.time() + print("=" * 72, flush=True) + print( + f"time={time.strftime('%H:%M:%S')} get_timestamp_ns={client.get_timestamp_ns()}", + flush=True, + ) + + for name in pose_names: + v = client.get_pose_by_name(name) + print( + f" get_pose_by_name({name!r}): " + f"{np.array2string(v, precision=4, suppress_small=True)}", + flush=True, + ) + + for name in key_names: + print( + f" get_key_value_by_name({name!r}): {client.get_key_value_by_name(name):.4f}", + flush=True, + ) + + for name in button_names: + print( + f" get_button_state_by_name({name!r}): {client.get_button_state_by_name(name)}", + flush=True, + ) + + for side in joy_sides: + j = client.get_joystick_state(side) + print(f" get_joystick_state({side!r}): {j}", flush=True) + + dt = time.time() - t0 + time.sleep(max(0.0, period - dt)) + except KeyboardInterrupt: + print("\nStopped.", flush=True) + finally: + client.close() + + +if __name__ == "__main__": + main() diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/torch_utils.py b/GR00T-WholeBodyControl/gear_sonic/utils/torch_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0e7d54221f045c352e92e90784dcecaf7fd96392 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/torch_utils.py @@ -0,0 +1,554 @@ +"""PyTorch and USD/Gf utility functions for IsaacGym-based RL. + +Provides quaternion arithmetic (multiply, apply, rotate, conjugate, inverse +transform, combine), Euler-angle conversions, tensor helpers (clamp, scale, +unscale, random float/direction), and a USD ``Gf.Matrix4d`` construction +helper. Most functions are compiled with ``@torch.jit.script`` for +performance. + +SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 +""" + +import numpy as np +from pxr import Gf +import torch + + +def set_env_attr(self, attr_name, attr_val, env_ids): + """Set a per-environment attribute on an env object for the given env ids. + + If the attribute already exists it is indexed by ``env_ids``; otherwise + the value is set as a plain attribute (useful for first-time initialisation + before the buffer exists). + + Args: + self: Environment object that owns the attribute. + attr_name: Name of the attribute to set. + attr_val: Value(s) to assign. + env_ids: Integer indices of the environments to update. + """ + if hasattr(self, attr_name): + getattr(self, attr_name)[env_ids] = attr_val + else: + setattr(self, attr_name, attr_val) + + +def to_torch( + x: np.ndarray | torch.Tensor | list, device: torch.device | str, dtype=None, requires_grad=False +): + """Convert a list, NumPy array, or Tensor to a ``torch.Tensor`` on ``device``. + + Args: + x: Input data. Lists are first converted to ``np.ndarray``. + device: Target device (e.g. ``"cuda:0"`` or ``torch.device("cpu")``). + dtype: Desired dtype. Defaults to ``torch.float`` for non-Tensor + inputs; for existing Tensors the current dtype is preserved when + ``dtype`` is ``None``. + requires_grad: Whether the result should track gradients. + + Returns: + A ``torch.Tensor`` on ``device`` with the requested dtype and + ``requires_grad`` setting. + """ + # Convert list to np.ndarray for shape and dtype handling + if isinstance(x, list): + x = np.array(x) + + if not isinstance(x, torch.Tensor): + if dtype is None: + dtype = torch.float + x = torch.tensor(x, device=device, dtype=dtype, requires_grad=requires_grad) + else: + # torch.Tensor + if dtype is None: + x = x.to(device=device) + if x.requires_grad != requires_grad: + x = x.detach().requires_grad_(requires_grad) + else: + x = x.to(dtype=dtype, device=device) + if x.requires_grad != requires_grad: + x = x.detach().requires_grad_(requires_grad) + return x + + +@torch.jit.script +def quat_mul(a, b): + """Multiply two batches of quaternions (xyzw convention). + + Args: + a: Quaternion tensor of shape ``(..., 4)`` in xyzw order. + b: Quaternion tensor of the same shape as ``a``. + + Returns: + Product quaternion tensor of the same shape as ``a``. + """ + assert a.shape == b.shape + shape = a.shape + a = a.reshape(-1, 4) + b = b.reshape(-1, 4) + + x1, y1, z1, w1 = a[:, 0], a[:, 1], a[:, 2], a[:, 3] + x2, y2, z2, w2 = b[:, 0], b[:, 1], b[:, 2], b[:, 3] + ww = (z1 + x1) * (x2 + y2) + yy = (w1 - y1) * (w2 + z2) + zz = (w1 + y1) * (w2 - z2) + xx = ww + yy + zz + qq = 0.5 * (xx + (z1 - x1) * (x2 - y2)) + w = qq - ww + (z1 - y1) * (y2 - z2) + x = qq - xx + (x1 + w1) * (x2 + w2) + y = qq - yy + (w1 - x1) * (y2 + z2) + z = qq - zz + (z1 + y1) * (w2 - x2) + + quat = torch.stack([x, y, z, w], dim=-1).view(shape) + + return quat + + +@torch.jit.script +def normalize(x, eps: float = 1e-9): + """L2-normalize a tensor along its last dimension. + + Args: + x: Input tensor of any shape. + eps: Minimum norm value to clamp to, preventing division by zero. + + Returns: + Tensor of the same shape as ``x`` with unit L2 norm along the last dim. + """ + return x / x.norm(p=2, dim=-1).clamp(min=eps, max=None).unsqueeze(-1) + + +@torch.jit.script +def quat_apply(a, b): + """Rotate a 3-D vector by a quaternion (xyzw convention). + + Args: + a: Quaternion tensor of shape ``(..., 4)`` in xyzw order. + b: Vector tensor of shape ``(..., 3)``. + + Returns: + Rotated vector tensor of the same shape as ``b``. + """ + shape = b.shape + a = a.reshape(-1, 4) + b = b.reshape(-1, 3) + xyz = a[:, :3] + t = xyz.cross(b, dim=-1) * 2 + return (b + a[:, 3:] * t + xyz.cross(t, dim=-1)).view(shape) + + +@torch.jit.script +def quat_rotate(q, v): + """Rotate batched 3-D vectors by batched quaternions (xyzw convention). + + Uses the expanded form ``v' = 2(q_w^2 - 0.5)v + 2(q·v)q + 2q_w(q×v)``. + + Args: + q: Quaternion tensor of shape ``(N, 4)`` in xyzw order. + v: Vector tensor of shape ``(N, 3)``. + + Returns: + Rotated vector tensor of shape ``(N, 3)``. + """ + shape = q.shape + q_w = q[:, -1] + q_vec = q[:, :3] + a = v * (2.0 * q_w**2 - 1.0).unsqueeze(-1) + b = torch.cross(q_vec, v, dim=-1) * q_w.unsqueeze(-1) * 2.0 + c = q_vec * torch.bmm(q_vec.view(shape[0], 1, 3), v.view(shape[0], 3, 1)).squeeze(-1) * 2.0 + return a + b + c + + +# @torch.jit.script +def quat_rotate_inverse(q, v): + """Rotate batched 3-D vectors by the *inverse* of batched quaternions. + + Equivalent to rotating by the conjugate quaternion (i.e. the transpose of + the rotation matrix). + + Args: + q: Quaternion tensor of shape ``(N, 4)`` in xyzw order. + v: Vector tensor of shape ``(N, 3)``. + + Returns: + Inversely-rotated vector tensor of shape ``(N, 3)``. + """ + shape = q.shape + q_w = q[:, -1] + q_vec = q[:, :3] + a = v * (2.0 * q_w**2 - 1.0).unsqueeze(-1) + b = torch.cross(q_vec, v, dim=-1) * q_w.unsqueeze(-1) * 2.0 + c = q_vec * torch.bmm(q_vec.view(shape[0], 1, 3), v.view(shape[0], 3, 1)).squeeze(-1) * 2.0 + return a - b + c + + +@torch.jit.script +def quat_conjugate(a): + """Return the conjugate of a batch of quaternions (xyzw convention). + + The conjugate negates the imaginary (xyz) part while keeping the real (w) + part, yielding the inverse rotation for unit quaternions. + + Args: + a: Quaternion tensor of shape ``(..., 4)`` in xyzw order. + + Returns: + Conjugate quaternion tensor of the same shape. + """ + shape = a.shape + a = a.reshape(-1, 4) + return torch.cat((-a[:, :3], a[:, -1:]), dim=-1).view(shape) + + +@torch.jit.script +def quat_unit(a): + """Normalize a batch of quaternions to unit length. + + Args: + a: Quaternion tensor of shape ``(..., 4)``. + + Returns: + Unit-length quaternion tensor of the same shape. + """ + return normalize(a) + + +@torch.jit.script +def quat_from_angle_axis(angle, axis): + """Construct a unit quaternion from an angle-axis representation. + + Args: + angle: Rotation angle in radians, shape ``(N,)``. + axis: Rotation axes, shape ``(N, 3)``. Need not be unit vectors. + + Returns: + Unit quaternion tensor of shape ``(N, 4)`` in xyzw order. + """ + theta = (angle / 2).unsqueeze(-1) + xyz = normalize(axis) * theta.sin() + w = theta.cos() + return quat_unit(torch.cat([xyz, w], dim=-1)) + + +@torch.jit.script +def normalize_angle(x): + """Wrap angles into the range ``(-pi, pi]``. + + Args: + x: Angle tensor (radians), any shape. + + Returns: + Wrapped angle tensor of the same shape. + """ + return torch.atan2(torch.sin(x), torch.cos(x)) + + +@torch.jit.script +def tf_inverse(q, t): + """Compute the inverse of a rigid transform (q, t). + + Args: + q: Rotation quaternion, shape ``(N, 4)`` in xyzw order. + t: Translation vector, shape ``(N, 3)``. + + Returns: + Tuple ``(q_inv, t_inv)`` representing the inverse transform. + """ + q_inv = quat_conjugate(q) + return q_inv, -quat_apply(q_inv, t) + + +@torch.jit.script +def tf_apply(q, t, v): + """Apply a rigid transform (q, t) to a batch of points v. + + Computes ``R(q) * v + t``. + + Args: + q: Rotation quaternion, shape ``(N, 4)`` in xyzw order. + t: Translation vector, shape ``(N, 3)``. + v: Points to transform, shape ``(N, 3)``. + + Returns: + Transformed points of shape ``(N, 3)``. + """ + return quat_apply(q, v) + t + + +@torch.jit.script +def tf_vector(q, v): + """Rotate a vector by a quaternion (no translation). + + Args: + q: Quaternion, shape ``(..., 4)`` in xyzw order. + v: Vector to rotate, shape ``(..., 3)``. + + Returns: + Rotated vector of the same shape as ``v``. + """ + return quat_apply(q, v) + + +@torch.jit.script +def tf_combine(q1, t1, q2, t2): + """Compose two rigid transforms T1 followed by T2. + + Computes the combined rotation ``q1 * q2`` and the combined translation + ``R(q1) * t2 + t1``. + + Args: + q1: First rotation quaternion, shape ``(N, 4)`` in xyzw order. + t1: First translation, shape ``(N, 3)``. + q2: Second rotation quaternion, shape ``(N, 4)`` in xyzw order. + t2: Second translation, shape ``(N, 3)``. + + Returns: + Tuple ``(q_combined, t_combined)`` of the composed transform. + """ + return quat_mul(q1, q2), quat_apply(q1, t2) + t1 + + +@torch.jit.script +def get_basis_vector(q, v): + """Rotate a basis vector ``v`` by quaternion ``q``. + + Args: + q: Quaternion tensor, shape ``(N, 4)`` in xyzw order. + v: Basis vector, shape ``(N, 3)``. + + Returns: + Rotated vector of shape ``(N, 3)``. + """ + return quat_rotate(q, v) + + +def get_axis_params(value, axis_idx, x_value=0.0, dtype=np.float64, n_dims=3): + """Construct a parameter list for a USD ``Vec`` along a specific axis. + + Creates an n-dimensional vector that is ``value`` along ``axis_idx`` and + zero everywhere else, then overrides index 0 with ``x_value``. + + Args: + value: Scalar value to place at position ``axis_idx``. + axis_idx: Index of the axis to set to ``value``. + x_value: Value to assign to index 0 after the axis fill. + dtype: NumPy dtype of the output array. + n_dims: Total number of dimensions in the vector. + + Returns: + List of ``n_dims`` floats suitable for passing to a USD ``Vec`` + constructor. + """ + zs = np.zeros((n_dims,)) + assert axis_idx < n_dims, "the axis dim should be within the vector dimensions" + zs[axis_idx] = 1.0 + params = np.where(zs == 1.0, value, zs) + params[0] = x_value + return list(params.astype(dtype)) + + +@torch.jit.script +def copysign(a, b): + # type: (float, Tensor) -> Tensor + """Copy the sign of tensor ``b`` onto scalar ``a``. + + Returns a tensor of the same shape as ``b`` with magnitude ``|a|`` and + sign matching each element of ``b``. + + Args: + a: Scalar magnitude. + b: Tensor whose signs are copied, shape ``(N,)``. + + Returns: + Tensor of shape ``(N,)`` equal to ``|a| * sign(b)``. + """ + a = torch.tensor(a, device=b.device, dtype=torch.float).repeat(b.shape[0]) + return torch.abs(a) * torch.sign(b) + + +@torch.jit.script +def get_euler_xyz(q): + """Extract Euler XYZ angles (roll, pitch, yaw) from a batch of quaternions. + + Uses the standard ZYX intrinsic decomposition. Handles the gimbal-lock + singularity at |sinp| >= 1 via ``copysign``. + + Args: + q: Quaternion tensor of shape ``(N, 4)`` in xyzw order. + + Returns: + Tuple ``(roll, pitch, yaw)`` each of shape ``(N,)`` in radians, + mapped to ``[0, 2*pi)``. + """ + qx, qy, qz, qw = 0, 1, 2, 3 + # roll (x-axis rotation) + sinr_cosp = 2.0 * (q[:, qw] * q[:, qx] + q[:, qy] * q[:, qz]) + cosr_cosp = ( + q[:, qw] * q[:, qw] - q[:, qx] * q[:, qx] - q[:, qy] * q[:, qy] + q[:, qz] * q[:, qz] + ) + roll = torch.atan2(sinr_cosp, cosr_cosp) + + # pitch (y-axis rotation) + sinp = 2.0 * (q[:, qw] * q[:, qy] - q[:, qz] * q[:, qx]) + pitch = torch.where(torch.abs(sinp) >= 1, copysign(np.pi / 2.0, sinp), torch.asin(sinp)) + + # yaw (z-axis rotation) + siny_cosp = 2.0 * (q[:, qw] * q[:, qz] + q[:, qx] * q[:, qy]) + cosy_cosp = ( + q[:, qw] * q[:, qw] + q[:, qx] * q[:, qx] - q[:, qy] * q[:, qy] - q[:, qz] * q[:, qz] + ) + yaw = torch.atan2(siny_cosp, cosy_cosp) + + return roll % (2 * np.pi), pitch % (2 * np.pi), yaw % (2 * np.pi) + + +@torch.jit.script +def quat_from_euler_xyz(roll, pitch, yaw): + """Construct unit quaternions from intrinsic XYZ Euler angles. + + Args: + roll: Rotation around X axis (radians), shape ``(N,)``. + pitch: Rotation around Y axis (radians), shape ``(N,)``. + yaw: Rotation around Z axis (radians), shape ``(N,)``. + + Returns: + Unit quaternion tensor of shape ``(N, 4)`` in xyzw order. + """ + cy = torch.cos(yaw * 0.5) + sy = torch.sin(yaw * 0.5) + cr = torch.cos(roll * 0.5) + sr = torch.sin(roll * 0.5) + cp = torch.cos(pitch * 0.5) + sp = torch.sin(pitch * 0.5) + + qw = cy * cr * cp + sy * sr * sp + qx = cy * sr * cp - sy * cr * sp + qy = cy * cr * sp + sy * sr * cp + qz = sy * cr * cp - cy * sr * sp + + return torch.stack([qx, qy, qz, qw], dim=-1) + + +@torch.jit.script +def torch_rand_float(lower, upper, shape, device): + # type: (float, float, Tuple[int, int], str) -> Tensor + """Sample uniform random floats in ``[lower, upper)``. + + Args: + lower: Lower bound of the uniform distribution. + upper: Upper bound of the uniform distribution. + shape: Output shape as a 2-tuple ``(rows, cols)``. + device: Target device string (e.g. ``"cuda:0"``). + + Returns: + Float tensor of the given shape on ``device``. + """ + return (upper - lower) * torch.rand(*shape, device=device) + lower + + +@torch.jit.script +def torch_random_dir_2(shape, device): + # type: (Tuple[int, int], str) -> Tensor + """Sample uniformly random unit vectors in 2-D. + + Args: + shape: Shape of the angle samples as a 2-tuple ``(rows, 1)``. + device: Target device string. + + Returns: + Float tensor of shape ``(rows, 2)`` containing ``(cos θ, sin θ)`` + with ``θ`` drawn uniformly from ``[-π, π)``. + """ + angle = torch_rand_float(-np.pi, np.pi, shape, device).squeeze(-1) + return torch.stack([torch.cos(angle), torch.sin(angle)], dim=-1) + + +@torch.jit.script +def tensor_clamp(t, min_t, max_t): + """Element-wise clamp of tensor ``t`` to the range ``[min_t, max_t]``. + + Unlike ``torch.clamp`` this version accepts tensors for the bounds so that + per-element limits are supported. + + Args: + t: Input tensor. + min_t: Lower bound tensor, same shape or broadcastable to ``t``. + max_t: Upper bound tensor, same shape or broadcastable to ``t``. + + Returns: + Clamped tensor of the same shape as ``t``. + """ + return torch.max(torch.min(t, max_t), min_t) + + +@torch.jit.script +def scale(x, lower, upper): + """Map values from ``[-1, 1]`` to ``[lower, upper]``. + + Args: + x: Input tensor in the normalised range ``[-1, 1]``. + lower: Target range lower bound. + upper: Target range upper bound. + + Returns: + Tensor scaled to ``[lower, upper]``. + """ + return 0.5 * (x + 1.0) * (upper - lower) + lower + + +@torch.jit.script +def unscale(x, lower, upper): + """Map values from ``[lower, upper]`` to ``[-1, 1]``. + + Inverse of :func:`scale`. + + Args: + x: Input tensor in the range ``[lower, upper]``. + lower: Source range lower bound. + upper: Source range upper bound. + + Returns: + Tensor normalised to ``[-1, 1]``. + """ + return (2.0 * x - upper - lower) / (upper - lower) + + +def unscale_np(x, lower, upper): + """NumPy equivalent of :func:`unscale`. + + Args: + x: Input array in the range ``[lower, upper]``. + lower: Source range lower bound. + upper: Source range upper bound. + + Returns: + Array normalised to ``[-1, 1]``. + """ + return (2.0 * x - upper - lower) / (upper - lower) + + +def euler_xyz_to_gf_matrix(angles): + """Convert Euler XYZ angles (in radians) to a USD Gf.Matrix4d. + + Args: + angles: [roll, pitch, yaw] in radians. Can be list, numpy array, or torch.Tensor. + + Returns: + Gf.Matrix4d with the rotation applied. + """ + if isinstance(angles, list | np.ndarray): + angles = torch.tensor(angles, dtype=torch.float32) + if angles.dim() == 1: + angles = angles.unsqueeze(0) # [1, 3] + + roll, pitch, yaw = angles[:, 0], angles[:, 1], angles[:, 2] + quat = quat_from_euler_xyz(roll, pitch, yaw) # [1, 4] as [qx, qy, qz, qw] + + # Convert to Gf.Quatd (w, x, y, z order) + qx, qy, qz, qw = quat[0].tolist() + gf_quat = Gf.Quatd(qw, qx, qy, qz) + + m = Gf.Matrix4d() + m.SetRotate(Gf.Rotation(gf_quat)) + return m