Add files using upload-large-folder tool

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."""

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)

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

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,
+    )

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()

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()

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()

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())

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())

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

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)