diff --git a/GR00T-WholeBodyControl/gear_sonic/camera/__pycache__/__init__.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/camera/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..82f9bf5775e8f8ef2693da5bec87725bf720cc59
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/camera/__pycache__/__init__.cpython-310.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/camera/__pycache__/composed_camera.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/camera/__pycache__/composed_camera.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bef949b07db5654b056d0844dfe8714e57ea0538
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/camera/__pycache__/composed_camera.cpython-310.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/camera/__pycache__/sensor.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/camera/__pycache__/sensor.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a40fc52e0080645c8074b8f5c8bada1fa9fafa3d
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/camera/__pycache__/sensor.cpython-310.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/camera/__pycache__/sensor_server.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/camera/__pycache__/sensor_server.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..45fbfff98c646cb386642af3b451bbcebbc959df
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/camera/__pycache__/sensor_server.cpython-310.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/camera/drivers/__init__.py b/GR00T-WholeBodyControl/gear_sonic/camera/drivers/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/GR00T-WholeBodyControl/gear_sonic/camera/drivers/dummy.py b/GR00T-WholeBodyControl/gear_sonic/camera/drivers/dummy.py
new file mode 100644
index 0000000000000000000000000000000000000000..19ac0cacde88ebba6c0c5a95590043d7ee5eb528
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/camera/drivers/dummy.py
@@ -0,0 +1,86 @@
+"""Dummy / replay sensor for testing without real camera hardware.
+
+``DummySensor`` generates random images.
+``ReplayDummySensor`` loops frames from a video file.
+"""
+
+import time
+from typing import Any
+
+import cv2
+import numpy as np
+
+try:
+    import gymnasium as gym
+except ImportError:
+    gym = None  # type: ignore[assignment]
+
+from gear_sonic.camera.sensor import Sensor
+from gear_sonic.camera.sensor_server import ImageMessageSchema
+
+
+class DummySensor(Sensor):
+    """Produces random 640x480 images at each read() call."""
+
+    def __init__(self):
+        pass
+
+    def read(self) -> dict[str, Any] | None:
+        return {
+            "color_image": np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8),
+            "timestamp": time.time(),
+        }
+
+    def serialize(self, data: dict[str, Any]) -> dict[str, Any]:
+        raise NotImplementedError("DummySensor does not support serialize()")
+
+    def close(self):
+        pass
+
+    def observation_space(self):
+        if gym is None:
+            return None
+        return gym.spaces.Dict(
+            {
+                "color_image": gym.spaces.Box(
+                    low=0, high=255, shape=(480, 640, 3), dtype=np.uint8
+                ),
+            }
+        )
+
+
+class ReplayDummySensor(DummySensor):
+    """Loops frames from a video file, useful for offline testing."""
+
+    def __init__(self, video_path: str):
+        self.video_path = video_path
+        self.image_ctr = 0
+        self.video_reader = cv2.VideoCapture(video_path)
+        self.frames = []
+        while self.video_reader.isOpened():
+            ret, frame = self.video_reader.read()
+            if not ret:
+                break
+            self.frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+
+    def read(self) -> dict[str, Any] | None:
+        self.image_ctr += 1
+        if self.image_ctr >= len(self.frames):
+            self.image_ctr = 0
+
+        img = self.frames[self.image_ctr]
+        img = cv2.resize(img, (640, 480))
+        return {
+            "color_image": img,
+            "timestamp": {"color_image": time.time()},
+        }
+
+    def serialize(self, data: dict[str, Any]) -> dict[str, Any]:
+        serialized_msg = ImageMessageSchema(
+            timestamps=data["timestamp"] if isinstance(data["timestamp"], dict) else {"color_image": data["timestamp"]},
+            images={"color_image": data["color_image"]},
+        )
+        return serialized_msg.serialize()
+
+    def close(self):
+        self.video_reader.release()
diff --git a/GR00T-WholeBodyControl/gear_sonic/camera/drivers/oak.py b/GR00T-WholeBodyControl/gear_sonic/camera/drivers/oak.py
new file mode 100644
index 0000000000000000000000000000000000000000..abba4eefcf455dad8c4c271f532bb712639dfdc2
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/camera/drivers/oak.py
@@ -0,0 +1,409 @@
+"""OAK (DepthAI) camera driver.
+
+Requires the ``depthai`` SDK — install with::
+
+    pip install depthai
+
+See https://docs.luxonis.com/ for hardware-specific instructions.
+"""
+
+import time
+from typing import Any
+
+import cv2
+import numpy as np
+
+try:
+    import gymnasium as gym
+except ImportError:
+    gym = None  # type: ignore[assignment]
+
+import depthai as dai
+
+from gear_sonic.camera.sensor import Sensor
+from gear_sonic.camera.sensor_server import (
+    CameraMountPosition,
+    ImageMessageSchema,
+    SensorServer,
+)
+
+
+class OAKConfig:
+    """Configuration for the OAK camera."""
+
+    color_image_dim: tuple[int, int] = (640, 480)
+    monochrome_image_dim: tuple[int, int] = (640, 480)
+    fps: int = 30
+    enable_color: bool = True
+    enable_mono_cameras: bool = False
+    mount_position: str = CameraMountPosition.EGO_VIEW.value
+    autofocus: bool = False
+    manual_focus: int = 130
+    use_mjpeg: bool = False
+    mjpeg_quality: int = 80
+
+
+class OAKSensor(Sensor, SensorServer):
+    """Sensor for the OAK camera family (OAK-D, OAK-1, etc.)."""
+
+    def __init__(
+        self,
+        run_as_server: bool = False,
+        port: int = 5555,
+        config: OAKConfig = OAKConfig(),
+        device_id: str | None = None,
+        mount_position: str = CameraMountPosition.EGO_VIEW.value,
+    ):
+        self.config = config
+        self.mount_position = mount_position
+        self._run_as_server = run_as_server
+
+        device_infos = dai.Device.getAllAvailableDevices()
+        assert len(device_infos) > 0, f"No OAK devices found for {mount_position}"
+        print(f"Device infos: {device_infos}")
+        if device_id is not None:
+            device_found = False
+            for device_info in device_infos:
+                if device_info.getDeviceId() == device_id:
+                    self.device = dai.Device(device_info, maxUsbSpeed=dai.UsbSpeed.SUPER_PLUS)
+                    device_found = True
+                    break
+            if not device_found:
+                raise ValueError(f"Device with ID {device_id} not found")
+        else:
+            self.device = dai.Device()
+
+        print(f"Connected to OAK device: {self.device.getDeviceName(), self.device.getDeviceId()}")
+        print(f"Device ID: {self.device.getDeviceId()}")
+
+        sockets: list[dai.CameraBoardSocket] = self.device.getConnectedCameras()
+        print(f"Available cameras: {[str(s) for s in sockets]}")
+
+        self.pipeline = dai.Pipeline(self.device)
+        self.output_queues = {}
+        self._use_mjpeg = config.use_mjpeg
+
+        # RGB camera (CAM_A)
+        if config.enable_color and dai.CameraBoardSocket.CAM_A in sockets:
+            self.cam_rgb = self.pipeline.create(dai.node.Camera)
+            cam_socket = dai.CameraBoardSocket.CAM_A
+            self.cam_rgb = self.cam_rgb.build(cam_socket)
+
+            if config.use_mjpeg:
+                cam_out = self.cam_rgb.requestOutput(
+                    config.color_image_dim,
+                    dai.ImgFrame.Type.NV12,
+                    fps=config.fps,
+                )
+                encoder = self.pipeline.create(dai.node.VideoEncoder)
+                encoder.setDefaultProfilePreset(
+                    config.fps, dai.VideoEncoderProperties.Profile.MJPEG
+                )
+                encoder.setQuality(config.mjpeg_quality)
+                cam_out.link(encoder.input)
+                self.output_queues["color"] = encoder.out.createOutputQueue(
+                    maxSize=3, blocking=False
+                )
+            else:
+                self.output_queues["color"] = self.cam_rgb.requestOutput(
+                    config.color_image_dim,
+                    fps=config.fps,
+                ).createOutputQueue(maxSize=3, blocking=False)
+            print(f"Enabled CAM_A (RGB){' with MJPEG encoding' if config.use_mjpeg else ''}")
+
+            if not config.autofocus:
+                ctrl_in = self.cam_rgb.inputControl.createInputQueue()
+                ctrl = dai.CameraControl()
+                ctrl.setAutoFocusMode(dai.CameraControl.AutoFocusMode.OFF)
+                ctrl.setManualFocus(config.manual_focus)
+                ctrl_in.send(ctrl)
+                print(f"Autofocus disabled, manual focus set to {config.manual_focus}")
+
+        # Monochrome cameras (CAM_B / CAM_C)
+        if config.enable_mono_cameras:
+            if dai.CameraBoardSocket.CAM_B in sockets:
+                self.cam_mono_left = self.pipeline.create(dai.node.Camera)
+                cam_socket = dai.CameraBoardSocket.CAM_B
+                self.cam_mono_left = self.cam_mono_left.build(cam_socket)
+
+                if config.use_mjpeg:
+                    cam_out = self.cam_mono_left.requestOutput(
+                        config.monochrome_image_dim,
+                        dai.ImgFrame.Type.NV12,
+                        fps=config.fps,
+                    )
+                    encoder = self.pipeline.create(dai.node.VideoEncoder)
+                    encoder.setDefaultProfilePreset(
+                        config.fps, dai.VideoEncoderProperties.Profile.MJPEG
+                    )
+                    encoder.setQuality(config.mjpeg_quality)
+                    cam_out.link(encoder.input)
+                    self.output_queues["mono_left"] = encoder.out.createOutputQueue(
+                        maxSize=3, blocking=False
+                    )
+                else:
+                    self.output_queues["mono_left"] = self.cam_mono_left.requestOutput(
+                        config.monochrome_image_dim,
+                        fps=config.fps,
+                    ).createOutputQueue(maxSize=3, blocking=False)
+                print("Enabled CAM_B (Monochrome Left)")
+
+            if dai.CameraBoardSocket.CAM_C in sockets:
+                self.cam_mono_right = self.pipeline.create(dai.node.Camera)
+                cam_socket = dai.CameraBoardSocket.CAM_C
+                self.cam_mono_right = self.cam_mono_right.build(cam_socket)
+
+                if config.use_mjpeg:
+                    cam_out = self.cam_mono_right.requestOutput(
+                        config.monochrome_image_dim,
+                        dai.ImgFrame.Type.NV12,
+                        fps=config.fps,
+                    )
+                    encoder = self.pipeline.create(dai.node.VideoEncoder)
+                    encoder.setDefaultProfilePreset(
+                        config.fps, dai.VideoEncoderProperties.Profile.MJPEG
+                    )
+                    encoder.setQuality(config.mjpeg_quality)
+                    cam_out.link(encoder.input)
+                    self.output_queues["mono_right"] = encoder.out.createOutputQueue(
+                        maxSize=3, blocking=False
+                    )
+                else:
+                    self.output_queues["mono_right"] = self.cam_mono_right.requestOutput(
+                        config.monochrome_image_dim,
+                        fps=config.fps,
+                    ).createOutputQueue(maxSize=3, blocking=False)
+                print("Enabled CAM_C (Monochrome Right)")
+
+        assert len(self.output_queues) > 0, "No output queues enabled"
+
+        self.pipeline.start()
+
+        print(f"[{mount_position}] Pipeline started, waiting for stabilization...")
+        time.sleep(2.0)
+
+        for _ in range(10):
+            test_frame = None
+            for queue_name, q in self.output_queues.items():
+                test_frame = q.tryGet()
+                if test_frame:
+                    print(f"[{mount_position}] First frame received from {queue_name}")
+                    break
+            if test_frame:
+                break
+            time.sleep(0.3)
+        else:
+            print(f"[{mount_position}] Warning: No frames received during init verification")
+
+        if run_as_server:
+            self.start_server(port)
+
+    def read(self) -> dict[str, Any] | None:
+        if not self.pipeline.isRunning():
+            print(f"[ERROR] OAK pipeline stopped for {self.mount_position}")
+            return None
+        if not self.device.isPipelineRunning():
+            print(f"[ERROR] OAK device disconnected for {self.mount_position}")
+            return None
+
+        timestamps = {}
+        images = {}
+        rgb_frame_time = None
+
+        def drain_queue_get_latest(queue):
+            latest_frame = None
+            while True:
+                frame = queue.tryGet()
+                if frame is None:
+                    break
+                latest_frame = frame
+            return latest_frame
+
+        expected_cameras = set(self.output_queues.keys())
+        received_cameras = set()
+
+        if "color" in self.output_queues:
+            try:
+                rgb_frame = drain_queue_get_latest(self.output_queues["color"])
+                if rgb_frame is None:
+                    return None
+                rgb_frame_time = rgb_frame.getTimestamp()
+                read_time = time.time()
+                frame_age = (dai.Clock.now() - rgb_frame_time).total_seconds()
+                capture_time = read_time - frame_age
+
+                if self._use_mjpeg:
+                    images[self.mount_position] = bytes(rgb_frame.getData())
+                else:
+                    images[self.mount_position] = rgb_frame.getCvFrame()[..., ::-1]
+                timestamps[self.mount_position] = capture_time
+                received_cameras.add("color")
+            except Exception as e:
+                print(f"[ERROR] Failed to read color frame from {self.mount_position}: {e}")
+                return None
+
+        if "mono_left" in self.output_queues:
+            try:
+                mono_left_frame = drain_queue_get_latest(self.output_queues["mono_left"])
+                if mono_left_frame is None:
+                    return None
+                mono_left_frame_time = mono_left_frame.getTimestamp()
+                read_time = time.time()
+                frame_age = (dai.Clock.now() - mono_left_frame_time).total_seconds()
+                capture_time = read_time - frame_age
+
+                key = f"{self.mount_position}_left_mono"
+                if self._use_mjpeg:
+                    images[key] = bytes(mono_left_frame.getData())
+                else:
+                    images[key] = mono_left_frame.getCvFrame()
+                timestamps[key] = capture_time
+                received_cameras.add("mono_left")
+            except Exception as e:
+                print(f"[ERROR] Failed to read mono_left frame from {self.mount_position}: {e}")
+                return None
+
+        if "mono_right" in self.output_queues:
+            try:
+                mono_right_frame = drain_queue_get_latest(self.output_queues["mono_right"])
+                if mono_right_frame is None:
+                    return None
+                mono_right_frame_time = mono_right_frame.getTimestamp()
+                read_time = time.time()
+                frame_age = (dai.Clock.now() - mono_right_frame_time).total_seconds()
+                capture_time = read_time - frame_age
+
+                key = f"{self.mount_position}_right_mono"
+                if self._use_mjpeg:
+                    images[key] = bytes(mono_right_frame.getData())
+                else:
+                    images[key] = mono_right_frame.getCvFrame()
+                timestamps[key] = capture_time
+                received_cameras.add("mono_right")
+            except Exception as e:
+                print(f"[ERROR] Failed to read mono_right frame from {self.mount_position}: {e}")
+                return None
+
+        if received_cameras != expected_cameras:
+            missing = expected_cameras - received_cameras
+            print(f"[ERROR] Missing frames from cameras: {missing} for {self.mount_position}")
+            return None
+
+        if rgb_frame_time is not None:
+            frame_age = (dai.Clock.now() - rgb_frame_time).total_seconds()
+            if frame_age > 0.1:
+                print(
+                    f"[{self.mount_position}] OAK frame age too large: {frame_age * 1000:.1f}ms"
+                )
+
+        return {"timestamps": timestamps, "images": images}
+
+    def serialize(self, data: dict[str, Any]) -> dict[str, Any]:
+        serialized_msg = ImageMessageSchema(timestamps=data["timestamps"], images=data["images"])
+        return serialized_msg.serialize()
+
+    def observation_space(self):
+        if gym is None:
+            return None
+        spaces = {}
+        if self.config.enable_color:
+            spaces["color_image"] = gym.spaces.Box(
+                low=0,
+                high=255,
+                shape=(self.config.color_image_dim[1], self.config.color_image_dim[0], 3),
+                dtype=np.uint8,
+            )
+        if self.config.enable_mono_cameras:
+            spaces["mono_left_image"] = gym.spaces.Box(
+                low=0,
+                high=255,
+                shape=(self.config.monochrome_image_dim[1], self.config.monochrome_image_dim[0]),
+                dtype=np.uint8,
+            )
+            spaces["mono_right_image"] = gym.spaces.Box(
+                low=0,
+                high=255,
+                shape=(self.config.monochrome_image_dim[1], self.config.monochrome_image_dim[0]),
+                dtype=np.uint8,
+            )
+        return gym.spaces.Dict(spaces)
+
+    def close(self):
+        if self._run_as_server:
+            self.stop_server()
+        if hasattr(self, "pipeline") and self.pipeline.isRunning():
+            self.pipeline.stop()
+        self.device.close()
+
+    def run_server(self):
+        if not self._run_as_server:
+            raise ValueError("run_as_server must be True to call run_server()")
+        while True:
+            frame = self.read()
+            if frame is None:
+                continue
+            msg = self.serialize(frame)
+            self.send_message({self.mount_position: msg})
+
+    def __del__(self):
+        self.close()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--server", action="store_true", help="Run as server")
+    parser.add_argument("--host", type=str, default="localhost", help="Server IP address")
+    parser.add_argument("--port", type=int, default=5555, help="Port number")
+    parser.add_argument("--device-id", type=str, default=None, help="Specific device ID")
+    parser.add_argument(
+        "--enable-mono", action="store_true", help="Enable monochrome cameras (CAM_B & CAM_C)"
+    )
+    parser.add_argument("--mount-position", type=str, default="ego_view", help="Mount position")
+    parser.add_argument("--show-image", action="store_true", help="Display images")
+    parser.add_argument("--use-mjpeg", action="store_true", help="Use MJPEG encoding on-device")
+    parser.add_argument(
+        "--mjpeg-quality", type=int, default=80, help="MJPEG quality 1-100 (default: 80)"
+    )
+    args = parser.parse_args()
+
+    oak_config = OAKConfig()
+    if args.enable_mono:
+        oak_config.enable_mono_cameras = True
+    if args.use_mjpeg:
+        oak_config.use_mjpeg = True
+        oak_config.mjpeg_quality = args.mjpeg_quality
+
+    if args.server:
+        oak = OAKSensor(
+            run_as_server=True,
+            port=args.port,
+            config=oak_config,
+            device_id=args.device_id,
+            mount_position=args.mount_position,
+        )
+        print(f"Starting OAK server on port {args.port}")
+        oak.run_server()
+    else:
+        oak = OAKSensor(run_as_server=False, config=oak_config, device_id=args.device_id)
+        print("Running OAK camera in standalone mode")
+
+        while True:
+            frame = oak.read()
+            if frame is None:
+                print("Waiting for frame...")
+                time.sleep(0.5)
+                continue
+
+            if args.show_image:
+                for key, img in frame.get("images", {}).items():
+                    if isinstance(img, np.ndarray):
+                        cv2.imshow(key, img[..., ::-1] if img.ndim == 3 and img.shape[2] == 3 else img)
+                if cv2.waitKey(1) == ord("q"):
+                    break
+
+            time.sleep(0.01)
+
+        cv2.destroyAllWindows()
+        oak.close()
diff --git a/GR00T-WholeBodyControl/gear_sonic/camera/drivers/realsense.py b/GR00T-WholeBodyControl/gear_sonic/camera/drivers/realsense.py
new file mode 100644
index 0000000000000000000000000000000000000000..7024594bbfa6ba41a57df2eb1ec992c397ca4028
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/camera/drivers/realsense.py
@@ -0,0 +1,173 @@
+"""Intel RealSense camera driver.
+
+Requires the ``pyrealsense2`` SDK — install with::
+
+    pip install pyrealsense2
+
+See https://github.com/IntelRealSense/librealsense for hardware-specific instructions.
+"""
+
+import time
+from typing import Any
+
+import numpy as np
+
+try:
+    import gymnasium as gym
+except ImportError:
+    gym = None  # type: ignore[assignment]
+
+import pyrealsense2 as rs
+
+from gear_sonic.camera.sensor import Sensor
+from gear_sonic.camera.sensor_server import (
+    CameraMountPosition,
+    ImageMessageSchema,
+    SensorServer,
+)
+
+
+class RealSenseConfig:
+    """Configuration for the RealSense camera."""
+
+    depth_image_dim: tuple[int, int] = (640, 480)
+    color_image_dim: tuple[int, int] = (640, 480)
+    fps: int = 30
+    mount_position: str = CameraMountPosition.EGO_VIEW.value
+
+
+class RealSenseSensor(Sensor, SensorServer):
+    """Sensor for Intel RealSense depth cameras."""
+
+    def __init__(
+        self,
+        run_as_server: bool = False,
+        port: int = 5555,
+        config: RealSenseConfig = RealSenseConfig(),
+        id: int = 0,
+        mount_position: str = CameraMountPosition.EGO_VIEW.value,
+    ):
+        devices = rs.context().query_devices()
+        if len(devices) == 0:
+            raise RuntimeError("No RealSense devices found")
+
+        for device in devices:
+            print(f"Device: {device.get_info(rs.camera_info.name)}")
+            print(f"    Serial number: {device.get_info(rs.camera_info.serial_number)}")
+            print(f"    Firmware version: {device.get_info(rs.camera_info.firmware_version)}")
+
+        self.pipeline = rs.pipeline()
+        self.config = rs.config()
+        devices = sorted(devices, key=lambda x: x.get_info(rs.camera_info.serial_number))
+        self.config.enable_device(devices[id].get_info(rs.camera_info.serial_number))
+
+        try:
+            self.config.enable_stream(
+                rs.stream.color,
+                config.color_image_dim[0],
+                config.color_image_dim[1],
+                rs.format.rgb8,
+                config.fps,
+            )
+            self.config.enable_stream(
+                rs.stream.depth,
+                config.depth_image_dim[0],
+                config.depth_image_dim[1],
+                rs.format.z16,
+                config.fps,
+            )
+            self.pipeline.start(self.config)
+        except Exception as e:
+            raise RuntimeError(f"Failed to start RealSense pipeline: {e}")
+
+        self._realsense_config = config
+        self._run_as_server = run_as_server
+        self.mount_position = mount_position
+        if self._run_as_server:
+            self.start_server(port)
+        print(
+            f"Done initializing RealSense sensor: "
+            f"{devices[id].get_info(rs.camera_info.serial_number)}"
+        )
+
+    def read(self) -> dict[str, Any] | None:
+        try:
+            frames = self.pipeline.wait_for_frames()
+        except Exception as e:
+            print(f"ERROR! Failed to wait for frames: {e}")
+            return None
+
+        color_frame = frames.get_color_frame()
+        depth_frame = frames.get_depth_frame()
+
+        if not color_frame or not depth_frame:
+            print("WARNING! No color or depth frame")
+            return None
+
+        try:
+            color_image = np.asanyarray(color_frame.get_data())
+            depth_image = np.asanyarray(depth_frame.get_data())
+        except Exception as e:
+            print(f"ERROR! Failed to convert frames to numpy arrays: {e}")
+            return None
+
+        if color_image.size == 0 or depth_image.size == 0:
+            print("WARNING! Empty color or depth image")
+            return None
+
+        current_time = time.time()
+        timestamps = {
+            self.mount_position: current_time,
+            f"{self.mount_position}_depth": current_time,
+        }
+        images = {
+            self.mount_position: color_image,
+            f"{self.mount_position}_depth": depth_image,
+        }
+        return {"timestamps": timestamps, "images": images}
+
+    def serialize(self, data: dict[str, Any]) -> dict[str, Any]:
+        serialized_msg = ImageMessageSchema(timestamps=data["timestamps"], images=data["images"])
+        return serialized_msg.serialize()
+
+    def observation_space(self):
+        if gym is None:
+            return None
+        return gym.spaces.Dict(
+            {
+                "color_image": gym.spaces.Box(
+                    low=0,
+                    high=255,
+                    shape=(
+                        self._realsense_config.color_image_dim[1],
+                        self._realsense_config.color_image_dim[0],
+                        3,
+                    ),
+                    dtype=np.uint8,
+                ),
+                "depth_image": gym.spaces.Box(
+                    low=0,
+                    high=255,
+                    shape=(
+                        self._realsense_config.depth_image_dim[1],
+                        self._realsense_config.depth_image_dim[0],
+                        1,
+                    ),
+                    dtype=np.uint16,
+                ),
+            }
+        )
+
+    def close(self):
+        if self._run_as_server:
+            self.stop_server()
+        self.pipeline.stop()
+
+    def run_server(self):
+        if not self._run_as_server:
+            raise ValueError("run_as_server must be True to call run_server()")
+        while True:
+            read_result = self.read()
+            if read_result is None:
+                continue
+            self.send_message({self.mount_position: self.serialize(read_result)})
diff --git a/GR00T-WholeBodyControl/gear_sonic/camera/drivers/usb_camera.py b/GR00T-WholeBodyControl/gear_sonic/camera/drivers/usb_camera.py
new file mode 100644
index 0000000000000000000000000000000000000000..9c196b5b646011cff3637751d74f8a76b77b8ad7
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/camera/drivers/usb_camera.py
@@ -0,0 +1,100 @@
+"""Generic USB webcam driver using OpenCV.
+
+No hardware SDK needed — works with any UVC-compatible camera visible as
+``/dev/video*``.  Only requires ``opencv-python``.
+"""
+
+import time
+from typing import Any
+
+import cv2
+import numpy as np
+
+try:
+    import gymnasium as gym
+except ImportError:
+    gym = None  # type: ignore[assignment]
+
+from gear_sonic.camera.sensor import Sensor
+from gear_sonic.camera.sensor_server import CameraMountPosition
+
+
+class USBCameraConfig:
+    """Configuration for generic USB camera."""
+
+    image_dim: tuple = (640, 480)
+    fps: int = 30
+    device_index: int = 0
+
+
+class USBCameraSensor(Sensor):
+    """Sensor for generic USB cameras using OpenCV VideoCapture."""
+
+    def __init__(
+        self,
+        config: USBCameraConfig = USBCameraConfig(),
+        mount_position: str = CameraMountPosition.EGO_VIEW.value,
+        device_index: int | None = None,
+    ):
+        self.config = config
+        self.mount_position = mount_position
+
+        idx = device_index if device_index is not None else config.device_index
+
+        self.cap = cv2.VideoCapture(idx)
+        if not self.cap.isOpened():
+            raise RuntimeError(f"Failed to open USB camera at index {idx}")
+
+        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, config.image_dim[0])
+        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, config.image_dim[1])
+        self.cap.set(cv2.CAP_PROP_FPS, config.fps)
+        self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
+
+        print(f"[{mount_position}] Warming up USB camera...")
+        for _ in range(10):
+            ret, _ = self.cap.read()
+            if ret:
+                break
+            time.sleep(0.1)
+
+        print(f"[{mount_position}] USB camera opened at index {idx}")
+        width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        print(f"  Resolution: {width}x{height}")
+        print(f"  FPS: {self.cap.get(cv2.CAP_PROP_FPS)}")
+
+    def read(self) -> dict[str, Any] | None:
+        ret, frame = self.cap.read()
+        if not ret or frame is None:
+            print(f"[{self.mount_position}] USB camera read failed: ret={ret}")
+            return None
+
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        return {
+            "timestamps": {self.mount_position: time.time()},
+            "images": {self.mount_position: frame_rgb},
+        }
+
+    def serialize(self, data: dict[str, Any]) -> dict[str, Any]:
+        from gear_sonic.camera.sensor_server import ImageMessageSchema
+
+        serialized_msg = ImageMessageSchema(timestamps=data["timestamps"], images=data["images"])
+        return serialized_msg.serialize()
+
+    def observation_space(self):
+        if gym is None:
+            return None
+        return gym.spaces.Dict(
+            {
+                "color_image": gym.spaces.Box(
+                    low=0,
+                    high=255,
+                    shape=(self.config.image_dim[1], self.config.image_dim[0], 3),
+                    dtype=np.uint8,
+                ),
+            }
+        )
+
+    def close(self):
+        if self.cap is not None:
+            self.cap.release()
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/left_ankle_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/left_ankle_pitch_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..69de8490184afc698f633e34bb74e65351bd4689
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/left_ankle_pitch_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/left_elbow_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/left_elbow_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..1a96d99ba469960173129084ab6dd3bf8a732a71
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/left_elbow_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/left_hand_thumb_0_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/left_hand_thumb_0_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..3028bb4d6e1ae3d30d2504259c08e4106bbacf63
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/left_hand_thumb_0_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/left_wrist_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/left_wrist_pitch_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..82cc224a8e41251d879502f9809e31d0988ec7f9
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/left_wrist_pitch_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/right_ankle_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/right_ankle_pitch_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..e77d8a2fe1e5d56fac049833d254d6ffa4f6b350
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/right_ankle_pitch_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/right_elbow_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/right_elbow_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..f259e3812efb9985d6463c04d3e8a4b53793a699
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/right_elbow_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/right_hand_thumb_0_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/right_hand_thumb_0_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..1cae7f18e16605cb9d6b1d1a0cf6e5c5c360a344
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/right_hand_thumb_0_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/right_wrist_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/right_wrist_pitch_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..da194543c40df9d492abb0553c06cc614a78caae
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/right_wrist_pitch_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/torso_constraint_L_rod_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/torso_constraint_L_rod_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..6747f3f9341bd72b3803385c135c3ce750322e9d
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/torso_constraint_L_rod_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/torso_constraint_R_rod_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/torso_constraint_R_rod_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..95cf415f72f1679a5867ca21a4af774f0b217cad
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/torso_constraint_R_rod_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/waist_roll_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/waist_roll_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..65831abd2a64bc8c36e31016964413a3d2116725
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/waist_roll_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/waist_roll_link_rev_1_0.STL b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/waist_roll_link_rev_1_0.STL
new file mode 100644
index 0000000000000000000000000000000000000000..a64f330c592582dc31cdf38f4d08cfff06681c5f
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/meshes/g1/waist_roll_link_rev_1_0.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/mjcf/g1_29dof_rev_1_0.xml b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/mjcf/g1_29dof_rev_1_0.xml
new file mode 100644
index 0000000000000000000000000000000000000000..9123bedd60eb0e90393c1fed4a945152259f0a08
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/mjcf/g1_29dof_rev_1_0.xml
@@ -0,0 +1,472 @@
+<mujoco model="g1_29dof_rev_1_0">
+    <compiler angle="radian" meshdir="../meshes/g1/" />
+
+    <asset>
+        <mesh name="pelvis" file="pelvis.STL" />
+        <mesh name="pelvis_contour_link" file="pelvis_contour_link.STL" />
+        <mesh name="left_hip_pitch_link" file="left_hip_pitch_link.STL" />
+        <mesh name="left_hip_roll_link" file="left_hip_roll_link.STL" />
+        <mesh name="left_hip_yaw_link" file="left_hip_yaw_link.STL" />
+        <mesh name="left_knee_link" file="left_knee_link.STL" />
+        <mesh name="left_ankle_pitch_link" file="left_ankle_pitch_link.STL" />
+        <mesh name="left_ankle_roll_link" file="left_ankle_roll_link.STL" />
+        <mesh name="right_hip_pitch_link" file="right_hip_pitch_link.STL" />
+        <mesh name="right_hip_roll_link" file="right_hip_roll_link.STL" />
+        <mesh name="right_hip_yaw_link" file="right_hip_yaw_link.STL" />
+        <mesh name="right_knee_link" file="right_knee_link.STL" />
+        <mesh name="right_ankle_pitch_link" file="right_ankle_pitch_link.STL" />
+        <mesh name="right_ankle_roll_link" file="right_ankle_roll_link.STL" />
+        <mesh name="waist_yaw_link" file="waist_yaw_link_rev_1_0.STL" />
+        <mesh name="waist_roll_link" file="waist_roll_link_rev_1_0.STL" />
+        <mesh name="torso_link" file="torso_link_rev_1_0.STL" />
+        <mesh name="logo_link" file="logo_link.STL" />
+        <mesh name="head_link" file="head_link.STL" />
+        <mesh name="left_shoulder_pitch_link" file="left_shoulder_pitch_link.STL" />
+        <mesh name="left_shoulder_roll_link" file="left_shoulder_roll_link.STL" />
+        <mesh name="left_shoulder_yaw_link" file="left_shoulder_yaw_link.STL" />
+        <mesh name="left_elbow_link" file="left_elbow_link.STL" />
+        <mesh name="left_wrist_roll_link" file="left_wrist_roll_link.STL" />
+        <mesh name="left_wrist_pitch_link" file="left_wrist_pitch_link.STL" />
+        <mesh name="left_wrist_yaw_link" file="left_wrist_yaw_link.STL" />
+        <mesh name="left_rubber_hand" file="left_rubber_hand.STL" />
+        <mesh name="right_shoulder_pitch_link" file="right_shoulder_pitch_link.STL" />
+        <mesh name="right_shoulder_roll_link" file="right_shoulder_roll_link.STL" />
+        <mesh name="right_shoulder_yaw_link" file="right_shoulder_yaw_link.STL" />
+        <mesh name="right_elbow_link" file="right_elbow_link.STL" />
+        <mesh name="right_wrist_roll_link" file="right_wrist_roll_link.STL" />
+        <mesh name="right_wrist_pitch_link" file="right_wrist_pitch_link.STL" />
+        <mesh name="right_wrist_yaw_link" file="right_wrist_yaw_link.STL" />
+        <mesh name="right_rubber_hand" file="right_rubber_hand.STL" />
+    </asset>
+
+    <worldbody>
+        <body name="pelvis" pos="0 0 0.793">
+            <inertial pos="0 0 -0.07605" quat="1 0 -0.000399148 0" mass="3.813"
+                diaginertia="0.010549 0.0093089 0.0079184" />
+            <joint name="floating_base_joint" type="free" limited="false" actuatorfrclimited="false" />
+            <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.2 0.2 0.2 1"
+                mesh="pelvis" />
+            <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1"
+                mesh="pelvis_contour_link" />
+            <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="pelvis_contour_link" />
+            <site name="imu_in_pelvis" size="0.01" pos="0.04525 0 -0.08339" />
+            <body name="left_hip_pitch_link" pos="0 0.064452 -0.1027">
+                <inertial pos="0.002741 0.047791 -0.02606"
+                    quat="0.954862 0.293964 0.0302556 0.030122" mass="1.35"
+                    diaginertia="0.00181517 0.00153422 0.00116212" />
+                <joint name="left_hip_pitch_joint" pos="0 0 0" axis="0 1 0" range="-2.5307 2.8798"
+                    actuatorfrcrange="-88 88" />
+                <geom type="mesh" contype="0" conaffinity="0" group="1" density="0"
+                    rgba="0.2 0.2 0.2 1" mesh="left_hip_pitch_link" />
+                <geom type="mesh" rgba="0.2 0.2 0.2 1" mesh="left_hip_pitch_link" />
+                <body name="left_hip_roll_link" pos="0 0.052 -0.030465"
+                    quat="0.996179 0 -0.0873386 0">
+                    <inertial pos="0.029812 -0.001045 -0.087934"
+                        quat="0.977808 -1.97119e-05 0.205576 -0.0403793" mass="1.52"
+                        diaginertia="0.00254986 0.00241169 0.00148755" />
+                    <joint name="left_hip_roll_joint" pos="0 0 0" axis="1 0 0"
+                        range="-0.5236 2.9671" actuatorfrcrange="-139 139" />
+                    <geom type="mesh" contype="0" conaffinity="0" group="1" density="0"
+                        rgba="0.7 0.7 0.7 1" mesh="left_hip_roll_link" />
+                    <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="left_hip_roll_link" />
+                    <body name="left_hip_yaw_link" pos="0.025001 0 -0.12412">
+                        <inertial pos="-0.057709 -0.010981 -0.15078"
+                            quat="0.600598 0.15832 0.223482 0.751181" mass="1.702"
+                            diaginertia="0.00776166 0.00717575 0.00160139" />
+                        <joint name="left_hip_yaw_joint" pos="0 0 0" axis="0 0 1"
+                            range="-2.7576 2.7576" actuatorfrcrange="-88 88" />
+                        <geom type="mesh" contype="0" conaffinity="0" group="1" density="0"
+                            rgba="0.7 0.7 0.7 1" mesh="left_hip_yaw_link" />
+                        <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="left_hip_yaw_link" />
+                        <body name="left_knee_link" pos="-0.078273 0.0021489 -0.17734"
+                            quat="0.996179 0 0.0873386 0">
+                            <inertial pos="0.005457 0.003964 -0.12074"
+                                quat="0.923418 -0.0327699 0.0158246 0.382067" mass="1.932"
+                                diaginertia="0.0113804 0.0112778 0.00146458" />
+                            <joint name="left_knee_joint" pos="0 0 0" axis="0 1 0"
+                                range="-0.087267 2.8798" actuatorfrcrange="-139 139" />
+                            <geom type="mesh" contype="0" conaffinity="0" group="1" density="0"
+                                rgba="0.7 0.7 0.7 1" mesh="left_knee_link" />
+                            <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="left_knee_link" />
+                            <body name="left_ankle_pitch_link" pos="0 -9.4445e-05 -0.30001">
+                                <inertial pos="-0.007269 0 0.011137"
+                                    quat="0.603053 0.369225 0.369225 0.603053" mass="0.074"
+                                    diaginertia="1.89e-05 1.40805e-05 6.9195e-06" />
+                                <joint name="left_ankle_pitch_joint" pos="0 0 0" axis="0 1 0"
+                                    range="-0.87267 0.5236" actuatorfrcrange="-50 50" />
+                                <geom type="mesh" contype="0" conaffinity="0" group="1" density="0"
+                                    rgba="0.7 0.7 0.7 1" mesh="left_ankle_pitch_link" />
+                                <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="left_ankle_pitch_link" />
+                                <body name="left_ankle_roll_link" pos="0 0 -0.017558">
+                                    <inertial pos="0.026505 0 -0.016425"
+                                        quat="-0.000481092 0.728482 -0.000618967 0.685065"
+                                        mass="0.608" diaginertia="0.00167218 0.0016161 0.000217621" />
+                                    <joint name="left_ankle_roll_joint" pos="0 0 0" axis="1 0 0"
+                                        range="-0.2618 0.2618" actuatorfrcrange="-50 50" />
+                                    <geom type="mesh" contype="0" conaffinity="0" group="1"
+                                        density="0" rgba="0.2 0.2 0.2 1" mesh="left_ankle_roll_link" />
+                                    <geom size="0.005" pos="-0.05 0.025 -0.03" rgba="0.2 0.2 0.2 1" />
+                                    <geom size="0.005" pos="-0.05 -0.025 -0.03" rgba="0.2 0.2 0.2 1" />
+                                    <geom size="0.005" pos="0.12 0.03 -0.03" rgba="0.2 0.2 0.2 1" />
+                                    <geom size="0.005" pos="0.12 -0.03 -0.03" rgba="0.2 0.2 0.2 1" />
+                                </body>
+                            </body>
+                        </body>
+                    </body>
+                </body>
+            </body>
+            <body name="right_hip_pitch_link" pos="0 -0.064452 -0.1027">
+                <inertial pos="0.002741 -0.047791 -0.02606"
+                    quat="0.954862 -0.293964 0.0302556 -0.030122" mass="1.35"
+                    diaginertia="0.00181517 0.00153422 0.00116212" />
+                <joint name="right_hip_pitch_joint" pos="0 0 0" axis="0 1 0" range="-2.5307 2.8798"
+                    actuatorfrcrange="-88 88" />
+                <geom type="mesh" contype="0" conaffinity="0" group="1" density="0"
+                    rgba="0.2 0.2 0.2 1" mesh="right_hip_pitch_link" />
+                <geom type="mesh" rgba="0.2 0.2 0.2 1" mesh="right_hip_pitch_link" />
+                <body name="right_hip_roll_link" pos="0 -0.052 -0.030465"
+                    quat="0.996179 0 -0.0873386 0">
+                    <inertial pos="0.029812 0.001045 -0.087934"
+                        quat="0.977808 1.97119e-05 0.205576 0.0403793" mass="1.52"
+                        diaginertia="0.00254986 0.00241169 0.00148755" />
+                    <joint name="right_hip_roll_joint" pos="0 0 0" axis="1 0 0"
+                        range="-2.9671 0.5236" actuatorfrcrange="-139 139" />
+                    <geom type="mesh" contype="0" conaffinity="0" group="1" density="0"
+                        rgba="0.7 0.7 0.7 1" mesh="right_hip_roll_link" />
+                    <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="right_hip_roll_link" />
+                    <body name="right_hip_yaw_link" pos="0.025001 0 -0.12412">
+                        <inertial pos="-0.057709 0.010981 -0.15078"
+                            quat="0.751181 0.223482 0.15832 0.600598" mass="1.702"
+                            diaginertia="0.00776166 0.00717575 0.00160139" />
+                        <joint name="right_hip_yaw_joint" pos="0 0 0" axis="0 0 1"
+                            range="-2.7576 2.7576" actuatorfrcrange="-88 88" />
+                        <geom type="mesh" contype="0" conaffinity="0" group="1" density="0"
+                            rgba="0.7 0.7 0.7 1" mesh="right_hip_yaw_link" />
+                        <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="right_hip_yaw_link" />
+                        <body name="right_knee_link" pos="-0.078273 -0.0021489 -0.17734"
+                            quat="0.996179 0 0.0873386 0">
+                            <inertial pos="0.005457 -0.003964 -0.12074"
+                                quat="0.923439 0.0345276 0.0116333 -0.382012" mass="1.932"
+                                diaginertia="0.011374 0.0112843 0.00146452" />
+                            <joint name="right_knee_joint" pos="0 0 0" axis="0 1 0"
+                                range="-0.087267 2.8798" actuatorfrcrange="-139 139" />
+                            <geom type="mesh" contype="0" conaffinity="0" group="1" density="0"
+                                rgba="0.7 0.7 0.7 1" mesh="right_knee_link" />
+                            <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="right_knee_link" />
+                            <body name="right_ankle_pitch_link" pos="0 9.4445e-05 -0.30001">
+                                <inertial pos="-0.007269 0 0.011137"
+                                    quat="0.603053 0.369225 0.369225 0.603053" mass="0.074"
+                                    diaginertia="1.89e-05 1.40805e-05 6.9195e-06" />
+                                <joint name="right_ankle_pitch_joint" pos="0 0 0" axis="0 1 0"
+                                    range="-0.87267 0.5236" actuatorfrcrange="-50 50" />
+                                <geom type="mesh" contype="0" conaffinity="0" group="1" density="0"
+                                    rgba="0.7 0.7 0.7 1" mesh="right_ankle_pitch_link" />
+                                <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="right_ankle_pitch_link" />
+                                <body name="right_ankle_roll_link" pos="0 0 -0.017558">
+                                    <inertial pos="0.026505 0 -0.016425"
+                                        quat="0.000481092 0.728482 0.000618967 0.685065"
+                                        mass="0.608" diaginertia="0.00167218 0.0016161 0.000217621" />
+                                    <joint name="right_ankle_roll_joint" pos="0 0 0" axis="1 0 0"
+                                        range="-0.2618 0.2618" actuatorfrcrange="-50 50" />
+                                    <geom type="mesh" contype="0" conaffinity="0" group="1"
+                                        density="0" rgba="0.2 0.2 0.2 1"
+                                        mesh="right_ankle_roll_link" />
+                                    <geom size="0.005" pos="-0.05 0.025 -0.03" rgba="0.2 0.2 0.2 1" />
+                                    <geom size="0.005" pos="-0.05 -0.025 -0.03" rgba="0.2 0.2 0.2 1" />
+                                    <geom size="0.005" pos="0.12 0.03 -0.03" rgba="0.2 0.2 0.2 1" />
+                                    <geom size="0.005" pos="0.12 -0.03 -0.03" rgba="0.2 0.2 0.2 1" />
+                                </body>
+                            </body>
+                        </body>
+                    </body>
+                </body>
+            </body>
+            <body name="waist_yaw_link">
+                <inertial pos="0.003494 0.000233 0.018034"
+                    quat="0.289697 0.591001 -0.337795 0.672821" mass="0.214"
+                    diaginertia="0.000163531 0.000107714 0.000102205" />
+                <joint name="waist_yaw_joint" pos="0 0 0" axis="0 0 1" range="-2.618 2.618"
+                    actuatorfrcrange="-88 88" />
+                <geom type="mesh" contype="0" conaffinity="0" group="1" density="0"
+                    rgba="0.7 0.7 0.7 1" mesh="waist_yaw_link" />
+                <body name="waist_roll_link" pos="-0.0039635 0 0.044">
+                    <inertial pos="0 2.3e-05 0" quat="0.5 0.5 -0.5 0.5" mass="0.086"
+                        diaginertia="8.245e-06 7.079e-06 6.339e-06" />
+                    <joint name="waist_roll_joint" pos="0 0 0" axis="1 0 0" range="-0.52 0.52"
+                        actuatorfrcrange="-50 50" />
+                    <geom type="mesh" contype="0" conaffinity="0" group="1" density="0"
+                        rgba="0.7 0.7 0.7 1" mesh="waist_roll_link" />
+                    <body name="torso_link">
+                        <inertial pos="0.00203158 0.000339683 0.184568"
+                            quat="0.999803 -6.03319e-05 0.0198256 0.00131986" mass="7.818"
+                            diaginertia="0.121847 0.109825 0.0273735" />
+                        <joint name="waist_pitch_joint" pos="0 0 0" axis="0 1 0" range="-0.52 0.52"
+                            actuatorfrcrange="-50 50" />
+                        <geom type="mesh" contype="0" conaffinity="0" group="1" density="0"
+                            rgba="0.7 0.7 0.7 1" mesh="torso_link" />
+                        <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="torso_link" />
+                        <geom pos="0.0039635 0 -0.044" quat="1 0 0 0" type="mesh" contype="0"
+                            conaffinity="0" group="1" density="0" rgba="0.2 0.2 0.2 1"
+                            mesh="logo_link" />
+                        <geom pos="0.0039635 0 -0.044" quat="1 0 0 0" type="mesh"
+                            rgba="0.2 0.2 0.2 1" mesh="logo_link" />
+                        <geom pos="0.0039635 0 -0.044" type="mesh" contype="0" conaffinity="0"
+                            group="1" density="0" rgba="0.2 0.2 0.2 1" mesh="head_link" />
+                        <geom pos="0.0039635 0 -0.044" type="mesh" rgba="0.2 0.2 0.2 1"
+                            mesh="head_link" />
+                        <site name="imu_in_torso" size="0.01" pos="-0.03959 -0.00224 0.14792" />
+                        <body name="left_shoulder_pitch_link" pos="0.0039563 0.10022 0.24778"
+                            quat="0.990264 0.139201 1.38722e-05 -9.86868e-05">
+                            <inertial pos="0 0.035892 -0.011628"
+                                quat="0.654152 0.0130458 -0.326267 0.68225" mass="0.718"
+                                diaginertia="0.000465864 0.000432842 0.000406394" />
+                            <joint name="left_shoulder_pitch_joint" pos="0 0 0" axis="0 1 0"
+                                range="-3.0892 2.6704" actuatorfrcrange="-25 25" />
+                            <geom type="mesh" contype="0" conaffinity="0" group="1" density="0"
+                                rgba="0.7 0.7 0.7 1" mesh="left_shoulder_pitch_link" />
+                            <geom size="0.03 0.025" pos="0 0.04 -0.01" quat="0.707107 0 0.707107 0"
+                                type="cylinder" rgba="0.7 0.7 0.7 1" />
+                            <body name="left_shoulder_roll_link" pos="0 0.038 -0.013831"
+                                quat="0.990268 -0.139172 0 0">
+                                <inertial pos="-0.000227 0.00727 -0.063243"
+                                    quat="0.701256 -0.0196223 -0.00710317 0.712604" mass="0.643"
+                                    diaginertia="0.000691311 0.000618011 0.000388977" />
+                                <joint name="left_shoulder_roll_joint" pos="0 0 0" axis="1 0 0"
+                                    range="-1.5882 2.2515" actuatorfrcrange="-25 25" />
+                                <geom type="mesh" contype="0" conaffinity="0" group="1" density="0"
+                                    rgba="0.7 0.7 0.7 1" mesh="left_shoulder_roll_link" />
+                                <geom size="0.03 0.015" pos="-0.004 0.006 -0.053" type="cylinder"
+                                    rgba="0.7 0.7 0.7 1" />
+                                <body name="left_shoulder_yaw_link" pos="0 0.00624 -0.1032">
+                                    <inertial pos="0.010773 -0.002949 -0.072009"
+                                        quat="0.716879 -0.0964829 -0.0679942 0.687134" mass="0.734"
+                                        diaginertia="0.00106187 0.00103217 0.000400661" />
+                                    <joint name="left_shoulder_yaw_joint" pos="0 0 0" axis="0 0 1"
+                                        range="-2.618 2.618" actuatorfrcrange="-25 25" />
+                                    <geom type="mesh" contype="0" conaffinity="0" group="1"
+                                        density="0" rgba="0.7 0.7 0.7 1"
+                                        mesh="left_shoulder_yaw_link" />
+                                    <geom type="mesh" rgba="0.7 0.7 0.7 1"
+                                        mesh="left_shoulder_yaw_link" />
+                                    <body name="left_elbow_link" pos="0.015783 0 -0.080518">
+                                        <inertial pos="0.064956 0.004454 -0.010062"
+                                            quat="0.541765 0.636132 0.388821 0.388129" mass="0.6"
+                                            diaginertia="0.000443035 0.000421612 0.000259353" />
+                                        <joint name="left_elbow_joint" pos="0 0 0" axis="0 1 0"
+                                            range="-1.0472 2.0944" actuatorfrcrange="-25 25" />
+                                        <geom type="mesh" contype="0" conaffinity="0" group="1"
+                                            density="0" rgba="0.7 0.7 0.7 1" mesh="left_elbow_link" />
+                                        <geom type="mesh" rgba="0.7 0.7 0.7 1"
+                                            mesh="left_elbow_link" />
+                                        <body name="left_wrist_roll_link" pos="0.1 0.00188791 -0.01">
+                                            <inertial pos="0.0171394 0.000537591 4.8864e-07"
+                                                quat="0.575338 0.411667 -0.574906 0.411094"
+                                                mass="0.085445"
+                                                diaginertia="5.48211e-05 4.96646e-05 3.57798e-05" />
+                                            <joint name="left_wrist_roll_joint" pos="0 0 0"
+                                                axis="1 0 0" range="-1.97222 1.97222"
+                                                actuatorfrcrange="-25 25" />
+                                            <geom type="mesh" contype="0" conaffinity="0" group="1"
+                                                density="0" rgba="0.7 0.7 0.7 1"
+                                                mesh="left_wrist_roll_link" />
+                                            <geom type="mesh" rgba="0.7 0.7 0.7 1"
+                                                mesh="left_wrist_roll_link" />
+                                            <body name="left_wrist_pitch_link" pos="0.038 0 0">
+                                                <inertial pos="0.0229999 -0.00111685 -0.00111658"
+                                                    quat="0.249998 0.661363 0.293036 0.643608"
+                                                    mass="0.48405"
+                                                    diaginertia="0.000430353 0.000429873 0.000164648" />
+                                                <joint name="left_wrist_pitch_joint" pos="0 0 0"
+                                                    axis="0 1 0" range="-1.61443 1.61443"
+                                                    actuatorfrcrange="-5 5" />
+                                                <geom type="mesh" contype="0" conaffinity="0"
+                                                    group="1" density="0" rgba="0.7 0.7 0.7 1"
+                                                    mesh="left_wrist_pitch_link" />
+                                                <geom type="mesh" rgba="0.7 0.7 0.7 1"
+                                                    mesh="left_wrist_pitch_link" />
+                                                <body name="left_wrist_yaw_link" pos="0.046 0 0">
+                                                    <inertial pos="0.0708244 0.000191745 0.00161742"
+                                                        quat="0.510571 0.526295 0.468078 0.493188"
+                                                        mass="0.254576"
+                                                        diaginertia="0.000646113 0.000559993 0.000147566" />
+                                                    <joint name="left_wrist_yaw_joint" pos="0 0 0"
+                                                        axis="0 0 1" range="-1.61443 1.61443"
+                                                        actuatorfrcrange="-5 5" />
+                                                    <geom type="mesh" contype="0" conaffinity="0"
+                                                        group="1" density="0" rgba="0.7 0.7 0.7 1"
+                                                        mesh="left_wrist_yaw_link" />
+                                                    <geom type="mesh" rgba="0.7 0.7 0.7 1"
+                                                        mesh="left_wrist_yaw_link" />
+                                                    <geom pos="0.0415 0.003 0" quat="1 0 0 0"
+                                                        type="mesh" contype="0" conaffinity="0"
+                                                        group="1" density="0" rgba="0.7 0.7 0.7 1"
+                                                        mesh="left_rubber_hand" />
+                                                </body>
+                                            </body>
+                                        </body>
+                                    </body>
+                                </body>
+                            </body>
+                        </body>
+                        <body name="right_shoulder_pitch_link" pos="0.0039563 -0.10021 0.24778"
+                            quat="0.990264 -0.139201 1.38722e-05 9.86868e-05">
+                            <inertial pos="0 -0.035892 -0.011628"
+                                quat="0.68225 -0.326267 0.0130458 0.654152" mass="0.718"
+                                diaginertia="0.000465864 0.000432842 0.000406394" />
+                            <joint name="right_shoulder_pitch_joint" pos="0 0 0" axis="0 1 0"
+                                range="-3.0892 2.6704" actuatorfrcrange="-25 25" />
+                            <geom type="mesh" contype="0" conaffinity="0" group="1" density="0"
+                                rgba="0.7 0.7 0.7 1" mesh="right_shoulder_pitch_link" />
+                            <geom size="0.03 0.025" pos="0 -0.04 -0.01" quat="0.707107 0 0.707107 0"
+                                type="cylinder" rgba="0.7 0.7 0.7 1" />
+                            <body name="right_shoulder_roll_link" pos="0 -0.038 -0.013831"
+                                quat="0.990268 0.139172 0 0">
+                                <inertial pos="-0.000227 -0.00727 -0.063243"
+                                    quat="0.712604 -0.00710317 -0.0196223 0.701256" mass="0.643"
+                                    diaginertia="0.000691311 0.000618011 0.000388977" />
+                                <joint name="right_shoulder_roll_joint" pos="0 0 0" axis="1 0 0"
+                                    range="-2.2515 1.5882" actuatorfrcrange="-25 25" />
+                                <geom type="mesh" contype="0" conaffinity="0" group="1" density="0"
+                                    rgba="0.7 0.7 0.7 1" mesh="right_shoulder_roll_link" />
+                                <geom size="0.03 0.015" pos="-0.004 -0.006 -0.053" type="cylinder"
+                                    rgba="0.7 0.7 0.7 1" />
+                                <body name="right_shoulder_yaw_link" pos="0 -0.00624 -0.1032">
+                                    <inertial pos="0.010773 0.002949 -0.072009"
+                                        quat="0.687134 -0.0679942 -0.0964829 0.716879" mass="0.734"
+                                        diaginertia="0.00106187 0.00103217 0.000400661" />
+                                    <joint name="right_shoulder_yaw_joint" pos="0 0 0" axis="0 0 1"
+                                        range="-2.618 2.618" actuatorfrcrange="-25 25" />
+                                    <geom type="mesh" contype="0" conaffinity="0" group="1"
+                                        density="0" rgba="0.7 0.7 0.7 1"
+                                        mesh="right_shoulder_yaw_link" />
+                                    <geom type="mesh" rgba="0.7 0.7 0.7 1"
+                                        mesh="right_shoulder_yaw_link" />
+                                    <body name="right_elbow_link" pos="0.015783 0 -0.080518">
+                                        <inertial pos="0.064956 -0.004454 -0.010062"
+                                            quat="0.388129 0.388821 0.636132 0.541765" mass="0.6"
+                                            diaginertia="0.000443035 0.000421612 0.000259353" />
+                                        <joint name="right_elbow_joint" pos="0 0 0" axis="0 1 0"
+                                            range="-1.0472 2.0944" actuatorfrcrange="-25 25" />
+                                        <geom type="mesh" contype="0" conaffinity="0" group="1"
+                                            density="0" rgba="0.7 0.7 0.7 1" mesh="right_elbow_link" />
+                                        <geom type="mesh" rgba="0.7 0.7 0.7 1"
+                                            mesh="right_elbow_link" />
+                                        <body name="right_wrist_roll_link"
+                                            pos="0.1 -0.00188791 -0.01">
+                                            <inertial pos="0.0171394 -0.000537591 4.8864e-07"
+                                                quat="0.411667 0.575338 -0.411094 0.574906"
+                                                mass="0.085445"
+                                                diaginertia="5.48211e-05 4.96646e-05 3.57798e-05" />
+                                            <joint name="right_wrist_roll_joint" pos="0 0 0"
+                                                axis="1 0 0" range="-1.97222 1.97222"
+                                                actuatorfrcrange="-25 25" />
+                                            <geom type="mesh" contype="0" conaffinity="0" group="1"
+                                                density="0" rgba="0.7 0.7 0.7 1"
+                                                mesh="right_wrist_roll_link" />
+                                            <geom type="mesh" rgba="0.7 0.7 0.7 1"
+                                                mesh="right_wrist_roll_link" />
+                                            <body name="right_wrist_pitch_link" pos="0.038 0 0">
+                                                <inertial pos="0.0229999 0.00111685 -0.00111658"
+                                                    quat="0.643608 0.293036 0.661363 0.249998"
+                                                    mass="0.48405"
+                                                    diaginertia="0.000430353 0.000429873 0.000164648" />
+                                                <joint name="right_wrist_pitch_joint" pos="0 0 0"
+                                                    axis="0 1 0" range="-1.61443 1.61443"
+                                                    actuatorfrcrange="-5 5" />
+                                                <geom type="mesh" contype="0" conaffinity="0"
+                                                    group="1" density="0" rgba="0.7 0.7 0.7 1"
+                                                    mesh="right_wrist_pitch_link" />
+                                                <geom type="mesh" rgba="0.7 0.7 0.7 1"
+                                                    mesh="right_wrist_pitch_link" />
+                                                <body name="right_wrist_yaw_link" pos="0.046 0 0">
+                                                    <inertial
+                                                        pos="0.0708244 -0.000191745 0.00161742"
+                                                        quat="0.493188 0.468078 0.526295 0.510571"
+                                                        mass="0.254576"
+                                                        diaginertia="0.000646113 0.000559993 0.000147566" />
+                                                    <joint name="right_wrist_yaw_joint" pos="0 0 0"
+                                                        axis="0 0 1" range="-1.61443 1.61443"
+                                                        actuatorfrcrange="-5 5" />
+                                                    <geom type="mesh" contype="0" conaffinity="0"
+                                                        group="1" density="0" rgba="0.7 0.7 0.7 1"
+                                                        mesh="right_wrist_yaw_link" />
+                                                    <geom type="mesh" rgba="0.7 0.7 0.7 1"
+                                                        mesh="right_wrist_yaw_link" />
+                                                    <geom pos="0.0415 -0.003 0" quat="1 0 0 0"
+                                                        type="mesh" contype="0" conaffinity="0"
+                                                        group="1" density="0" rgba="0.7 0.7 0.7 1"
+                                                        mesh="right_rubber_hand" />
+                                                </body>
+                                            </body>
+                                        </body>
+                                    </body>
+                                </body>
+                            </body>
+                        </body>
+                    </body>
+                </body>
+            </body>
+        </body>
+    </worldbody>
+
+    <actuator>
+        <motor name="left_hip_pitch_joint" joint="left_hip_pitch_joint" />
+        <motor name="left_hip_roll_joint" joint="left_hip_roll_joint" />
+        <motor name="left_hip_yaw_joint" joint="left_hip_yaw_joint" />
+        <motor name="left_knee_joint" joint="left_knee_joint" />
+        <motor name="left_ankle_pitch_joint" joint="left_ankle_pitch_joint" />
+        <motor name="left_ankle_roll_joint" joint="left_ankle_roll_joint" />
+        <motor name="right_hip_pitch_joint" joint="right_hip_pitch_joint" />
+        <motor name="right_hip_roll_joint" joint="right_hip_roll_joint" />
+        <motor name="right_hip_yaw_joint" joint="right_hip_yaw_joint" />
+        <motor name="right_knee_joint" joint="right_knee_joint" />
+        <motor name="right_ankle_pitch_joint" joint="right_ankle_pitch_joint" />
+        <motor name="right_ankle_roll_joint" joint="right_ankle_roll_joint" />
+        <motor name="waist_yaw_joint" joint="waist_yaw_joint" />
+        <motor name="waist_roll_joint" joint="waist_roll_joint" />
+        <motor name="waist_pitch_joint" joint="waist_pitch_joint" />
+        <motor name="left_shoulder_pitch_joint" joint="left_shoulder_pitch_joint" />
+        <motor name="left_shoulder_roll_joint" joint="left_shoulder_roll_joint" />
+        <motor name="left_shoulder_yaw_joint" joint="left_shoulder_yaw_joint" />
+        <motor name="left_elbow_joint" joint="left_elbow_joint" />
+        <motor name="left_wrist_roll_joint" joint="left_wrist_roll_joint" />
+        <motor name="left_wrist_pitch_joint" joint="left_wrist_pitch_joint" />
+        <motor name="left_wrist_yaw_joint" joint="left_wrist_yaw_joint" />
+        <motor name="right_shoulder_pitch_joint" joint="right_shoulder_pitch_joint" />
+        <motor name="right_shoulder_roll_joint" joint="right_shoulder_roll_joint" />
+        <motor name="right_shoulder_yaw_joint" joint="right_shoulder_yaw_joint" />
+        <motor name="right_elbow_joint" joint="right_elbow_joint" />
+        <motor name="right_wrist_roll_joint" joint="right_wrist_roll_joint" />
+        <motor name="right_wrist_pitch_joint" joint="right_wrist_pitch_joint" />
+        <motor name="right_wrist_yaw_joint" joint="right_wrist_yaw_joint" />
+    </actuator>
+
+    <sensor>
+        <gyro name="imu-torso-angular-velocity" site="imu_in_torso" noise="5e-4" cutoff="34.9" />
+        <accelerometer name="imu-torso-linear-acceleration" site="imu_in_torso" noise="1e-2"
+            cutoff="157" />
+        <gyro name="imu-pelvis-angular-velocity" site="imu_in_pelvis" noise="5e-4" cutoff="34.9" />
+        <accelerometer name="imu-pelvis-linear-acceleration" site="imu_in_pelvis" noise="1e-2"
+            cutoff="157" />
+    </sensor>
+
+
+    <!-- setup scene -->
+    <statistic center="1.0 0.7 1.0" extent="0.8" />
+    <visual>
+        <headlight diffuse="0.6 0.6 0.6" ambient="0.1 0.1 0.1" specular="0.9 0.9 0.9" />
+        <rgba haze="0.15 0.25 0.35 1" />
+        <global azimuth="-140" elevation="-20" />
+    </visual>
+    <asset>
+        <texture type="skybox" builtin="flat" rgb1="0 0 0" rgb2="0 0 0" width="512" height="3072" />
+        <texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4"
+            rgb2="0.1 0.2 0.3" markrgb="0.8 0.8 0.8" width="300" height="300" />
+        <material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5"
+            reflectance="0.2" />
+    </asset>
+    <worldbody>
+        <light pos="1 0 3.5" dir="0 0 -1" directional="true" />
+        <geom name="floor" size="0 0 0.05" type="plane" material="groundplane" />
+    </worldbody>
+</mujoco>
\ No newline at end of file
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/mjcf/h2.xml b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/mjcf/h2.xml
new file mode 100644
index 0000000000000000000000000000000000000000..b431d759ecb0b2bbb69329a098c4b1eb39bba218
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/mjcf/h2.xml
@@ -0,0 +1,270 @@
+<mujoco model="H1031">
+  <compiler angle="radian" meshdir="meshes/"/>
+
+  <asset>
+    <mesh name="pelvis" content_type="model/stl" file="pelvis.stl"/>
+    <mesh name="pelvis_collision_link" content_type="model/stl" file="pelvis_collision.stl"/>
+    <mesh name="left_hip_pitch_link" content_type="model/stl" file="left_hip_pitch_link.stl"/>
+    <mesh name="left_hip_roll_link" content_type="model/stl" file="left_hip_roll_link.stl"/>
+    <mesh name="left_hip_yaw_link" content_type="model/stl" file="left_hip_yaw_link.stl"/>
+    <mesh name="left_knee_link" content_type="model/stl" file="left_knee_link.stl"/>
+    <mesh name="left_ankle_roll_link" content_type="model/stl" file="left_ankle_roll_link.stl"/>
+    <mesh name="left_ankle_pitch_link" content_type="model/stl" file="left_ankle_pitch_link.stl"/>
+    <mesh name="right_hip_pitch_link" content_type="model/stl" file="right_hip_pitch_link.stl"/>
+    <mesh name="right_hip_roll_link" content_type="model/stl" file="right_hip_roll_link.stl"/>
+    <mesh name="right_hip_yaw_link" content_type="model/stl" file="right_hip_yaw_link.stl"/>
+    <mesh name="right_knee_link" content_type="model/stl" file="right_knee_link.stl"/>
+    <mesh name="right_ankle_roll_link" content_type="model/stl" file="right_ankle_roll_link.stl"/>
+    <mesh name="right_ankle_pitch_link" content_type="model/stl" file="right_ankle_pitch_link.stl"/>
+    <mesh name="waist_yaw_link" content_type="model/stl" file="waist_yaw_link.stl"/>
+    <mesh name="waist_roll_link" content_type="model/stl" file="waist_roll_link.stl"/>
+    <mesh name="torso_link" content_type="model/stl" file="torso_link.stl"/>
+    <mesh name="torso_collision" content_type="model/stl" file="torso_collision.stl"/>
+    <mesh name="head_pitch_link" content_type="model/stl" file="head_pitch_link.stl"/>
+    <mesh name="head_yaw_link" content_type="model/stl" file="head_yaw_link.stl"/>
+    <mesh name="left_shoulder_pitch_link" content_type="model/stl" file="left_shoulder_pitch_link.stl"/>
+    <mesh name="left_shoulder_roll_link" content_type="model/stl" file="left_shoulder_roll_link.stl"/>
+    <mesh name="left_shoulder_yaw_link" content_type="model/stl" file="left_shoulder_yaw_link.stl"/>
+    <mesh name="left_elbow_link" content_type="model/stl" file="left_elbow_link.stl"/>
+    <mesh name="left_wrist_roll_link" content_type="model/stl" file="left_wrist_roll_link.stl"/>
+    <mesh name="left_wrist_pitch_link" content_type="model/stl" file="left_wrist_pitch_link.stl"/>
+    <mesh name="left_rubber_hand" content_type="model/stl" file="left_rubber_hand.stl"/>
+    <mesh name="right_shoulder_pitch_link" content_type="model/stl" file="right_shoulder_pitch_link.stl"/>
+    <mesh name="right_shoulder_roll_link" content_type="model/stl" file="right_shoulder_roll_link.stl"/>
+    <mesh name="right_shoulder_yaw_link" content_type="model/stl" file="right_shoulder_yaw_link.stl"/>
+    <mesh name="right_elbow_link" content_type="model/stl" file="right_elbow_link.stl"/>
+    <mesh name="right_wrist_roll_link" content_type="model/stl" file="right_wrist_roll_link.stl"/>
+    <mesh name="right_wrist_pitch_link" content_type="model/stl" file="right_wrist_pitch_link.stl"/>
+    <mesh name="right_rubber_hand" content_type="model/stl" file="right_rubber_hand.stl"/>
+  </asset>
+
+  <worldbody>
+    <body name="pelvis" pos="0 0 1.04">
+      <site name="imu" size="0.01" pos="0.0 0.0 0.0"/>
+      <inertial pos="0.00046244 -0.00004214 -0.01473435" quat="1 0 0 0" mass="7.60468095" diaginertia="0.03548205 0.02290677 0.02753519"/>
+      <joint name="floating_base_joint" type="free" limited="false" actuatorfrclimited="false"/>
+      <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.2 0.2 0.2 1" mesh="pelvis"/>
+      <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="pelvis_collision_link"/>
+      <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="pelvis_collision_link"/>
+      <body name="left_hip_pitch_link" pos="0 0.083873 -0.048424" quat="0.965926 -0.25882 0 0">
+        <inertial pos="0.00300595 0.0593147 -0.0170915" quat="0.899485 0.435934 -0.0222136 -0.0198741" mass="2.79753" diaginertia="0.00510297 0.00410085 0.0034593"/>
+        <joint name="left_hip_pitch_joint" pos="0 0 0" axis="0 1 0" range="-2.4526 2.77542" actuatorfrcrange="-360 360"/>
+        <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="left_hip_pitch_link"/>
+        <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="left_hip_pitch_link"/>
+        <body name="left_hip_roll_link" pos="0.041851 0.0635 -0.019" quat="0.965926 0.25882 0 0">
+          <inertial pos="-0.0302619 -0.0046893 -0.0444985" quat="-0.0935812 0.633472 -0.0477755 0.766598" mass="0.627146" diaginertia="0.00200252 0.00167509 0.00128023"/>
+          <joint name="left_hip_roll_joint" pos="0 0 0" axis="1 0 0" range="-0.467441 2.16886" actuatorfrcrange="-360 360"/>
+          <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="left_hip_roll_link"/>
+          <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="left_hip_roll_link"/>
+          <body name="left_hip_yaw_link" pos="-0.028187 -0.0095 -0.073">
+            <inertial pos="0.00069408 -0.0108659 -0.113015" quat="0.621993 -0.0405952 0.0198637 0.781717" mass="7.08752" diaginertia="0.0514012 0.050596 0.0130067"/>
+            <joint name="left_hip_yaw_joint" pos="0 0 0" axis="0 0 1" range="-2.827 2.827" actuatorfrcrange="-360 360"/>
+            <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="left_hip_yaw_link"/>
+            <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="left_hip_yaw_link"/>
+            <body name="left_knee_link" pos="-0.010341 0.018489 -0.3191">
+              <inertial pos="0.00461331 -0.0245322 -0.167155" quat="0.689662 0.00325283 0.00290593 0.724118" mass="3.5534" diaginertia="0.0512388 0.0510884 0.00333998"/>
+              <joint name="left_knee_joint" pos="0 0 0" axis="0 1 0" range="-0.08725 2.53025" actuatorfrcrange="-360 360"/>
+              <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="left_knee_link"/>
+              <geom size="0.04 0.175" pos="0.01 -0.025 -0.25" type="cylinder" rgba="0.7 0.7 0.7 1"/>
+              <body name="left_ankle_roll_link" pos="0.03272 -0.023002 -0.49662">
+                <inertial pos="-0.0174063 6.424e-05 -0.0002146" quat="0.499157 0.502573 -0.497809 0.500449" mass="0.123182" diaginertia="4.00501e-05 3.43207e-05 1.44993e-05"/>
+                <joint name="left_ankle_roll_joint" pos="0 0 0" axis="1 0 0" range="-0.349066 0.296706" actuatorfrcrange="-19 19"/>
+                <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="left_ankle_roll_link"/>
+                <body name="left_ankle_pitch_link" pos="-0.0225 -0.0373 0">
+                  <inertial pos="0.032059 0.0370053 -0.0288282" quat="6.19452e-05 0.72345 -3.0004e-05 0.690377" mass="0.844967" diaginertia="0.00438317 0.00410964 0.000599978"/>
+                  <joint name="left_ankle_pitch_joint" pos="0 0 0" axis="0 1 0" range="-1.13446 0.610865" actuatorfrcrange="-66.88 66.88"/>
+                  <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.4 0.4 0.4 1" mesh="left_ankle_pitch_link"/>
+                  <geom type="mesh" rgba="0.4 0.4 0.4 1" mesh="left_ankle_pitch_link"/>
+                </body>
+              </body>
+            </body>
+          </body>
+        </body>
+      </body>
+      <body name="right_hip_pitch_link" pos="0 -0.083873 -0.048424" quat="0.965926 0.25882 0 0">
+        <inertial pos="0.00300595 -0.0593147 -0.0170915" quat="0.435934 0.899485 0.0198741 0.0222136" mass="2.79753" diaginertia="0.00510297 0.00410085 0.0034593"/>
+        <joint name="right_hip_pitch_joint" pos="0 0 0" axis="0 1 0" range="-2.4526 2.77542" actuatorfrcrange="-360 360"/>
+        <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="right_hip_pitch_link"/>
+        <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="right_hip_pitch_link"/>
+        <body name="right_hip_roll_link" pos="0.041851 -0.0635 -0.019" quat="0.965926 -0.25882 0 0">
+          <inertial pos="-0.0302619 0.0046893 -0.0444985" quat="0.0935812 0.633472 0.0477755 0.766598" mass="0.627146" diaginertia="0.00200252 0.00167509 0.00128023"/>
+          <joint name="right_hip_roll_joint" pos="0 0 0" axis="1 0 0" range="-2.16886 0.467441" actuatorfrcrange="-360 360"/>
+          <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="right_hip_roll_link"/>
+          <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="right_hip_roll_link"/>
+          <body name="right_hip_yaw_link" pos="-0.028187 0.0095 -0.073">
+            <inertial pos="0.00069408 0.0108659 -0.113015" quat="0.781717 0.0198637 -0.0405952 0.621993" mass="7.08752" diaginertia="0.0514012 0.050596 0.0130067"/>
+            <joint name="right_hip_yaw_joint" pos="0 0 0" axis="0 0 1" range="-2.827 2.827" actuatorfrcrange="-360 360"/>
+            <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="right_hip_yaw_link"/>
+            <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="right_hip_yaw_link"/>
+            <body name="right_knee_link" pos="-0.010341 -0.018489 -0.3191">
+              <inertial pos="0.00461331 0.0245322 -0.167155" quat="0.724118 0.00290593 0.00325283 0.689662" mass="3.5534" diaginertia="0.0512388 0.0510884 0.00333998"/>
+              <joint name="right_knee_joint" pos="0 0 0" axis="0 1 0" range="-0.08725 2.53025" actuatorfrcrange="-360 360"/>
+              <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="right_knee_link"/>
+              <geom size="0.04 0.175" pos="0.01 0.025 -0.25" type="cylinder" rgba="0.7 0.7 0.7 1"/>
+              <body name="right_ankle_roll_link" pos="0.03272 0.023002 -0.49662">
+                <inertial pos="-0.0174063 -6.424e-05 -0.0002146" quat="0.502573 0.499157 -0.500449 0.497809" mass="0.123182" diaginertia="4.00501e-05 3.43207e-05 1.44993e-05"/>
+                <joint name="right_ankle_roll_joint" pos="0 0 0" axis="1 0 0" range="-0.296706 0.349066" actuatorfrcrange="-19 19"/>
+                <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="right_ankle_roll_link"/>
+                <body name="right_ankle_pitch_link" pos="-0.0225 0.0373 0">
+                  <inertial pos="0.032059 -0.0370053 -0.0288282" quat="-6.19452e-05 0.72345 3.0004e-05 0.690377" mass="0.844967" diaginertia="0.00438317 0.00410964 0.000599978"/>
+                  <joint name="right_ankle_pitch_joint" pos="0 0 0" axis="0 1 0" range="-1.13446 0.610865" actuatorfrcrange="-66.88 66.88"/>
+                  <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.4 0.4 0.4 1" mesh="right_ankle_pitch_link"/>
+                  <geom type="mesh" rgba="0.4 0.4 0.4 1" mesh="right_ankle_pitch_link"/>
+                </body>
+              </body>
+            </body>
+          </body>
+        </body>
+      </body>
+      <body name="waist_yaw_link" pos="0 0 0.089837">
+        <inertial pos="0.00626994 0 0.0169696" quat="0 0.717728 0 0.696323" mass="0.279762" diaginertia="0.000311422 0.0002473 0.000123348"/>
+        <joint name="waist_yaw_joint" pos="0 0 0" axis="0 0 1" range="-1.7453 1.7453" actuatorfrcrange="-120 120"/>
+        <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="waist_yaw_link"/>
+        <body name="waist_roll_link">
+          <inertial pos="0 0.00128474 0" quat="0 0.707107 0 0.707107" mass="0.0658504" diaginertia="1.58e-05 1.041e-05 8.09e-06"/>
+          <joint name="waist_roll_joint" pos="0 0 0" axis="1 0 0" range="-0.5236 0.5236" actuatorfrcrange="-180 180"/>
+          <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="waist_roll_link"/>
+          <body name="torso_link">
+            <inertial pos="0.0134247 0.0054719 0.210501" quat="0.999399 -0.0341887 0.00399831 0.00403179" mass="15.8703" diaginertia="0.273064 0.244383 0.0840858"/>
+            <joint name="waist_pitch_joint" pos="0 0 0" axis="0 1 0" range="-0.43633 0.5236" actuatorfrcrange="-180 180"/>
+            <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.4 0.4 0.4 1" mesh="torso_link"/>
+            <geom type="mesh" rgba="0.4 0.4 0.4 1" mesh="torso_collision"/>
+            <body name="head_pitch_link" pos="-0.0002283 2.7112e-05 0.42181">
+              <inertial pos="5.279e-08 -0.00035318 0.045125" quat="0.999996 -0.000955249 -0.00262397 -0.000580497" mass="1.82" diaginertia="0.00538697 0.00532442 0.00140982"/>
+              <joint name="head_pitch_joint" pos="0 0 0" axis="0 1 0" range="-0.5236 0.83775" actuatorfrcrange="-50 50"/>
+              <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="head_pitch_link"/>
+              <body name="head_yaw_link" pos="0 0.0003 0.1325">
+                <inertial pos="0.0219 0 0.0187" quat="0.667311 -0.227767 -0.231751 0.670156" mass="1.3" diaginertia="0.00535274 0.00472336 0.00378091"/>
+                <joint name="head_yaw_joint" pos="0 0 0" axis="0 0 1" range="-1.7453 1.7453" actuatorfrcrange="-50 50"/>
+                <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="head_yaw_link"/>
+                <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="head_yaw_link"/>
+              </body>
+            </body>
+            <body name="left_shoulder_pitch_link" pos="0.0015 0.11254 0.35812" quat="0.990959 0.134164 0 0">
+              <inertial pos="0.00296153 0.0520268 -0.025045" quat="0.955852 0.29138 -0.0379121 0.00288293" mass="1.06056" diaginertia="0.00122192 0.00110581 0.000458834"/>
+              <joint name="left_shoulder_pitch_joint" pos="0 0 0" axis="0 1 0" range="-2.61799 1.8326" actuatorfrcrange="-120 120"/>
+              <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="left_shoulder_pitch_link"/>
+              <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="left_shoulder_pitch_link"/>
+              <body name="left_shoulder_roll_link" pos="0.027682 0.061638 -0.039909" quat="0.990959 -0.134164 0 0">
+                <inertial pos="-0.02061 0.00415443 -0.0629422" quat="0.678366 0.110828 0.0778018 0.722138" mass="1.34511" diaginertia="0.00210225 0.00198513 0.00105666"/>
+                <joint name="left_shoulder_roll_joint" pos="0 0 0" axis="1 0 0" range="-0.516617 2.63545" actuatorfrcrange="-54 54"/>
+                <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="left_shoulder_roll_link"/>
+                <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="left_shoulder_roll_link"/>
+                <body name="left_shoulder_yaw_link" pos="-0.030632 0.0037809 -0.10422">
+                  <inertial pos="0.00364187 0.0186515 -0.0558555" quat="0.995617 -0.0729116 -0.0584518 -0.00371547" mass="0.915717" diaginertia="0.00306267 0.00257349 0.00143042"/>
+                  <joint name="left_shoulder_yaw_joint" pos="0 0 0" axis="0 0 1" range="-2.61799 2.61799" actuatorfrcrange="-54 54"/>
+                  <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="left_shoulder_yaw_link"/>
+                  <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="left_shoulder_yaw_link"/>
+                  <body name="left_elbow_link" pos="0.016955 0.031006 -0.12719">
+                    <inertial pos="0.0459969 -0.0283296 -0.00661325" quat="0.520177 0.507229 0.507638 0.463074" mass="1.69432" diaginertia="0.00548751 0.00533836 0.000695052"/>
+                    <joint name="left_elbow_joint" pos="0 0 0" axis="0 1 0" range="-0.986111 3.07178" actuatorfrcrange="-54 54"/>
+                    <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="left_elbow_link"/>
+                    <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="left_elbow_link"/>
+                    <body name="left_wrist_roll_link" pos="0.12826 -0.030815 -0.013299">
+                      <inertial pos="0.0337195 0.00237482 0.00031544" quat="-0.242945 0.638599 0.286806 0.671499" mass="0.969837" diaginertia="0.000926999 0.000903832 0.000852829"/>
+                      <joint name="left_wrist_roll_joint" pos="0 0 0" axis="1 0 0" range="-2.61799 2.61799" actuatorfrcrange="-54 54"/>
+                      <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="left_wrist_roll_link"/>
+                      <body name="left_wrist_pitch_link" pos="0.082003 -0.031003 0.00020786">
+                        <inertial pos="0.01821 0.0314035 0.0115687" quat="0.977698 -0.037912 -0.203807 0.033629" mass="0.946953" diaginertia="0.000918192 0.00082104 0.000648419"/>
+                        <joint name="left_wrist_pitch_joint" pos="0 0 0" axis="0 1 0" range="-0.436332 0.436332" actuatorfrcrange="-25 25"/>
+                        <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="left_wrist_pitch_link"/>
+                        <body name="left_wrist_yaw_link" pos="0.026025 0.030996 0.026981">
+                          <inertial pos="0.00932726 0.00247855 -0.0196166" quat="0.528225 0.503579 0.545204 0.41248" mass="0.775752" diaginertia="0.0021247 0.00196098 0.000336621"/>
+                          <joint name="left_wrist_yaw_joint" pos="0 0 0" axis="0 0 1" range="-1.22173 1.22173" actuatorfrcrange="-25 25"/>
+                          <geom size="0.028 0.03" pos="0 0 -0.029" type="cylinder" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1"/>
+                          <geom size="0.04" pos="0 0 -0.029" rgba="0.7 0.7 0.7 1"/>
+                          <geom pos="0.025 0 -0.03" quat="1 0 0 0" type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.4 0.4 0.4 1" mesh="left_rubber_hand"/>
+                        </body>
+                      </body>
+                    </body>
+                  </body>
+                </body>
+              </body>
+            </body>
+            <body name="right_shoulder_pitch_link" pos="0.0015 -0.1125 0.35812" quat="0.990959 -0.134164 0 0">
+              <inertial pos="0.00296153 -0.0520268 -0.025045" quat="0.955852 -0.291379 -0.0379149 -0.00289698" mass="1.06056" diaginertia="0.00122192 0.00110581 0.000458835"/>
+              <joint name="right_shoulder_pitch_joint" pos="0 0 0" axis="0 1 0" range="-2.61799 1.8326" actuatorfrcrange="-120 120"/>
+              <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="right_shoulder_pitch_link"/>
+              <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="right_shoulder_pitch_link"/>
+              <body name="right_shoulder_roll_link" pos="0.027682 -0.061638 -0.039909" quat="0.990959 0.134164 0 0">
+                <inertial pos="-0.02061 -0.00415443 -0.0629422" quat="0.722138 0.0778018 0.110828 0.678366" mass="1.34511" diaginertia="0.00210225 0.00198513 0.00105666"/>
+                <joint name="right_shoulder_roll_joint" pos="0 0 0" axis="1 0 0" range="-2.63545 0.516617" actuatorfrcrange="-54 54"/>
+                <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="right_shoulder_roll_link"/>
+                <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="right_shoulder_roll_link"/>
+                <body name="right_shoulder_yaw_link" pos="-0.030632 -0.0037809 -0.10422">
+                  <inertial pos="0.00364187 -0.0186515 -0.0558555" quat="0.995617 0.0729116 -0.0584518 0.00371547" mass="0.915717" diaginertia="0.00306267 0.00257349 0.00143042"/>
+                  <joint name="right_shoulder_yaw_joint" pos="0 0 0" axis="0 0 1" range="-2.61799 2.61799" actuatorfrcrange="-54 54"/>
+                  <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="right_shoulder_yaw_link"/>
+                  <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="right_shoulder_yaw_link"/>
+                  <body name="right_elbow_link" pos="0.016955 -0.031006 -0.12719">
+                    <inertial pos="0.0459969 0.0283296 -0.00661325" quat="0.463074 0.507638 0.507229 0.520177" mass="1.69432" diaginertia="0.00548751 0.00533836 0.000695052"/>
+                    <joint name="right_elbow_joint" pos="0 0 0" axis="0 1 0" range="-0.986111 3.07178" actuatorfrcrange="-54 54"/>
+                    <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="right_elbow_link"/>
+                    <geom type="mesh" rgba="0.7 0.7 0.7 1" mesh="right_elbow_link"/>
+                    <body name="right_wrist_roll_link" pos="0.12826 0.030815 -0.013299">
+                      <inertial pos="0.0337195 -0.00237482 0.00031544" quat="0.242945 0.638599 -0.286806 0.671499" mass="0.969837" diaginertia="0.000926999 0.000903832 0.000852829"/>
+                      <joint name="right_wrist_roll_joint" pos="0 0 0" axis="1 0 0" range="-2.61799 2.61799" actuatorfrcrange="-54 54"/>
+                      <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="right_wrist_roll_link"/>
+                      <body name="right_wrist_pitch_link" pos="0.082003 0.031003 0.00020786">
+                        <inertial pos="0.01821 -0.0314035 0.0115687" quat="0.977698 0.037912 -0.203807 -0.033629" mass="0.946953" diaginertia="0.000918192 0.00082104 0.000648419"/>
+                        <joint name="right_wrist_pitch_joint" pos="0 0 0" axis="0 1 0" range="-0.436332 0.436332" actuatorfrcrange="-25 25"/>
+                        <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1" mesh="right_wrist_pitch_link"/>
+                        <body name="right_wrist_yaw_link" pos="0.026025 -0.030996 0.026981">
+                          <inertial pos="0.00932726 -0.00247855 -0.0196166" quat="0.41248 0.545204 0.503579 0.528225" mass="0.775752" diaginertia="0.0021247 0.00196098 0.000336621"/>
+                          <joint name="right_wrist_yaw_joint" pos="0 0 0" axis="0 0 1" range="-1.22173 1.22173" actuatorfrcrange="-25 25"/>
+                          <geom size="0.028 0.03" pos="0 0 -0.029" type="cylinder" contype="0" conaffinity="0" group="1" density="0" rgba="0.7 0.7 0.7 1"/>
+                          <geom size="0.04" pos="0 0 -0.029" rgba="0.7 0.7 0.7 1"/>
+                          <geom pos="0.025 0 -0.03" quat="1 0 0 0" type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.4 0.4 0.4 1" mesh="right_rubber_hand"/>
+                        </body>
+                      </body>
+                    </body>
+                  </body>
+                </body>
+              </body>
+            </body>
+          </body>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+
+
+  <actuator>
+    <motor name="left_hip_pitch" joint="left_hip_pitch_joint" gear="1"/>
+    <motor name="left_hip_roll" joint="left_hip_roll_joint" gear="1"/>
+    <motor name="left_hip_yaw" joint="left_hip_yaw_joint" gear="1"/>
+    <motor name="left_knee" joint="left_knee_joint" gear="1"/>
+    <motor name="left_ankle_roll" joint="left_ankle_roll_joint" gear="1"/>
+    <motor name="left_ankle_pitch" joint="left_ankle_pitch_joint" gear="1"/>
+
+    <motor name="right_hip_pitch" joint="right_hip_pitch_joint" gear="1"/>
+    <motor name="right_hip_roll" joint="right_hip_roll_joint" gear="1"/>
+    <motor name="right_hip_yaw" joint="right_hip_yaw_joint" gear="1"/>
+    <motor name="right_knee" joint="right_knee_joint" gear="1"/>
+    <motor name="right_ankle_roll" joint="right_ankle_roll_joint" gear="1"/>
+    <motor name="right_ankle_pitch" joint="right_ankle_pitch_joint" gear="1"/>
+
+    <motor name="waist_yaw" joint="waist_yaw_joint" gear="1"/>
+    <motor name="waist_roll" joint="waist_roll_joint" gear="1"/>
+    <motor name="waist_pitch" joint="waist_pitch_joint" gear="1"/>
+
+    <motor name="head_pitch" joint="head_pitch_joint" gear="1"/>
+    <motor name="head_yaw" joint="head_yaw_joint" gear="1"/>
+
+    <motor name="left_shoulder_pitch" joint="left_shoulder_pitch_joint" gear="1"/>
+    <motor name="left_shoulder_roll" joint="left_shoulder_roll_joint" gear="1"/>
+    <motor name="left_shoulder_yaw" joint="left_shoulder_yaw_joint" gear="1"/>
+    <motor name="left_elbow" joint="left_elbow_joint" gear="1"/>
+    <motor name="left_wrist_roll" joint="left_wrist_roll_joint" gear="1"/>
+    <motor name="left_wrist_pitch" joint="left_wrist_pitch_joint" gear="1"/>
+    <motor name="left_wrist_yaw" joint="left_wrist_yaw_joint" gear="1"/>
+
+    <motor name="right_shoulder_pitch" joint="right_shoulder_pitch_joint" gear="1"/>
+    <motor name="right_shoulder_roll" joint="right_shoulder_roll_joint" gear="1"/>
+    <motor name="right_shoulder_yaw" joint="right_shoulder_yaw_joint" gear="1"/>
+    <motor name="right_elbow" joint="right_elbow_joint" gear="1"/>
+    <motor name="right_wrist_roll" joint="right_wrist_roll_joint" gear="1"/>
+    <motor name="right_wrist_pitch" joint="right_wrist_pitch_joint" gear="1"/>
+    <motor name="right_wrist_yaw" joint="right_wrist_yaw_joint" gear="1"/>
+  </actuator>
+</mujoco>
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/g1/main.urdf b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/g1/main.urdf
new file mode 100644
index 0000000000000000000000000000000000000000..094bb600225d3697ca8adb5f8c2844b57700a730
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/g1/main.urdf
@@ -0,0 +1,1498 @@
+<?xml version="1.0"?>
+
+<!-- This URDF is adapted from the BeyondMimic project.
+     See: https://github.com/HybridRobotics/whole_body_tracking -->
+
+<robot name="g1">
+
+    <link name="pelvis">
+        <inertial>
+            <origin xyz="0 0 -0.07605" rpy="0 0 0" />
+            <mass value="3.813" />
+            <inertia ixx="0.010549" ixy="0" ixz="2.1E-06" iyy="0.0093089" iyz="0" izz="0.0079184" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/pelvis.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <origin xyz="0 0 -0.08" />
+            <geometry>
+                <sphere radius="0.07" />
+            </geometry>
+        </collision>
+    </link>
+    <link name="pelvis_contour_link">
+        <inertial>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <mass value="0.001" />
+            <inertia ixx="1e-7" ixy="0" ixz="0" iyy="1e-7" iyz="0" izz="1e-7" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/pelvis_contour_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+    </link>
+    <joint name="pelvis_contour_joint" type="fixed">
+        <parent link="pelvis" />
+        <child link="pelvis_contour_link" />
+    </joint>
+
+    <!-- Legs -->
+    <link name="left_hip_pitch_link">
+        <inertial>
+            <origin xyz="0.002741 0.047791 -0.02606" rpy="0 0 0" />
+            <mass value="1.35" />
+            <inertia ixx="0.001811" ixy="3.68E-05" ixz="-3.44E-05" iyy="0.0014193" iyz="0.000171"
+                izz="0.0012812" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hip_pitch_link.STL" />
+            </geometry>
+            <material name="dark">
+                <color rgba="0.2 0.2 0.2 1" />
+            </material>
+        </visual>
+    </link>
+    <joint name="left_hip_pitch_joint" type="revolute">
+        <origin xyz="0 0.064452 -0.1027" rpy="0 0 0" />
+        <parent link="pelvis" />
+        <child link="left_hip_pitch_link" />
+        <axis xyz="0 1 0" />
+        <limit lower="-2.5307" upper="2.8798" effort="88" velocity="32" />
+    </joint>
+    <link name="left_hip_roll_link">
+        <inertial>
+            <origin xyz="0.029812 -0.001045 -0.087934" rpy="0 0 0" />
+            <mass value="1.52" />
+            <inertia ixx="0.0023773" ixy="-3.8E-06" ixz="-0.0003908" iyy="0.0024123" iyz="1.84E-05"
+                izz="0.0016595" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hip_roll_link.STL" />
+            </geometry>
+            <material name="white">
+                <color rgba="0.7 0.7 0.7 1" />
+            </material>
+        </visual>
+        <collision>
+            <origin xyz="0.02 0 -0.1" />
+            <geometry>
+                <cylinder radius="0.05" length="0.2" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="left_hip_roll_joint" type="revolute">
+        <origin xyz="0 0.052 -0.030465" rpy="0 -0.1749 0" />
+        <parent link="left_hip_pitch_link" />
+        <child link="left_hip_roll_link" />
+        <axis xyz="1 0 0" />
+        <limit lower="-0.5236" upper="2.9671" effort="139" velocity="20" />
+    </joint>
+    <link name="left_hip_yaw_link">
+        <inertial>
+            <origin xyz="-0.057709 -0.010981 -0.15078" rpy="0 0 0" />
+            <mass value="1.702" />
+            <inertia ixx="0.0057774" ixy="-0.0005411" ixz="-0.0023948" iyy="0.0076124"
+                iyz="-0.0007072" izz="0.003149" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hip_yaw_link.STL" />
+            </geometry>
+            <material name="white">
+                <color rgba="0.7 0.7 0.7 1" />
+            </material>
+        </visual>
+    </link>
+    <joint name="left_hip_yaw_joint" type="revolute">
+        <origin xyz="0.025001 0 -0.12412" rpy="0 0 0" />
+        <parent link="left_hip_roll_link" />
+        <child link="left_hip_yaw_link" />
+        <axis xyz="0 0 1" />
+        <limit lower="-2.7576" upper="2.7576" effort="88" velocity="32" />
+    </joint>
+    <link name="left_knee_link">
+        <inertial>
+            <origin xyz="0.005457 0.003964 -0.12074" rpy="0 0 0" />
+            <mass value="1.932" />
+            <inertia ixx="0.011329" ixy="4.82E-05" ixz="-4.49E-05" iyy="0.011277" iyz="-0.0007146"
+                izz="0.0015168" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_knee_link.STL" />
+            </geometry>
+            <material name="white">
+                <color rgba="0.7 0.7 0.7 1" />
+            </material>
+        </visual>
+        <collision>
+            <origin xyz="0.02 0 -0.1" />
+            <geometry>
+                <cylinder radius="0.05" length="0.2" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="left_knee_joint" type="revolute">
+        <origin xyz="-0.078273 0.0021489 -0.17734" rpy="0 0.1749 0" />
+        <parent link="left_hip_yaw_link" />
+        <child link="left_knee_link" />
+        <axis xyz="0 1 0" />
+        <limit lower="-0.087267" upper="2.8798" effort="139" velocity="20" />
+    </joint>
+    <link name="left_ankle_pitch_link">
+        <inertial>
+            <origin xyz="-0.007269 0 0.011137" rpy="0 0 0" />
+            <mass value="0.074" />
+            <inertia ixx="8.4E-06" ixy="0" ixz="-2.9E-06" iyy="1.89E-05" iyz="0" izz="1.26E-05" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_ankle_pitch_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+    </link>
+    <joint name="left_ankle_pitch_joint" type="revolute">
+        <origin xyz="0 -9.4445E-05 -0.30001" rpy="0 0 0" />
+        <parent link="left_knee_link" />
+        <child link="left_ankle_pitch_link" />
+        <axis xyz="0 1 0" />
+        <limit lower="-0.87267" upper="0.5236" effort="50" velocity="37" />
+    </joint>
+    <link name="left_ankle_roll_link">
+        <inertial>
+            <origin xyz="0.026505 0 -0.016425" rpy="0 0 0" />
+            <mass value="0.608" />
+            <inertia ixx="0.0002231" ixy="2E-07" ixz="8.91E-05" iyy="0.0016161" iyz="-1E-07"
+                izz="0.0016667" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_ankle_roll_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <origin xyz="0.075 -0.026 -0.025" rpy="0 1.5708 0" />
+            <geometry>
+                <cylinder radius="0.01" length="0.05" />
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.0395 -0.018 -0.025" rpy="0 -1.5708 0" />
+            <geometry>
+                <cylinder radius="0.008" length="0.167" />
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.039 -0.01 -0.025" rpy="0 -1.5708 0" />
+            <geometry>
+                <cylinder radius="0.01" length="0.182" />
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.039 0 -0.025" rpy="0 -1.5708 0" />
+            <geometry>
+                <cylinder radius="0.01" length="0.186" />
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.039 0.01 -0.025" rpy="0 -1.5708 0" />
+            <geometry>
+                <cylinder radius="0.01" length="0.182" />
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.0395 0.018 -0.025" rpy="0 -1.5708 0" />
+            <geometry>
+                <cylinder radius="0.008" length="0.167" />
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.075 0.026 -0.025" rpy="0 1.5708 0" />
+            <geometry>
+                <cylinder radius="0.01" length="0.05" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="left_ankle_roll_joint" type="revolute">
+        <origin xyz="0 0 -0.017558" rpy="0 0 0" />
+        <parent link="left_ankle_pitch_link" />
+        <child link="left_ankle_roll_link" />
+        <axis xyz="1 0 0" />
+        <limit lower="-0.2618" upper="0.2618" effort="50" velocity="37" />
+    </joint>
+    <link name="right_hip_pitch_link">
+        <inertial>
+            <origin xyz="0.002741 -0.047791 -0.02606" rpy="0 0 0" />
+            <mass value="1.35" />
+            <inertia ixx="0.001811" ixy="-3.68E-05" ixz="-3.44E-05" iyy="0.0014193" iyz="-0.000171"
+                izz="0.0012812" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_hip_pitch_link.STL" />
+            </geometry>
+            <material name="dark">
+                <color rgba="0.2 0.2 0.2 1" />
+            </material>
+        </visual>
+    </link>
+    <joint name="right_hip_pitch_joint" type="revolute">
+        <origin xyz="0 -0.064452 -0.1027" rpy="0 0 0" />
+        <parent link="pelvis" />
+        <child link="right_hip_pitch_link" />
+        <axis xyz="0 1 0" />
+        <limit lower="-2.5307" upper="2.8798" effort="88" velocity="32" />
+    </joint>
+    <link name="right_hip_roll_link">
+        <inertial>
+            <origin xyz="0.029812 0.001045 -0.087934" rpy="0 0 0" />
+            <mass value="1.52" />
+            <inertia ixx="0.0023773" ixy="3.8E-06" ixz="-0.0003908" iyy="0.0024123" iyz="-1.84E-05"
+                izz="0.0016595" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_hip_roll_link.STL" />
+            </geometry>
+            <material name="white">
+                <color rgba="0.7 0.7 0.7 1" />
+            </material>
+        </visual>
+        <collision>
+            <origin xyz="0.02 0 -0.1" />
+            <geometry>
+                <cylinder radius="0.05" length="0.2" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="right_hip_roll_joint" type="revolute">
+        <origin xyz="0 -0.052 -0.030465" rpy="0 -0.1749 0" />
+        <parent link="right_hip_pitch_link" />
+        <child link="right_hip_roll_link" />
+        <axis xyz="1 0 0" />
+        <limit lower="-2.9671" upper="0.5236" effort="139" velocity="20" />
+    </joint>
+    <link name="right_hip_yaw_link">
+        <inertial>
+            <origin xyz="-0.057709 0.010981 -0.15078" rpy="0 0 0" />
+            <mass value="1.702" />
+            <inertia ixx="0.0057774" ixy="0.0005411" ixz="-0.0023948" iyy="0.0076124"
+                iyz="0.0007072" izz="0.003149" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_hip_yaw_link.STL" />
+            </geometry>
+            <material name="white">
+                <color rgba="0.7 0.7 0.7 1" />
+            </material>
+        </visual>
+    </link>
+    <joint name="right_hip_yaw_joint" type="revolute">
+        <origin xyz="0.025001 0 -0.12412" rpy="0 0 0" />
+        <parent link="right_hip_roll_link" />
+        <child link="right_hip_yaw_link" />
+        <axis xyz="0 0 1" />
+        <limit lower="-2.7576" upper="2.7576" effort="88" velocity="32" />
+    </joint>
+    <link name="right_knee_link">
+        <inertial>
+            <origin xyz="0.005457 -0.003964 -0.12074" rpy="0 0 0" />
+            <mass value="1.932" />
+            <inertia ixx="0.011329" ixy="-4.82E-05" ixz="4.49E-05" iyy="0.011277" iyz="0.0007146"
+                izz="0.0015168" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_knee_link.STL" />
+            </geometry>
+            <material name="white">
+                <color rgba="0.7 0.7 0.7 1" />
+            </material>
+        </visual>
+        <collision>
+            <origin xyz="0.02 0 -0.1" />
+            <geometry>
+                <cylinder radius="0.05" length="0.2" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="right_knee_joint" type="revolute">
+        <origin xyz="-0.078273 -0.0021489 -0.17734" rpy="0 0.1749 0" />
+        <parent link="right_hip_yaw_link" />
+        <child link="right_knee_link" />
+        <axis xyz="0 1 0" />
+        <limit lower="-0.087267" upper="2.8798" effort="139" velocity="20" />
+    </joint>
+    <link name="right_ankle_pitch_link">
+        <inertial>
+            <origin xyz="-0.007269 0 0.011137" rpy="0 0 0" />
+            <mass value="0.074" />
+            <inertia ixx="8.4E-06" ixy="0" ixz="-2.9E-06" iyy="1.89E-05" iyz="0" izz="1.26E-05" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_ankle_pitch_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+    </link>
+    <joint name="right_ankle_pitch_joint" type="revolute">
+        <origin xyz="0 9.4445E-05 -0.30001" rpy="0 0 0" />
+        <parent link="right_knee_link" />
+        <child link="right_ankle_pitch_link" />
+        <axis xyz="0 1 0" />
+        <limit lower="-0.87267" upper="0.5236" effort="50" velocity="37" />
+    </joint>
+    <link name="right_ankle_roll_link">
+        <inertial>
+            <origin xyz="0.026505 0 -0.016425" rpy="0 0 0" />
+            <mass value="0.608" />
+            <inertia ixx="0.0002231" ixy="-2E-07" ixz="8.91E-05" iyy="0.0016161" iyz="1E-07"
+                izz="0.0016667" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_ankle_roll_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <origin xyz="0.075 -0.026 -0.025" rpy="0 1.5708 0" />
+            <geometry>
+                <cylinder radius="0.01" length="0.05" />
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.0395 -0.018 -0.025" rpy="0 -1.5708 0" />
+            <geometry>
+                <cylinder radius="0.008" length="0.167" />
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.039 -0.01 -0.025" rpy="0 -1.5708 0" />
+            <geometry>
+                <cylinder radius="0.01" length="0.182" />
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.039 0 -0.025" rpy="0 -1.5708 0" />
+            <geometry>
+                <cylinder radius="0.01" length="0.186" />
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.039 0.01 -0.025" rpy="0 -1.5708 0" />
+            <geometry>
+                <cylinder radius="0.01" length="0.182" />
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.0395 0.018 -0.025" rpy="0 -1.5708 0" />
+            <geometry>
+                <cylinder radius="0.008" length="0.167" />
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.075 0.026 -0.025" rpy="0 1.5708 0" />
+            <geometry>
+                <cylinder radius="0.01" length="0.05" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="right_ankle_roll_joint" type="revolute">
+        <origin xyz="0 0 -0.017558" rpy="0 0 0" />
+        <parent link="right_ankle_pitch_link" />
+        <child link="right_ankle_roll_link" />
+        <axis xyz="1 0 0" />
+        <limit lower="-0.2618" upper="0.2618" effort="50" velocity="37" />
+    </joint>
+
+    <!-- Torso -->
+    <link name="waist_yaw_link">
+        <inertial>
+            <origin xyz="0.003494 0.000233 0.018034" rpy="0 0 0" />
+            <mass value="0.214" />
+            <inertia ixx="0.00010673" ixy="2.703E-06" ixz="-7.631E-06" iyy="0.00010422"
+                iyz="-2.01E-07"
+                izz="0.0001625" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/waist_yaw_link_rev_1_0.STL" />
+            </geometry>
+            <material name="white">
+                <color rgba="0.7 0.7 0.7 1" />
+            </material>
+        </visual>
+    </link>
+    <joint name="waist_yaw_joint" type="revolute">
+        <origin xyz="0 0 0" rpy="0 0 0" />
+        <parent link="pelvis" />
+        <child link="waist_yaw_link" />
+        <axis xyz="0 0 1" />
+        <limit lower="-2.618" upper="2.618" effort="88" velocity="32" />
+    </joint>
+    <link name="waist_roll_link">
+        <inertial>
+            <origin xyz="0 2.3E-05 0" rpy="0 0 0" />
+            <mass value="0.086" />
+            <inertia ixx="7.079E-06" ixy="0" ixz="0" iyy="6.339E-06" iyz="0" izz="8.245E-06" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/waist_roll_link_rev_1_0.STL" />
+            </geometry>
+            <material name="white">
+                <color rgba="0.7 0.7 0.7 1" />
+            </material>
+        </visual>
+    </link>
+    <joint name="waist_roll_joint" type="revolute">
+        <origin xyz="-0.0039635 0 0.044" rpy="0 0 0" />
+        <parent link="waist_yaw_link" />
+        <child link="waist_roll_link" />
+        <axis xyz="1 0 0" />
+        <limit lower="-0.52" upper="0.52" effort="50" velocity="37" />
+    </joint>
+    <link name="torso_link">
+        <inertial>
+            <origin xyz="0.000931 0.000346 0.15082" rpy="0 0 0" />
+            <mass value="6.78" />
+            <inertia ixx="0.05905" ixy="3.3302E-05" ixz="-0.0017715" iyy="0.047014"
+                iyz="-2.2399E-05" izz="0.025652" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/torso_link_rev_1_0.STL" />
+            </geometry>
+            <material name="white">
+                <color rgba="0.7 0.7 0.7 1" />
+            </material>
+        </visual>
+        <collision>
+            <origin xyz="0.005 0 0.22" rpy="1.5708 0 0" />
+            <geometry>
+                <cylinder radius="0.073" length="0.064" />
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.005 0 0.13" rpy="1.5708 0 0" />
+            <geometry>
+                <cylinder radius="0.07" length="0.056" />
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.005 0 0.06" rpy="1.5708 0 0" />
+            <geometry>
+                <cylinder radius="0.065" length="0.04" />
+            </geometry>
+        </collision>
+        <collision>
+            <origin xyz="0.01 0 0.415" rpy="3.1416 0 0" />
+            <geometry>
+                <cylinder radius="0.068" length="0.01" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="waist_pitch_joint" type="revolute">
+        <origin xyz="0 0 0" rpy="0 0 0" />
+        <parent link="waist_roll_link" />
+        <child link="torso_link" />
+        <axis xyz="0 1 0" />
+        <limit lower="-0.52" upper="0.52" effort="50" velocity="37" />
+    </joint>
+
+    <!-- LOGO -->
+    <joint name="logo_joint" type="fixed">
+        <origin xyz="0.0039635 0 -0.044" rpy="0 0 0" />
+        <parent link="torso_link" />
+        <child link="logo_link" />
+    </joint>
+    <link name="logo_link">
+        <inertial>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <mass value="0.001" />
+            <inertia ixx="1e-7" ixy="0" ixz="0" iyy="1e-7" iyz="0" izz="1e-7" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/logo_link.STL" />
+            </geometry>
+            <material name="dark">
+                <color rgba="0.2 0.2 0.2 1" />
+            </material>
+        </visual>
+    </link>
+
+    <!-- Head -->
+    <link name="head_link">
+        <inertial>
+            <origin xyz="0.005267 0.000299 0.449869" rpy="0 0 0" />
+            <mass value="1.036" />
+            <inertia ixx="0.004085051" ixy="-2.543E-06" ixz="-6.9455E-05" iyy="0.004185212"
+                iyz="-3.726E-06"
+                izz="0.001807911" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/head_link.STL" />
+            </geometry>
+            <material name="dark">
+                <color rgba="0.2 0.2 0.2 1" />
+            </material>
+        </visual>
+    </link>
+    <joint name="head_joint" type="fixed">
+        <origin xyz="0.0039635 0 -0.044" rpy="0 0 0" />
+        <parent link="torso_link" />
+        <child link="head_link" />
+    </joint>
+
+
+    <!-- IMU -->
+    <link name="imu_in_torso"></link>
+    <joint name="imu_in_torso_joint" type="fixed">
+        <origin xyz="-0.03959 -0.00224 0.14792" rpy="0 0 0" />
+        <parent link="torso_link" />
+        <child link="imu_in_torso" />
+    </joint>
+
+    <link name="imu_in_pelvis"></link>
+    <joint name="imu_in_pelvis_joint" type="fixed">
+        <origin xyz="0.04525 0 -0.08339" rpy="0 0 0" />
+        <parent link="pelvis" />
+        <child link="imu_in_pelvis" />
+    </joint>
+
+    <!-- mid360 -->
+    <link name="mid360_link"></link>
+    <joint name="mid360_joint" type="fixed">
+        <origin xyz="0.0002835 0.00003 0.41618" rpy="0 3.101 3.1415" />
+        <parent link="torso_link" />
+        <child link="mid360_link" />
+    </joint>
+
+    <!-- Arm -->
+    <link name="left_shoulder_pitch_link">
+        <inertial>
+            <origin xyz="0 0.035892 -0.011628" rpy="0 0 0" />
+            <mass value="0.718" />
+            <inertia ixx="0.0004291" ixy="-9.2E-06" ixz="6.4E-06" iyy="0.000453" iyz="2.26E-05"
+                izz="0.000423" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh
+                    filename="package://robot_description/meshes/g1/left_shoulder_pitch_link.STL" />
+            </geometry>
+            <material name="white">
+                <color rgba="0.7 0.7 0.7 1" />
+            </material>
+        </visual>
+    </link>
+    <joint name="left_shoulder_pitch_joint" type="revolute">
+        <origin xyz="0.0039563 0.10022 0.24778" rpy="0.27931 5.4949E-05 -0.00019159" />
+        <parent link="torso_link" />
+        <child link="left_shoulder_pitch_link" />
+        <axis xyz="0 1 0" />
+        <limit lower="-3.0892" upper="2.6704" effort="25" velocity="37" />
+    </joint>
+    <link name="left_shoulder_roll_link">
+        <inertial>
+            <origin xyz="-0.000227 0.00727 -0.063243" rpy="0 0 0" />
+            <mass value="0.643" />
+            <inertia ixx="0.0006177" ixy="-1E-06" ixz="8.7E-06" iyy="0.0006912" iyz="-5.3E-06"
+                izz="0.0003894" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_shoulder_roll_link.STL" />
+            </geometry>
+            <material name="white">
+                <color rgba="0.7 0.7 0.7 1" />
+            </material>
+        </visual>
+    </link>
+    <joint name="left_shoulder_roll_joint" type="revolute">
+        <origin xyz="0 0.038 -0.013831" rpy="-0.27925 0 0" />
+        <parent link="left_shoulder_pitch_link" />
+        <child link="left_shoulder_roll_link" />
+        <axis xyz="1 0 0" />
+        <limit lower="-1.5882" upper="2.2515" effort="25" velocity="37" />
+    </joint>
+    <link name="left_shoulder_yaw_link">
+        <inertial>
+            <origin xyz="0.010773 -0.002949 -0.072009" rpy="0 0 0" />
+            <mass value="0.734" />
+            <inertia ixx="0.0009988" ixy="7.9E-06" ixz="0.0001412" iyy="0.0010605" iyz="-2.86E-05"
+                izz="0.0004354" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_shoulder_yaw_link.STL" />
+            </geometry>
+            <material name="white">
+                <color rgba="0.7 0.7 0.7 1" />
+            </material>
+        </visual>
+        <collision>
+            <origin xyz="0 0 -0.015" rpy="3.1416 0 0" />
+            <geometry>
+                <cylinder radius="0.035" length="0.13" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="left_shoulder_yaw_joint" type="revolute">
+        <origin xyz="0 0.00624 -0.1032" rpy="0 0 0" />
+        <parent link="left_shoulder_roll_link" />
+        <child link="left_shoulder_yaw_link" />
+        <axis xyz="0 0 1" />
+        <limit lower="-2.618" upper="2.618" effort="25" velocity="37" />
+    </joint>
+    <link name="left_elbow_link">
+        <inertial>
+            <origin xyz="0.064956 0.004454 -0.010062" rpy="0 0 0" />
+            <mass value="0.6" />
+            <inertia ixx="0.0002891" ixy="6.53E-05" ixz="1.72E-05" iyy="0.0004152" iyz="-5.6E-06"
+                izz="0.0004197" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_elbow_link.STL" />
+            </geometry>
+            <material name="white">
+                <color rgba="0.7 0.7 0.7 1" />
+            </material>
+        </visual>
+        <collision>
+            <origin xyz="0.055 0 -0.01" rpy="0 -1.5708 0" />
+            <geometry>
+                <cylinder radius="0.035" length="0.13" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="left_elbow_joint" type="revolute">
+        <origin xyz="0.015783 0 -0.080518" rpy="0 0 0" />
+        <parent link="left_shoulder_yaw_link" />
+        <child link="left_elbow_link" />
+        <axis xyz="0 1 0" />
+        <limit lower="-1.0472" upper="2.0944" effort="25" velocity="37" />
+    </joint>
+    <joint name="left_wrist_roll_joint" type="revolute">
+        <origin xyz="0.100 0.00188791 -0.010" rpy="0 0 0" />
+        <axis xyz="1 0 0" />
+        <parent link="left_elbow_link" />
+        <child link="left_wrist_roll_link" />
+        <limit effort="25" velocity="37" lower="-1.972222054" upper="1.972222054" />
+    </joint>
+    <link name="left_wrist_roll_link">
+        <inertial>
+            <origin xyz="0.01713944778 0.00053759094 0.00000048864" rpy="0 0 0" />
+            <mass value="0.08544498" />
+            <inertia ixx="0.00004821544023" ixy="-0.00000424511021" ixz="0.00000000510599"
+                iyy="0.00003722899093"
+                iyz="-0.00000000123525" izz="0.00005482106541" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_wrist_roll_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+    </link>
+    <joint name="left_wrist_pitch_joint" type="revolute">
+        <origin xyz="0.038 0 0" rpy="0 0 0" />
+        <axis xyz="0 1 0" />
+        <parent link="left_wrist_roll_link" />
+        <child link="left_wrist_pitch_link" />
+        <limit effort="5" velocity="22" lower="-1.614429558" upper="1.614429558" />
+    </joint>
+    <link name="left_wrist_pitch_link">
+        <inertial>
+            <origin xyz="0.02299989837 -0.00111685314 -0.00111658096" rpy="0 0 0" />
+            <mass value="0.48404956" />
+            <inertia ixx="0.00016579646273" ixy="-0.00001231206746" ixz="0.00001231699194"
+                iyy="0.00042954057410"
+                iyz="0.00000081417712" izz="0.00042953697654" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_wrist_pitch_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+    </link>
+    <joint name="left_wrist_yaw_joint" type="revolute">
+        <origin xyz="0.046 0 0" rpy="0 0 0" />
+        <axis xyz="0 0 1" />
+        <parent link="left_wrist_pitch_link" />
+        <child link="left_wrist_yaw_link" />
+        <limit effort="5" velocity="22" lower="-1.614429558" upper="1.614429558" />
+    </joint>
+    <link name="left_wrist_yaw_link">
+        <inertial>
+            <origin xyz="0.02200381568 0.00049485096 0.00053861123" rpy="0 0 0" />
+            <mass value="0.08457647" />
+            <inertia ixx="0.00004929128828" ixy="-0.00000045735494" ixz="0.00000445867591"
+                iyy="0.00005973338134"
+                iyz="0.00000043217198" izz="0.00003928083826" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_wrist_yaw_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_wrist_yaw_link.STL" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="left_hand_palm_joint" type="fixed">
+        <origin xyz="0.0415 0.003 0" rpy="0 0 0" />
+        <parent link="left_wrist_yaw_link" />
+        <child link="left_hand_palm_link" />
+    </joint>
+    <link name="left_hand_palm_link">
+        <inertial>
+            <origin xyz="0.06214634836 -0.00050869656 -0.00058171093" rpy="0 0 0" />
+            <mass value="0.37283854" />
+            <inertia ixx="0.00027535181027" ixy="-0.00001595519465" ixz="-0.00000242161890"
+                iyy="0.00053951827219" iyz="-0.00000042279435" izz="0.00039623390907" />
+        </inertial>
+        <collision>
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hand_palm_link.STL" />
+            </geometry>
+        </collision>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hand_palm_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+
+    </link>
+    <joint name="left_hand_thumb_0_joint" type="fixed">
+        <origin xyz="0.0255 0 0" rpy="0 0 0" />
+        <axis xyz="0 1 0" />
+        <parent link="left_hand_palm_link" />
+        <child link="left_hand_thumb_0_link" />
+        <limit effort="2.45" velocity="6.857" lower="-1.04719755" upper="1.04719755" />
+    </joint>
+    <link name="left_hand_thumb_0_link">
+        <inertial>
+            <origin xyz="-0.00088424580 -0.00863407079 0.00094429336" rpy="0 0 0" />
+            <mass value="0.08623657" />
+            <inertia ixx="0.00001602919238" ixy="0.00000010683177" ixz="0.00000016728875"
+                iyy="0.00001451795012" iyz="-0.00000051094752" izz="0.00001637877663" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hand_thumb_0_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hand_thumb_0_link.STL" />
+            </geometry>
+        </collision>
+
+    </link>
+    <joint name="left_hand_thumb_1_joint" type="fixed">
+        <origin xyz="-0.0025 -0.0193 0" rpy="0 0 0" />
+        <axis xyz="0 0 1" />
+        <parent link="left_hand_thumb_0_link" />
+        <child link="left_hand_thumb_1_link" />
+        <limit effort="1.4" velocity="12" lower="-0.72431163" upper="1.04719755" />
+    </joint>
+    <link name="left_hand_thumb_1_link">
+        <inertial>
+            <origin xyz="-0.00082788768 -0.03547435774 -0.00038089960" rpy="0 0 0" />
+            <mass value="0.05885070" />
+            <inertia ixx="0.00001274699945" ixy="-0.00000050770784" ixz="0.00000016088850"
+                iyy="0.00000601573947" iyz="-0.00000027839003" izz="0.00001234543582" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hand_thumb_1_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hand_thumb_1_link.STL" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="left_hand_thumb_2_joint" type="fixed">
+        <origin xyz="0 -0.0458 0" rpy="0 0 0" />
+        <axis xyz="0 0 1" />
+        <parent link="left_hand_thumb_1_link" />
+        <child link="left_hand_thumb_2_link" />
+        <limit effort="1.4" velocity="12" lower="0" upper="1.74532925" />
+    </joint>
+    <link name="left_hand_thumb_2_link">
+        <inertial>
+            <origin xyz="-0.00171735242 -0.02628192939 0.00010778879" rpy="0 0 0" />
+            <mass value="0.02030626" />
+            <inertia ixx="0.00000461267817" ixy="-0.00000003422130" ixz="-0.00000000823881"
+                iyy="0.00000153561368" iyz="-0.00000002549885" izz="0.00000386625776" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hand_thumb_2_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hand_thumb_2_link.STL" />
+            </geometry>
+        </collision>
+
+    </link>
+    <joint name="left_hand_middle_0_joint" type="fixed">
+        <origin xyz="0.0777 0.0016 -0.0285" rpy="0 0 0" />
+        <axis xyz="0 0 1" />
+        <parent link="left_hand_palm_link" />
+        <child link="left_hand_middle_0_link" />
+        <limit effort="1.4" velocity="12" lower="-1.57079632" upper="0" />
+    </joint>
+    <link name="left_hand_middle_0_link">
+        <inertial>
+            <origin xyz="0.03547435774 0.00082788768 0.00038089960" rpy="0 0 0" />
+            <mass value="0.05885070" />
+            <inertia ixx="0.00000601573947" ixy="-0.00000050770784" ixz="-0.00000027839003"
+                iyy="0.00001274699945" iyz="0.00000016088850" izz="0.00001234543582" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hand_middle_0_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hand_middle_0_link.STL" />
+            </geometry>
+        </collision>
+
+    </link>
+    <joint name="left_hand_middle_1_joint" type="fixed">
+        <origin xyz="0.0458 0 0" rpy="0 0 0" />
+        <axis xyz="0 0 1" />
+        <parent link="left_hand_middle_0_link" />
+        <child link="left_hand_middle_1_link" />
+        <limit effort="1.4" velocity="12" lower="-1.74532925" upper="0" />
+    </joint>
+    <link name="left_hand_middle_1_link">
+        <inertial>
+            <origin xyz="0.02628192939 0.00171735242 -0.00010778879" rpy="0 0 0" />
+            <mass value="0.02030626" />
+            <inertia ixx="0.00000153561368" ixy="-0.00000003422130" ixz="-0.00000002549885"
+                iyy="0.00000461267817" iyz="-0.00000000823881" izz="0.00000386625776" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hand_middle_1_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hand_middle_1_link.STL" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="left_hand_index_0_joint" type="fixed">
+        <origin xyz="0.0777 0.0016 0.0285" rpy="0 0 0" />
+        <axis xyz="0 0 1" />
+        <parent link="left_hand_palm_link" />
+        <child link="left_hand_index_0_link" />
+        <limit effort="1.4" velocity="12" lower="-1.57079632" upper="0" />
+    </joint>
+    <link name="left_hand_index_0_link">
+        <inertial>
+            <origin xyz="0.03547435774 0.00082788768 0.00038089960" rpy="0 0 0" />
+            <mass value="0.05885070" />
+            <inertia ixx="0.00000601573947" ixy="-0.00000050770784" ixz="-0.00000027839003"
+                iyy="0.00001274699945" iyz="0.00000016088850" izz="0.00001234543582" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hand_index_0_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hand_index_0_link.STL" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="left_hand_index_1_joint" type="fixed">
+        <origin xyz="0.0458 0 0" rpy="0 0 0" />
+        <axis xyz="0 0 1" />
+        <parent link="left_hand_index_0_link" />
+        <child link="left_hand_index_1_link" />
+        <limit effort="1.4" velocity="12" lower="-1.74532925" upper="0" />
+    </joint>
+    <link name="left_hand_index_1_link">
+        <inertial>
+            <origin xyz="0.02628192939 0.00171735242 -0.00010778879" rpy="0 0 0" />
+            <mass value="0.02030626" />
+            <inertia ixx="0.00000153561368" ixy="-0.00000003422130" ixz="-0.00000002549885"
+                iyy="0.00000461267817" iyz="-0.00000000823881" izz="0.00000386625776" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hand_index_1_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/left_hand_index_1_link.STL" />
+            </geometry>
+        </collision>
+    </link>
+
+    <link name="right_shoulder_pitch_link">
+        <inertial>
+            <origin xyz="0 -0.035892 -0.011628" rpy="0 0 0" />
+            <mass value="0.718" />
+            <inertia ixx="0.0004291" ixy="9.2E-06" ixz="6.4E-06" iyy="0.000453" iyz="-2.26E-05"
+                izz="0.000423" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh
+                    filename="package://robot_description/meshes/g1/right_shoulder_pitch_link.STL" />
+            </geometry>
+            <material name="white">
+                <color rgba="0.7 0.7 0.7 1" />
+            </material>
+        </visual>
+    </link>
+    <joint name="right_shoulder_pitch_joint" type="revolute">
+        <origin xyz="0.0039563 -0.10021 0.24778" rpy="-0.27931 5.4949E-05 0.00019159" />
+        <parent link="torso_link" />
+        <child link="right_shoulder_pitch_link" />
+        <axis xyz="0 1 0" />
+        <limit lower="-3.0892" upper="2.6704" effort="25" velocity="37" />
+    </joint>
+    <link name="right_shoulder_roll_link">
+        <inertial>
+            <origin xyz="-0.000227 -0.00727 -0.063243" rpy="0 0 0" />
+            <mass value="0.643" />
+            <inertia ixx="0.0006177" ixy="1E-06" ixz="8.7E-06" iyy="0.0006912" iyz="5.3E-06"
+                izz="0.0003894" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh
+                    filename="package://robot_description/meshes/g1/right_shoulder_roll_link.STL" />
+            </geometry>
+            <material name="white">
+                <color rgba="0.7 0.7 0.7 1" />
+            </material>
+        </visual>
+    </link>
+    <joint name="right_shoulder_roll_joint" type="revolute">
+        <origin xyz="0 -0.038 -0.013831" rpy="0.27925 0 0" />
+        <parent link="right_shoulder_pitch_link" />
+        <child link="right_shoulder_roll_link" />
+        <axis xyz="1 0 0" />
+        <limit lower="-2.2515" upper="1.5882" effort="25" velocity="37" />
+    </joint>
+    <link name="right_shoulder_yaw_link">
+        <inertial>
+            <origin xyz="0.010773 0.002949 -0.072009" rpy="0 0 0" />
+            <mass value="0.734" />
+            <inertia ixx="0.0009988" ixy="-7.9E-06" ixz="0.0001412" iyy="0.0010605" iyz="2.86E-05"
+                izz="0.0004354" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_shoulder_yaw_link.STL" />
+            </geometry>
+            <material name="white">
+                <color rgba="0.7 0.7 0.7 1" />
+            </material>
+        </visual>
+        <collision>
+            <origin xyz="0 0 -0.015" rpy="3.1416 0 0" />
+            <geometry>
+                <cylinder radius="0.035" length="0.13" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="right_shoulder_yaw_joint" type="revolute">
+        <origin xyz="0 -0.00624 -0.1032" rpy="0 0 0" />
+        <parent link="right_shoulder_roll_link" />
+        <child link="right_shoulder_yaw_link" />
+        <axis xyz="0 0 1" />
+        <limit lower="-2.618" upper="2.618" effort="25" velocity="37" />
+    </joint>
+    <link name="right_elbow_link">
+        <inertial>
+            <origin xyz="0.064956 -0.004454 -0.010062" rpy="0 0 0" />
+            <mass value="0.6" />
+            <inertia ixx="0.0002891" ixy="-6.53E-05" ixz="1.72E-05" iyy="0.0004152" iyz="5.6E-06"
+                izz="0.0004197" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_elbow_link.STL" />
+            </geometry>
+            <material name="white">
+                <color rgba="0.7 0.7 0.7 1" />
+            </material>
+        </visual>
+        <collision>
+            <origin xyz="0.055 0 -0.01" rpy="0 -1.5708 0" />
+            <geometry>
+                <cylinder radius="0.035" length="0.13" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="right_elbow_joint" type="revolute">
+        <origin xyz="0.015783 0 -0.080518" rpy="0 0 0" />
+        <parent link="right_shoulder_yaw_link" />
+        <child link="right_elbow_link" />
+        <axis xyz="0 1 0" />
+        <limit lower="-1.0472" upper="2.0944" effort="25" velocity="37" />
+    </joint>
+    <joint name="right_wrist_roll_joint" type="revolute">
+        <origin xyz="0.100 -0.00188791 -0.010" rpy="0 0 0" />
+        <axis xyz="1 0 0" />
+        <parent link="right_elbow_link" />
+        <child link="right_wrist_roll_link" />
+        <limit effort="25" velocity="37" lower="-1.972222054" upper="1.972222054" />
+    </joint>
+    <link name="right_wrist_roll_link">
+        <inertial>
+            <origin xyz="0.01713944778 -0.00053759094 0.00000048864" rpy="0 0 0" />
+            <mass value="0.08544498" />
+            <inertia ixx="0.00004821544023" ixy="0.00000424511021" ixz="0.00000000510599"
+                iyy="0.00003722899093"
+                iyz="0.00000000123525" izz="0.00005482106541" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_wrist_roll_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+    </link>
+    <joint name="right_wrist_pitch_joint" type="revolute">
+        <origin xyz="0.038 0 0" rpy="0 0 0" />
+        <axis xyz="0 1 0" />
+        <parent link="right_wrist_roll_link" />
+        <child link="right_wrist_pitch_link" />
+        <limit effort="5" velocity="22" lower="-1.614429558" upper="1.614429558" />
+    </joint>
+    <link name="right_wrist_pitch_link">
+        <inertial>
+            <origin xyz="0.02299989837 0.00111685314 -0.00111658096" rpy="0 0 0" />
+            <mass value="0.48404956" />
+            <inertia ixx="0.00016579646273" ixy="0.00001231206746" ixz="0.00001231699194"
+                iyy="0.00042954057410"
+                iyz="-0.00000081417712" izz="0.00042953697654" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_wrist_pitch_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+    </link>
+    <joint name="right_wrist_yaw_joint" type="revolute">
+        <origin xyz="0.046 0 0" rpy="0 0 0" />
+        <axis xyz="0 0 1" />
+        <parent link="right_wrist_pitch_link" />
+        <child link="right_wrist_yaw_link" />
+        <limit effort="5" velocity="22" lower="-1.614429558" upper="1.614429558" />
+    </joint>
+    <link name="right_wrist_yaw_link">
+        <inertial>
+            <origin xyz="0.02200381568 -0.00049485096 0.00053861123" rpy="0 0 0" />
+            <mass value="0.08457647" />
+            <inertia ixx="0.00004929128828" ixy="0.00000045735494" ixz="0.00000445867591"
+                iyy="0.00005973338134"
+                iyz="-0.00000043217198" izz="0.00003928083826" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_wrist_yaw_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_wrist_yaw_link.STL" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="right_hand_palm_joint" type="fixed">
+        <origin xyz="0.0415 -0.003 0" rpy="0 0 0" />
+        <parent link="right_wrist_yaw_link" />
+        <child link="right_hand_palm_link" />
+    </joint>
+    <link name="right_hand_palm_link">
+        <inertial>
+            <origin xyz="0.06214634836 0.00050869656 -0.00058171093" rpy="0 0 0" />
+            <mass value="0.37283854" />
+            <inertia ixx="0.00027535181027" ixy="0.00001595519465" ixz="-0.00000242161890"
+                iyy="0.00053951827219" iyz="0.00000042279435" izz="0.00039623390907" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_hand_palm_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_hand_palm_link.STL" />
+            </geometry>
+        </collision>
+
+    </link>
+    <joint name="right_hand_thumb_0_joint" type="fixed">
+        <origin xyz="0.0255 0 0" rpy="0 0 0" />
+        <axis xyz="0 1 0" />
+        <parent link="right_hand_palm_link" />
+        <child link="right_hand_thumb_0_link" />
+        <limit effort="2.45" velocity="6.857" lower="-1.04719755" upper="1.04719755" />
+    </joint>
+    <link name="right_hand_thumb_0_link">
+        <inertial>
+            <origin xyz="-0.00088424580 0.00863407079 0.00094429336" rpy="0 0 0" />
+            <mass value="0.08623657" />
+            <inertia ixx="0.00001602919238" ixy="-0.00000010683177" ixz="0.00000016728875"
+                iyy="0.00001451795012" iyz="0.00000051094752" izz="0.00001637877663" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_hand_thumb_0_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_hand_thumb_0_link.STL" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="right_hand_thumb_1_joint" type="fixed">
+        <origin xyz="-0.0025 0.0193 0" rpy="0 0 0" />
+        <axis xyz="0 0 1" />
+        <parent link="right_hand_thumb_0_link" />
+        <child link="right_hand_thumb_1_link" />
+        <limit effort="1.4" velocity="12" lower="-1.04719755" upper="0.72431163" />
+    </joint>
+    <link name="right_hand_thumb_1_link">
+        <inertial>
+            <origin xyz="-0.00082788768 0.03547435774 -0.00038089960" rpy="0 0 0" />
+            <mass value="0.05885070" />
+            <inertia ixx="0.00001274699945" ixy="0.00000050770784" ixz="0.00000016088850"
+                iyy="0.00000601573947" iyz="0.00000027839003" izz="0.00001234543582" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_hand_thumb_1_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_hand_thumb_1_link.STL" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="right_hand_thumb_2_joint" type="fixed">
+        <origin xyz="0 0.0458 0" rpy="0 0 0" />
+        <axis xyz="0 0 1" />
+        <parent link="right_hand_thumb_1_link" />
+        <child link="right_hand_thumb_2_link" />
+        <limit effort="1.4" velocity="12" lower="-1.74532925" upper="0" />
+    </joint>
+    <link name="right_hand_thumb_2_link">
+        <inertial>
+            <origin xyz="-0.00171735242 0.02628192939 0.00010778879" rpy="0 0 0" />
+            <mass value="0.02030626" />
+            <inertia ixx="0.00000461267817" ixy="0.00000003422130" ixz="-0.00000000823881"
+                iyy="0.00000153561368" iyz="0.00000002549885" izz="0.00000386625776" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_hand_thumb_2_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <geometry>
+                <mesh
+                    filename="package://robot_description/meshes/g1/right_hand_thumb_2_link.STL" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="right_hand_middle_0_joint" type="fixed">
+        <origin xyz="0.0777 -0.0016 -0.0285" rpy="0 0 0" />
+        <axis xyz="0 0 1" />
+        <parent link="right_hand_palm_link" />
+        <child link="right_hand_middle_0_link" />
+        <limit effort="1.4" velocity="12" lower="0" upper="1.57079632" />
+    </joint>
+    <link name="right_hand_middle_0_link">
+        <inertial>
+            <origin xyz="0.03547435774 -0.00082788768 0.00038089960" rpy="0 0 0" />
+            <mass value="0.05885070" />
+            <inertia ixx="0.00000601573947" ixy="0.00000050770784" ixz="-0.00000027839003"
+                iyy="0.00001274699945" iyz="-0.00000016088850" izz="0.00001234543582" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh
+                    filename="package://robot_description/meshes/g1/right_hand_middle_0_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <geometry>
+                <mesh
+                    filename="package://robot_description/meshes/g1/right_hand_middle_0_link.STL" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="right_hand_middle_1_joint" type="fixed">
+        <origin xyz="0.0458 0 0" rpy="0 0 0" />
+        <axis xyz="0 0 1" />
+        <parent link="right_hand_middle_0_link" />
+        <child link="right_hand_middle_1_link" />
+        <limit effort="1.4" velocity="12" lower="0" upper="1.74532925" />
+    </joint>
+    <link name="right_hand_middle_1_link">
+        <inertial>
+            <origin xyz="0.02628192939 -0.00171735242 -0.00010778879" rpy="0 0 0" />
+            <mass value="0.02030626" />
+            <inertia ixx="0.00000153561368" ixy="0.00000003422130" ixz="-0.00000002549885"
+                iyy="0.00000461267817" iyz="0.00000000823881" izz="0.00000386625776" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh
+                    filename="package://robot_description/meshes/g1/right_hand_middle_1_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <geometry>
+                <mesh
+                    filename="package://robot_description/meshes/g1/right_hand_middle_1_link.STL" />
+            </geometry>
+        </collision>
+
+    </link>
+    <joint name="right_hand_index_0_joint" type="fixed">
+        <origin xyz="0.0777 -0.0016 0.0285" rpy="0 0 0" />
+        <axis xyz="0 0 1" />
+        <parent link="right_hand_palm_link" />
+        <child link="right_hand_index_0_link" />
+        <limit effort="1.4" velocity="12" lower="0" upper="1.57079632" />
+    </joint>
+    <link name="right_hand_index_0_link">
+        <inertial>
+            <origin xyz="0.03547435774 -0.00082788768 0.00038089960" rpy="0 0 0" />
+            <mass value="0.05885070" />
+            <inertia ixx="0.00000601573947" ixy="0.00000050770784" ixz="-0.00000027839003"
+                iyy="0.00001274699945" iyz="-0.00000016088850" izz="0.00001234543582" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_hand_index_0_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_hand_index_0_link.STL" />
+            </geometry>
+        </collision>
+    </link>
+    <joint name="right_hand_index_1_joint" type="fixed">
+        <origin xyz="0.0458 0 0" rpy="0 0 0" />
+        <axis xyz="0 0 1" />
+        <parent link="right_hand_index_0_link" />
+        <child link="right_hand_index_1_link" />
+        <limit effort="1.4" velocity="12" lower="0" upper="1.74532925" />
+    </joint>
+    <link name="right_hand_index_1_link">
+        <inertial>
+            <origin xyz="0.02628192939 -0.00171735242 -0.00010778879" rpy="0 0 0" />
+            <mass value="0.02030626" />
+            <inertia ixx="0.00000153561368" ixy="0.00000003422130" ixz="-0.00000002549885"
+                iyy="0.00000461267817" iyz="0.00000000823881" izz="0.00000386625776" />
+        </inertial>
+        <visual>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_hand_index_1_link.STL" />
+            </geometry>
+            <material name="black">
+                <color rgba="0.05 0.05 0.05 1" />
+            </material>
+        </visual>
+        <collision>
+            <geometry>
+                <mesh filename="package://robot_description/meshes/g1/right_hand_index_1_link.STL" />
+            </geometry>
+        </collision>
+    </link>
+
+
+    <link name="LL_FOOT">
+        <inertial>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <mass value="1e-9" />
+            <inertia ixx="0" ixy="0" ixz="0" iyy="0" iyz="0" izz="0" />
+        </inertial>
+    </link>
+    <joint name="LL_FOOT_frame" type="fixed">
+        <origin xyz="0.04 0 -0.037" rpy="0 0 0" />
+        <parent link="left_ankle_roll_link" />
+        <child link="LL_FOOT" />
+    </joint>
+
+    <link name="LR_FOOT">
+        <inertial>
+            <origin xyz="0 0 0" rpy="0 0 0" />
+            <mass value="1e-9" />
+            <inertia ixx="0" ixy="0" ixz="0" iyy="0" iyz="0" izz="0" />
+        </inertial>
+    </link>
+    <joint name="LR_FOOT_frame" type="fixed">
+        <origin xyz="0.04 0 -0.037" rpy="0 0 0" />
+        <parent link="right_ankle_roll_link" />
+        <child link="LR_FOOT" />
+    </joint>
+</robot>
\ No newline at end of file
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/h2.urdf b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/h2.urdf
new file mode 100644
index 0000000000000000000000000000000000000000..c114dc56dc4fb0100a3b300a6680b72af4fd054b
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/h2.urdf
@@ -0,0 +1,885 @@
+<robot name="H1031">
+  <material name="dark">
+    <color rgba="0.4 0.4 0.4 1"/>
+  </material>
+  <material name="white">
+    <color rgba="0.7 0.7 0.7 1"/>
+  </material>
+
+  <mujoco>
+    <compiler meshdir="meshes" discardvisual="false"/>
+  </mujoco>
+
+  <!-- <link name="world"></link>
+  <joint name="floating_base_joint" type="floating">
+    <parent link="world"/>
+    <child link="pelvis"/>
+  </joint> -->
+
+  <link name="pelvis">
+    <inertial>
+      <origin xyz="0.00046244 -0.00004214 -0.01473435" rpy="0 0 0"/>
+      <mass value="7.60468095"/>
+      <inertia ixx="0.03548205" ixy="-0.00060277" ixz="0.00014978" iyy="0.02290677" iyz="-0.00000038" izz="0.02753519"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/pelvis.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/pelvis_collision.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="left_hip_pitch_link">
+    <inertial>
+      <origin xyz="0.00300595 0.05931469 -0.01709150" rpy="0 0 0"/>
+      <mass value="2.79752530"/>
+      <inertia ixx="0.00509731" ixy="-0.00008393" ixz="0.00004539" iyy="0.00371022" iyz="0.00030980" izz="0.00385560"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_hip_pitch_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_hip_pitch_link.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_hip_pitch_joint" type="revolute">
+    <origin xyz="0 0.083873 -0.048424" rpy="-0.5236 0 0"/>
+    <parent link="pelvis"/>
+    <child link="left_hip_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-2.4525975" upper="2.7754225" effort="360" velocity="20"/>
+  </joint>
+  <link name="left_hip_roll_link">
+    <inertial>
+      <origin xyz="-0.03026193 -0.00468930 -0.04449852" rpy="0 0 0"/>
+      <mass value="0.62714607"/>
+      <inertia ixx="0.00130633" ixy="-0.00000552" ixz="-0.00013133" iyy="0.00168791" iyz="-0.00006773" izz="0.00196361"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_hip_roll_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_hip_roll_link.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_hip_roll_joint" type="revolute">
+    <origin xyz="0.041851 0.0635 -0.019" rpy="0.5236 0 0"/>
+    <parent link="left_hip_pitch_link"/>
+    <child link="left_hip_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-0.467441" upper="2.1688605" effort="360" velocity="20"/>
+  </joint>
+  <link name="left_hip_yaw_link">
+    <inertial>
+      <origin xyz="0.00069408 -0.01086587 -0.11301456" rpy="0 0 0"/>
+      <mass value="7.08752046"/>
+      <inertia ixx="0.05057956" ixy="-0.00005547" ixz="0.00146675" iyy="0.05110490" iyz="-0.00312202" izz="0.01331946"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_hip_yaw_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_hip_yaw_link.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_hip_yaw_joint" type="revolute">
+    <origin xyz="-0.028187 -0.0095 -0.073" rpy="0 0 0"/>
+    <parent link="left_hip_roll_link"/>
+    <child link="left_hip_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.827" upper="2.827" effort="360" velocity="20"/>
+  </joint>
+  <link name="left_knee_link">
+    <inertial>
+      <origin xyz="0.00461331 -0.02453217 -0.16715459" rpy="0 0 0"/>
+      <mass value="3.55340244"/>
+      <inertia ixx="0.05108514" ixy="-0.00000720" ixz="-0.00041631" iyy="0.05123842" iyz="0.00001339" izz="0.00334361"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_knee_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0.01 -0.025 -0.25" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.04" length="0.35"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_knee_joint" type="revolute">
+    <origin xyz="-0.010341 0.018489 -0.3191" rpy="0 0 0"/>
+    <parent link="left_hip_yaw_link"/>
+    <child link="left_knee_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.08725" upper="2.53025" effort="360" velocity="20"/>
+  </joint>
+  <link name="left_ankle_roll_link">
+    <inertial>
+      <origin xyz="-0.01740632 0.00006424 -0.00021460" rpy="0 0 0"/>
+      <mass value="0.12318157"/>
+      <inertia ixx="0.00003432" ixy="0.00000012" ixz="0.00000002" iyy="0.00001450" iyz="-0.00000002" izz="0.00004005"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_ankle_roll_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+  </link>
+  <joint name="left_ankle_roll_joint" type="revolute">
+    <origin xyz="0.03272 -0.023002 -0.49662" rpy="0 0 0"/>
+    <parent link="left_knee_link"/>
+    <child link="left_ankle_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-0.3490658503" upper="0.296705972783" effort="19" velocity="100.7"/>
+  </joint>
+  <link name="left_ankle_pitch_link">
+    <inertial>
+      <origin xyz="0.03205900 0.03700531 -0.02882818" rpy="0 0 0"/>
+      <mass value="0.84496747"/>
+      <inertia ixx="0.00060825" ixy="0.00000046" ixz="0.00017671" iyy="0.00410964" iyz="-0.00000001" izz="0.00437490"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_ankle_pitch_link.stl"/>
+      </geometry>
+      <material name="dark"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_ankle_pitch_link.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_ankle_pitch_joint" type="revolute">
+    <origin xyz="-0.0225 -0.0373 0" rpy="0 0 0"/>
+    <parent link="left_ankle_roll_link"/>
+    <child link="left_ankle_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-1.13446401358" upper="0.61086523808" effort="66.88" velocity="28.61"/>
+  </joint>
+  <link name="right_hip_pitch_link">
+    <inertial>
+      <origin xyz="0.00300595 -0.05931469 -0.01709150" rpy="0 0 0"/>
+      <mass value="2.79752530"/>
+      <inertia ixx="0.00509731" ixy="0.00008393" ixz="0.00004539" iyy="0.00371022" iyz="-0.00030980" izz="0.00385560"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_hip_pitch_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_hip_pitch_link.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_hip_pitch_joint" type="revolute">
+    <origin xyz="0 -0.083873 -0.048424" rpy="0.5236 0 0"/>
+    <parent link="pelvis"/>
+    <child link="right_hip_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-2.4525975" upper="2.7754225" effort="360" velocity="20"/>
+  </joint>
+  <link name="right_hip_roll_link">
+    <inertial>
+      <origin xyz="-0.03026193 0.00468930 -0.04449852" rpy="0 0 0"/>
+      <mass value="0.62714607"/>
+      <inertia ixx="0.00130633" ixy="0.00000552" ixz="-0.00013133" iyy="0.00168791" iyz="0.00006773" izz="0.00196361"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_hip_roll_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_hip_roll_link.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_hip_roll_joint" type="revolute">
+    <origin xyz="0.041851 -0.0635 -0.019" rpy="-0.5236 0 0"/>
+    <parent link="right_hip_pitch_link"/>
+    <child link="right_hip_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-2.1688605" upper="0.467441" effort="360" velocity="20"/>
+  </joint>
+  <link name="right_hip_yaw_link">
+    <inertial>
+      <origin xyz="0.00069408 0.01086587 -0.11301456" rpy="0 0 0"/>
+      <mass value="7.08752046"/>
+      <inertia ixx="0.05057956" ixy="0.00005547" ixz="0.00146675" iyy="0.05110490" iyz="0.00312202" izz="0.01331946"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_hip_yaw_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_hip_yaw_link.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_hip_yaw_joint" type="revolute">
+    <origin xyz="-0.028187 0.0095 -0.073" rpy="0 0 0"/>
+    <parent link="right_hip_roll_link"/>
+    <child link="right_hip_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.827" upper="2.827" effort="360" velocity="20"/>
+  </joint>
+  <link name="right_knee_link">
+    <inertial>
+      <origin xyz="0.00461331 0.02453217 -0.16715459" rpy="0 0 0"/>
+      <mass value="3.55340244"/>
+      <inertia ixx="0.05108514" ixy="0.00000720" ixz="-0.00041631" iyy="0.05123842" iyz="-0.00001339" izz="0.00334361"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_knee_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0.01 0.025 -0.25" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.04" length="0.35"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_knee_joint" type="revolute">
+    <origin xyz="-0.010341 -0.018489 -0.3191" rpy="0 0 0"/>
+    <parent link="right_hip_yaw_link"/>
+    <child link="right_knee_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.08725" upper="2.53025" effort="360" velocity="20"/>
+  </joint>
+  <link name="right_ankle_roll_link">
+    <inertial>
+      <origin xyz="-0.01740632 -0.00006424 -0.00021460" rpy="0 0 0"/>
+      <mass value="0.12318157"/>
+      <inertia ixx="0.00003432" ixy="-0.00000012" ixz="0.00000002" iyy="0.00001450" iyz="0.00000002" izz="0.00004005"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_ankle_roll_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+  </link>
+  <joint name="right_ankle_roll_joint" type="revolute">
+    <origin xyz="0.03272 0.023002 -0.49662" rpy="0 0 0"/>
+    <parent link="right_knee_link"/>
+    <child link="right_ankle_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-0.296705972783" upper="0.3490658503" effort="19" velocity="100.7"/>
+  </joint>
+  <link name="right_ankle_pitch_link">
+    <inertial>
+      <origin xyz="0.03205900 -0.03700531 -0.02882818" rpy="0 0 0"/>
+      <mass value="0.84496747"/>
+      <inertia ixx="0.00060825" ixy="-0.00000046" ixz="0.00017671" iyy="0.00410964" iyz="0.00000001" izz="0.00437490"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_ankle_pitch_link.stl"/>
+      </geometry>
+      <material name="dark"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_ankle_pitch_link.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_ankle_pitch_joint" type="revolute">
+    <origin xyz="-0.0225 0.0373 0" rpy="0 0 0"/>
+    <parent link="right_ankle_roll_link"/>
+    <child link="right_ankle_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-1.13446401358" upper="0.61086523808" effort="66.88" velocity="28.61"/>
+  </joint>
+  <link name="waist_yaw_link">
+    <inertial>
+      <origin xyz="0.00626994 0 0.01696958" rpy="0 0 0"/>
+      <mass value="0.27976220"/>
+      <inertia ixx="0.00012352" ixy="0.00000000" ixz="0.00000569" iyy="0.00024730" iyz="0.00000000" izz="0.00031125"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/waist_yaw_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+  </link>
+  <joint name="waist_yaw_joint" type="revolute">
+    <origin xyz="0 0 0.089837" rpy="0 0 0"/>
+    <parent link="pelvis"/>
+    <child link="waist_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-1.7453" upper="1.7453" effort="120" velocity="28.375"/>
+  </joint>
+  <link name="waist_roll_link">
+    <inertial>
+      <origin xyz="0.00000000 0.00128474 0.00000000" rpy="0 0 0"/>
+      <mass value="0.06585040"/>
+      <inertia ixx="0.00000809" ixy="0.00000000" ixz="0.00000000" iyy="0.00001041" iyz="0.00000000" izz="0.00001580"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/waist_roll_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+  </link>
+  <joint name="waist_roll_joint" type="revolute">
+    <origin xyz="0 0 0" rpy="0 0 0"/>
+    <parent link="waist_yaw_link"/>
+    <child link="waist_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-0.5236" upper="0.5236" effort="180" velocity="28.375"/>
+  </joint>
+  <link name="torso_link">
+    <inertial>
+      <origin xyz="0.01342470 0.00547190 0.21050133" rpy="0 0 0"/>
+      <mass value="15.87026294"/>
+      <inertia ixx="0.27305029" ixy="0.00013871" ixz="-0.00147103" iyy="0.24363561" iyz="-0.01093518" izz="0.08484661"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/torso_link.stl"/>
+      </geometry>
+      <material name="dark"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/torso_collision.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="waist_pitch_joint" type="revolute">
+    <origin xyz="0 0 0" rpy="0 0 0"/>
+    <parent link="waist_roll_link"/>
+    <child link="torso_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.43633" upper="0.5236" effort="180" velocity="28.375"/>
+  </joint>
+
+  <!-- head -->
+
+  <link name="head_pitch_link">
+    <inertial>
+      <origin xyz="5.279E-08 -0.00035318 0.045125" rpy="0 0 0"/>
+      <mass value="1.82"/>
+      <inertia ixx="0.005386858" ixy="-0.000000033" ixz="0.000020867" iyy="0.005324406" iyz="-0.000007491" izz="0.001409942"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/head_pitch_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <!-- <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/head_pitch_link.stl"/>
+      </geometry>
+    </collision> -->
+  </link>
+  <joint name="head_pitch_joint" type="revolute">
+    <origin xyz="-0.0002283 2.7112E-05 0.42181" rpy="0 0 0"/>
+    <parent link="torso_link"/>
+    <child link="head_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.5236" upper="0.83775" effort="50" velocity="10"/>
+  </joint>
+
+  <link name="head_yaw_link">
+    <inertial>
+      <origin xyz="0.0219 0 0.0187" rpy="0 0 0"/>
+      <mass value="1.3"/>
+      <inertia ixx="0.004367405" ixy="-0.00000739" ixz="0.000456883" iyy="0.005352667" iyz="0.000007463" izz="0.004136928"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/head_yaw_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/head_yaw_link.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="head_yaw_joint" type="revolute">
+    <origin xyz="0 0.0003 0.1325" rpy="0 0 0"/>
+    <parent link="head_pitch_link"/>
+    <child link="head_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-1.7453" upper="1.7453" effort="50" velocity="10"/>
+  </joint>
+
+  <!-- arm -->
+
+  <link name="left_shoulder_pitch_link">
+    <inertial>
+      <origin xyz="0.00296153 0.05202683 -0.02504502" rpy="0 0 0"/>
+      <mass value="1.06055553"/>
+      <inertia ixx="0.00121801" ixy="-0.000027444" ixz="0.00004648" iyy="0.00090494" iyz="0.00029813" izz="0.00066362"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_shoulder_pitch_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_shoulder_pitch_link.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_shoulder_pitch_joint" type="revolute">
+    <origin xyz="0.0015 0.11254 0.35812" rpy="0.26914 0 0"/>
+    <parent link="torso_link"/>
+    <child link="left_shoulder_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-2.6179939" upper="1.8325957" effort="120" velocity="28.375"/>
+  </joint>
+  <link name="left_shoulder_roll_link">
+    <inertial>
+      <origin xyz="-0.02061003 0.00415443 -0.06294221" rpy="0 0 0"/>
+      <mass value="1.34511054"/>
+      <inertia ixx="0.00191998" ixy="0.00000294" ixz="-0.00023793" iyy="0.00210021" iyz="0.00004035" izz="0.00112386"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_shoulder_roll_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_shoulder_roll_link.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_shoulder_roll_joint" type="revolute">
+    <origin xyz="0.027682 0.061638 -0.039909" rpy="-0.26914 0 0"/>
+    <parent link="left_shoulder_pitch_link"/>
+    <child link="left_shoulder_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-0.5166175" upper="2.6354472" effort="54" velocity="34.3"/>
+  </joint>
+  <link name="left_shoulder_yaw_link">
+    <inertial>
+      <origin xyz="0.00364187 0.01865147 -0.05585548" rpy="0 0 0"/>
+      <mass value="0.91571732"/>
+      <inertia ixx="0.00304064" ixy="0.00001983" ixz="0.00018692" iyy="0.00254925" iyz="-0.00016348" izz="0.00147669"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_shoulder_yaw_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_shoulder_yaw_link.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_shoulder_yaw_joint" type="revolute">
+    <origin xyz="-0.030632 0.0037809 -0.10422" rpy="0 0 0"/>
+    <parent link="left_shoulder_roll_link"/>
+    <child link="left_shoulder_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.6179939" upper="2.6179939" effort="54" velocity="34.3"/>
+  </joint>
+  <link name="left_elbow_link">
+    <inertial>
+      <origin xyz="0.04599687 -0.02832961 -0.00661325" rpy="0 0 0"/>
+      <mass value="1.69431869"/>
+      <inertia ixx="0.00071506" ixy="0.00027494" ixz="0.00013832" iyy="0.00547116" iyz="-0.00001668" izz="0.00533470"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_elbow_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_elbow_link.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_elbow_joint" type="revolute">
+    <origin xyz="0.016955 0.031006 -0.12719" rpy="0 0 0"/>
+    <parent link="left_shoulder_yaw_link"/>
+    <child link="left_elbow_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.986111" upper="3.0717795" effort="54" velocity="34.3"/>
+  </joint>
+  <link name="left_wrist_roll_link">
+    <inertial>
+      <origin xyz="0.03371947 0.00237482 0.00031544" rpy="0 0 0"/>
+      <mass value="0.96983674"/>
+      <inertia ixx="0.00087762" ixy="-0.00002554" ixz="-0.00000226" iyy="0.00087920" iyz="0.00000022" izz="0.00092684"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_wrist_roll_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <!-- <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_wrist_roll_link.stl"/>
+      </geometry>
+    </collision> -->
+  </link>
+  <joint name="left_wrist_roll_joint" type="revolute">
+    <origin xyz="0.12826 -0.030815 -0.013299" rpy="0 0 0"/>
+    <parent link="left_elbow_link"/>
+    <child link="left_wrist_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-2.6179939" upper="2.6179939" effort="54" velocity="34.3"/>
+  </joint>
+  <link name="left_wrist_pitch_link">
+    <inertial>
+      <origin xyz="0.01820997 0.03140354 0.01156873" rpy="0 0 0"/>
+      <mass value="0.94695328"/>
+      <inertia ixx="0.00087455" ixy="0.00001140" ixz="0.00009847" iyy="0.00082105" iyz="-0.00000641" izz="0.00069205"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_wrist_pitch_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+  </link>
+  <joint name="left_wrist_pitch_joint" type="revolute">
+    <origin xyz="0.082003 -0.031003 0.00020786" rpy="0 0 0"/>
+    <parent link="left_wrist_roll_link"/>
+    <child link="left_wrist_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.4363323" upper="0.4363323" effort="25" velocity="50.125"/>
+  </joint>
+  <link name="left_wrist_yaw_link">
+    <inertial>
+      <origin xyz="-0.01011731 0.00400457 -0.01730713" rpy="0 0 0"/>
+      <mass value="0.60575230"/>
+      <inertia ixx="0.00024292" ixy="0.00002474" ixz="0.00003864" iyy="0.00076758" iyz="-0.00003303" izz="0.00067790"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 -0.029" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.028" length="0.06"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 -0.029" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.04"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_wrist_yaw_joint" type="revolute">
+    <origin xyz="0.026025 0.030996 0.026981" rpy="0 0 0"/>
+    <parent link="left_wrist_pitch_link"/>
+    <child link="left_wrist_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-1.2217305" upper="1.2217305" effort="25" velocity="50.125"/>
+  </joint>
+  <joint name="left_hand_palm_joint" type="fixed">
+    <origin xyz="0.025 0 -0.03" rpy="0 0 0"/>
+    <parent link="left_wrist_yaw_link"/>
+    <child link="left_rubber_hand"/>
+  </joint>
+  <link name="left_rubber_hand">
+    <inertial>
+      <origin xyz="0.05361310808 -0.00295905240 0.00215413091" rpy="0 0 0"/>
+      <mass value="0.170"/>
+      <inertia ixx="0.00010099485234748" ixy="0.00003618590790516" ixz="-0.00000074301518642" iyy="0.00028135871571621" iyz="0.00000330189743286" izz="0.00021894770413514"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/left_rubber_hand.stl"/>
+      </geometry>
+      <material name="dark"/>
+    </visual>
+  </link>
+  <link name="right_shoulder_pitch_link">
+    <inertial>
+      <origin xyz="0.00296153 -0.05202683 -0.02504502" rpy="0 0 0"/>
+      <mass value="1.06055553"/>
+      <inertia ixx="0.00121801" ixy="0.00002744" ixz="0.00004648" iyy="0.00090494" iyz="-0.00029813" izz="0.00066362"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_shoulder_pitch_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_shoulder_pitch_link.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_shoulder_pitch_joint" type="revolute">
+    <origin xyz="0.0015 -0.1125 0.35812" rpy="-0.26914 0 0"/>
+    <parent link="torso_link"/>
+    <child link="right_shoulder_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-2.6179939" upper="1.8325957" effort="120" velocity="28.375"/>
+  </joint>
+  <link name="right_shoulder_roll_link">
+    <inertial>
+      <origin xyz="-0.02061003 -0.00415443 -0.06294221" rpy="0 0 0"/>
+      <mass value="1.34511054"/>
+      <inertia ixx="0.00191998" ixy="-0.00000294" ixz="-0.00023793" iyy="0.00210021" iyz="-0.00004035" izz="0.00112386"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_shoulder_roll_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_shoulder_roll_link.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_shoulder_roll_joint" type="revolute">
+    <origin xyz="0.027682 -0.061638 -0.039909" rpy="0.26914 0 0"/>
+    <parent link="right_shoulder_pitch_link"/>
+    <child link="right_shoulder_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-2.6354472" upper="0.5166175" effort="54" velocity="34.3"/>
+  </joint>
+  <link name="right_shoulder_yaw_link">
+    <inertial>
+      <origin xyz="0.00364187 -0.01865147 -0.05585548" rpy="0 0 0"/>
+      <mass value="0.91571732"/>
+      <inertia ixx="0.00304064" ixy="-0.00001983" ixz="0.00018692" iyy="0.00254925" iyz="0.00016348" izz="0.00147669"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_shoulder_yaw_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_shoulder_yaw_link.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_shoulder_yaw_joint" type="revolute">
+    <origin xyz="-0.030632 -0.0037809 -0.10422" rpy="0 0 0"/>
+    <parent link="right_shoulder_roll_link"/>
+    <child link="right_shoulder_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.6179939" upper="2.6179939" effort="54" velocity="34.3"/>
+  </joint>
+  <link name="right_elbow_link">
+    <inertial>
+      <origin xyz="0.04599687 0.02832961 -0.00661325" rpy="0 0 0"/>
+      <mass value="1.69431869"/>
+      <inertia ixx="0.00071506" ixy="-0.00027494" ixz="0.00013832" iyy="0.00547116" iyz="0.00001668" izz="0.00533470"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_elbow_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_elbow_link.stl"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_elbow_joint" type="revolute">
+    <origin xyz="0.016955 -0.031006 -0.12719" rpy="0 0 0"/>
+    <parent link="right_shoulder_yaw_link"/>
+    <child link="right_elbow_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.986111" upper="3.0717795" effort="54" velocity="34.3"/>
+  </joint>
+  <link name="right_wrist_roll_link">
+    <inertial>
+      <origin xyz="0.03371947 -0.00237482 0.00031544" rpy="0 0 0"/>
+      <mass value="0.96983674"/>
+      <inertia ixx="0.00087762" ixy="0.00002554" ixz="-0.00000226" iyy="0.00087920" iyz="-0.00000022" izz="0.00092684"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_wrist_roll_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <!-- <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_wrist_roll_link.stl"/>
+      </geometry>
+    </collision> -->
+  </link>
+  <joint name="right_wrist_roll_joint" type="revolute">
+    <origin xyz="0.12826 0.030815 -0.013299" rpy="0 0 0"/>
+    <parent link="right_elbow_link"/>
+    <child link="right_wrist_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-2.6179939" upper="2.6179939" effort="54" velocity="34.3"/>
+  </joint>
+  <link name="right_wrist_pitch_link">
+    <inertial>
+      <origin xyz="0.01820997 -0.03140354 0.01156873" rpy="0 0 0"/>
+      <mass value="0.94695328"/>
+      <inertia ixx="0.00087455" ixy="-0.00001140" ixz="0.00009847" iyy="0.00082105" iyz="0.00000641" izz="0.00069205"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_wrist_pitch_link.stl"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+  </link>
+  <joint name="right_wrist_pitch_joint" type="revolute">
+    <origin xyz="0.082003 0.031003 0.00020786" rpy="0 0 0"/>
+    <parent link="right_wrist_roll_link"/>
+    <child link="right_wrist_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.4363323" upper="0.4363323" effort="25" velocity="50.125"/>
+  </joint>
+  <link name="right_wrist_yaw_link">
+    <inertial>
+      <origin xyz="-0.01011731 -0.00400457 -0.01730713" rpy="0 0 0"/>
+      <mass value="0.60575230"/>
+      <inertia ixx="0.00024292" ixy="-0.00002474" ixz="0.00003864" iyy="0.00076758" iyz="0.00003303" izz="0.00067790"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 -0.029" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.028" length="0.06"/>
+      </geometry>
+      <material name="white"/>
+    </visual>
+    <collision>
+      <origin xyz="0 0 -0.029" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.04"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_wrist_yaw_joint" type="revolute">
+    <origin xyz="0.026025 -0.030996 0.026981" rpy="0 0 0"/>
+    <parent link="right_wrist_pitch_link"/>
+    <child link="right_wrist_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-1.2217305" upper="1.2217305" effort="25" velocity="50.125"/>
+  </joint>
+  <joint name="right_hand_palm_joint" type="fixed">
+    <origin xyz="0.025 0 -0.03" rpy="0 0 0"/>
+    <parent link="right_wrist_yaw_link"/>
+    <child link="right_rubber_hand"/>
+  </joint>
+  <link name="right_rubber_hand">
+    <inertial>
+      <origin xyz="0.05361310808 0.00295905240 0.00215413091" rpy="0 0 0"/>
+      <mass value="0.170"/>
+      <inertia ixx="0.00010099485234748" ixy="-0.00003618590790516" ixz="-0.00000074301518642" iyy="0.00028135871571621" iyz="-0.00000330189743286" izz="0.00021894770413514"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/right_rubber_hand.stl"/>
+      </geometry>
+      <material name="dark"/>
+    </visual>
+  </link>
+</robot>
\ No newline at end of file
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/head_pitch_link.stl b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/head_pitch_link.stl
new file mode 100644
index 0000000000000000000000000000000000000000..21ea7629b5045a8ec3ecd3408c5a8ae243f7b23b
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/head_pitch_link.stl differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/left_ankle_roll_link.stl b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/left_ankle_roll_link.stl
new file mode 100644
index 0000000000000000000000000000000000000000..bf1707d917c0e9d6b92bdf66434685b2d64b8e65
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/left_ankle_roll_link.stl differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/left_wrist_pitch_link.stl b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/left_wrist_pitch_link.stl
new file mode 100644
index 0000000000000000000000000000000000000000..3d337ae7297bdd7c747d69eea76b861dda2876ee
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/left_wrist_pitch_link.stl differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/right_ankle_roll_link.stl b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/right_ankle_roll_link.stl
new file mode 100644
index 0000000000000000000000000000000000000000..72eebc5f41598b2e7309c3969c9558ff273c6329
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/right_ankle_roll_link.stl differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/right_wrist_pitch_link.stl b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/right_wrist_pitch_link.stl
new file mode 100644
index 0000000000000000000000000000000000000000..8f272b8905084c51eb305c04e5ccec5ea0480c8c
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/right_wrist_pitch_link.stl differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/waist_roll_link.stl b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/waist_roll_link.stl
new file mode 100644
index 0000000000000000000000000000000000000000..67c67d744e6da16e7f0940c84da828fca2902ea9
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/waist_roll_link.stl differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/waist_yaw_link.stl b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/waist_yaw_link.stl
new file mode 100644
index 0000000000000000000000000000000000000000..ae2d6d6bbdbe85378d45d24324d5686d525597d0
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/assets/robot_description/urdf/h2/meshes/waist_yaw_link.stl differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_29dof_cylinder/configuration/main_nodex_physics.usd b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_29dof_cylinder/configuration/main_nodex_physics.usd
new file mode 100644
index 0000000000000000000000000000000000000000..b224bcfc2a6b117c23195f789754edbee19ac783
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_29dof_cylinder/configuration/main_nodex_physics.usd differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_29dof_cylinder/configuration/main_nodex_sensor.usd b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_29dof_cylinder/configuration/main_nodex_sensor.usd
new file mode 100644
index 0000000000000000000000000000000000000000..a842196f531d5e4bb56492a9000b7f09d5b1c944
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_29dof_cylinder/configuration/main_nodex_sensor.usd differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_29dof_rev_1_0/configuration/g1_29dof_rev_1_0_physics.usd b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_29dof_rev_1_0/configuration/g1_29dof_rev_1_0_physics.usd
new file mode 100644
index 0000000000000000000000000000000000000000..b96a97e9d7358096d2c833236160ce300343828f
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_29dof_rev_1_0/configuration/g1_29dof_rev_1_0_physics.usd differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_29dof_rev_1_0/configuration/g1_29dof_rev_1_0_sensor.usd b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_29dof_rev_1_0/configuration/g1_29dof_rev_1_0_sensor.usd
new file mode 100644
index 0000000000000000000000000000000000000000..45403a1089e6ef28bec167e11da0455596d53819
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_29dof_rev_1_0/configuration/g1_29dof_rev_1_0_sensor.usd differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_sonic_usd/configuration/g1_43dof_main_physics.usd b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_sonic_usd/configuration/g1_43dof_main_physics.usd
new file mode 100644
index 0000000000000000000000000000000000000000..2e4ea219e2dba562fb96ed2acaa6cd1c0efa47d4
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_sonic_usd/configuration/g1_43dof_main_physics.usd differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_sonic_usd/configuration/g1_43dof_main_robot.usd b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_sonic_usd/configuration/g1_43dof_main_robot.usd
new file mode 100644
index 0000000000000000000000000000000000000000..70f2d08bde4b92e646b1dee1d020122078097db2
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_sonic_usd/configuration/g1_43dof_main_robot.usd differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_sonic_usd/configuration/g1_43dof_main_sensor.usd b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_sonic_usd/configuration/g1_43dof_main_sensor.usd
new file mode 100644
index 0000000000000000000000000000000000000000..c1d897ca11c31e5b1b2b4dfc210be75313b8ffa9
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/g1_sonic_usd/configuration/g1_43dof_main_sensor.usd differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/g1_10dof_sausage.urdf b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/g1_10dof_sausage.urdf
new file mode 100644
index 0000000000000000000000000000000000000000..a00a6726242c7683261f4c6423b5c7018e4b55ce
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/g1_10dof_sausage.urdf
@@ -0,0 +1,1299 @@
+<robot name="g1">
+
+  <link name="pelvis">
+    <inertial>
+      <origin xyz="0 0 -0.07605" rpy="0 0 0"/>
+      <mass value="2.86"/>
+      <inertia ixx="0.0079143" ixy="0" ixz="1.6E-06" iyy="0.0069837" iyz="0" izz="0.0059404"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/pelvis.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/pelvis.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="pelvis_contour_link">
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/pelvis_contour_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+  </link>
+  <joint name="pelvis_contour_joint" type="fixed">
+    <parent link="pelvis"/>
+    <child link="pelvis_contour_link"/>
+  </joint>
+
+  <!-- Leg -->
+  <link name="left_hip_pitch_link">
+    <inertial>
+      <origin xyz="0.001962 0.049392 -0.000941" rpy="0 0 0"/>
+      <mass value="1.299"/>
+      <inertia ixx="0.0013873" ixy="-1.63E-05" ixz="-1E-06" iyy="0.0009059" iyz="-4.24E-05" izz="0.0009196"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_hip_pitch_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0.01 0.05 0" rpy="0 1.57 0"/>
+      <geometry>
+        <cylinder radius="0.03" length="0.05"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_hip_pitch_joint" type="revolute">
+    <origin xyz="0 0.06445 -0.1027" rpy="0 -0.34907 0"/>
+    <parent link="pelvis"/>
+    <child link="left_hip_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-2.35" upper="3.05" effort="88" velocity="32"/>
+  </joint>
+  <link name="left_hip_roll_link">
+    <inertial>
+      <origin xyz="0.024757 -0.001036 -0.086323" rpy="0 0 0"/>
+      <mass value="1.446"/>
+      <inertia ixx="0.0022702" ixy="-3.7E-06" ixz="-0.0003789" iyy="0.002304" iyz="-1.8E-05" izz="0.0016001"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_hip_roll_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0.02 0 -0.1" rpy="0 0. 0"/>
+      <geometry>
+        <cylinder radius="0.04" length="0.05"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_hip_roll_joint" type="revolute">
+    <origin xyz="0 0.0523 0" rpy="0 0 0"/>
+    <parent link="left_hip_pitch_link"/>
+    <child link="left_hip_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-0.26" upper="2.53" effort="88" velocity="32"/>
+  </joint>
+  <link name="left_hip_yaw_link">
+    <inertial>
+      <origin xyz="-0.053554 -0.011477 -0.14067" rpy="0 0 0"/>
+      <mass value="2.052"/>
+      <inertia ixx="0.0087264" ixy="-0.0004402" ixz="-0.0036676" iyy="0.011374" iyz="-0.0006654" izz="0.004279"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_hip_yaw_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="-0.05 0 -0.14" rpy="0 0.5 0"/>
+      <geometry>
+        <cylinder radius="0.03" length="0.12"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_hip_yaw_joint" type="revolute">
+    <origin xyz="0.01966 -0.0012139 -0.1241" rpy="0 0 0"/>
+    <parent link="left_hip_roll_link"/>
+    <child link="left_hip_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.75" upper="2.75" effort="88" velocity="32"/>
+  </joint>
+  <link name="left_knee_link">
+    <inertial>
+      <origin xyz="0.005505 0.006534 -0.116629" rpy="0 0 0"/>
+      <mass value="2.252"/>
+      <inertia ixx="0.012443837" ixy="0.000053496" ixz="-0.000437641" iyy="0.012674902" iyz="-0.000682499" izz="0.001986501"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_knee_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0.007 0.005 -0.15" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.03" length="0.13"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_knee_joint" type="revolute">
+    <origin xyz="-0.078292 -0.0017335 -0.177225" rpy="0 0.5096 0"/>
+    <parent link="left_hip_yaw_link"/>
+    <child link="left_knee_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.33489" upper="2.5449" effort="139" velocity="20"/>
+  </joint>
+  <link name="left_ankle_pitch_link">
+    <inertial>
+      <origin xyz="-0.007269 0 0.011137" rpy="0 0 0"/>
+      <mass value="0.074"/>
+      <inertia ixx="8.4E-06" ixy="0" ixz="-2.9E-06" iyy="1.89E-05" iyz="0" izz="1.26E-05"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_ankle_pitch_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0.02" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.01" length="0.01"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_ankle_pitch_joint" type="revolute">
+    <origin xyz="0 0.0040687 -0.30007" rpy="0 -0.16053 0"/>
+    <parent link="left_knee_link"/>
+    <child link="left_ankle_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.68" upper="0.73" effort="40" velocity="53"/>
+  </joint>
+  <link name="left_ankle_roll_link">
+    <inertial>
+      <origin xyz="0.024762 2E-05 -0.012526" rpy="0 0 0"/>
+      <mass value="0.391"/>
+      <inertia ixx="0.0001552" ixy="-1E-07" ixz="7.51E-05" iyy="0.0010657" iyz="1E-07" izz="0.001098"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_ankle_roll_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0.035 0 -0.015" rpy="0 0 0"/>
+      <geometry>
+        <box size="0.19 0.04 0.03"/>
+      </geometry>
+    </collision>
+    <!-- <collision>
+      <origin xyz="-0.06 0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="-0.06 -0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="0.13 0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="0.13 -0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision> -->
+  </link>
+  <joint name="left_ankle_roll_joint" type="revolute">
+    <origin xyz="0 0 -0.017558" rpy="0 0 0"/>
+    <parent link="left_ankle_pitch_link"/>
+    <child link="left_ankle_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-0.2618" upper="0.2618" effort="40" velocity="53"/>
+  </joint>
+  <link name="right_hip_pitch_link">
+    <inertial>
+      <origin xyz="0.001962 -0.049392 -0.000941" rpy="0 0 0"/>
+      <mass value="1.299"/>
+      <inertia ixx="0.0013873" ixy="1.63E-05" ixz="-1E-06" iyy="0.0009059" iyz="4.24E-05" izz="0.0009196"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_hip_pitch_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0.01 -0.05 0" rpy="0 1.57 0"/>
+      <geometry>
+        <cylinder radius="0.03" length="0.05"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_hip_pitch_joint" type="revolute">
+    <origin xyz="0 -0.06445 -0.1027" rpy="0 -0.34907 0"/>
+    <parent link="pelvis"/>
+    <child link="right_hip_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-2.35" upper="3.05" effort="88" velocity="32"/>
+  </joint>
+  <link name="right_hip_roll_link">
+    <inertial>
+      <origin xyz="0.024757 0.001036 -0.086323" rpy="0 0 0"/>
+      <mass value="1.446"/>
+      <inertia ixx="0.0022702" ixy="3.7E-06" ixz="-0.0003789" iyy="0.002304" iyz="1.8E-05" izz="0.0016001"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_hip_roll_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0.02 0 -0.1" rpy="0 0. 0"/>
+      <geometry>
+        <cylinder radius="0.04" length="0.05"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_hip_roll_joint" type="revolute">
+    <origin xyz="0 -0.0523 0" rpy="0 0 0"/>
+    <parent link="right_hip_pitch_link"/>
+    <child link="right_hip_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-2.53" upper="0.26" effort="88" velocity="32"/>
+  </joint>
+  <link name="right_hip_yaw_link">
+    <inertial>
+      <origin xyz="-0.053554 0.011477 -0.14067" rpy="0 0 0"/>
+      <mass value="2.052"/>
+      <inertia ixx="0.0087264" ixy="0.0004402" ixz="-0.0036676" iyy="0.011374" iyz="0.0006654" izz="0.004279"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_hip_yaw_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="-0.05 0 -0.14" rpy="0 0.5 0"/>
+      <geometry>
+        <cylinder radius="0.03" length="0.12"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_hip_yaw_joint" type="revolute">
+    <origin xyz="0.01966 0.0012139 -0.1241" rpy="0 0 0"/>
+    <parent link="right_hip_roll_link"/>
+    <child link="right_hip_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.75" upper="2.75" effort="88" velocity="32"/>
+  </joint>
+  <link name="right_knee_link">
+    <inertial>
+      <origin xyz="0.005505 -0.006534 -0.116629" rpy="0 0 0"/>
+      <mass value="2.252"/>
+      <inertia ixx="0.012443837" ixy="-0.000053496" ixz="-0.000437641" iyy="0.012674902" iyz="0.000682499" izz="0.001986501"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_knee_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0.007 -0.005 -0.15" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.03" length="0.13"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_knee_joint" type="revolute">
+    <origin xyz="-0.078292 0.0017335 -0.177225" rpy="0 0.5096 0"/>
+    <parent link="right_hip_yaw_link"/>
+    <child link="right_knee_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.33489" upper="2.5449" effort="139" velocity="20"/>
+  </joint>
+  <link name="right_ankle_pitch_link">
+    <inertial>
+      <origin xyz="-0.007269 0 0.011137" rpy="0 0 0"/>
+      <mass value="0.074"/>
+      <inertia ixx="8.4E-06" ixy="0" ixz="-2.9E-06" iyy="1.89E-05" iyz="0" izz="1.26E-05"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_ankle_pitch_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0.02" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.01" length="0.01"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_ankle_pitch_joint" type="revolute">
+    <origin xyz="0 -0.0040687 -0.30007" rpy="0 -0.16053 0"/>
+    <parent link="right_knee_link"/>
+    <child link="right_ankle_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.68" upper="0.73" effort="40" velocity="53"/>
+  </joint>
+  <link name="right_ankle_roll_link">
+    <inertial>
+      <origin xyz="0.024762 -2E-05 -0.012526" rpy="0 0 0"/>
+      <mass value="0.391"/>
+      <inertia ixx="0.0001552" ixy="1E-07" ixz="7.51E-05" iyy="0.0010657" iyz="-1E-07" izz="0.001098"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_ankle_roll_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0.035 0 -0.015" rpy="0 0 0"/>
+      <geometry>
+        <box size="0.19 0.04 0.03"/>
+      </geometry>
+    </collision>
+      <!-- <geometry>
+        <mesh filename="../mesh/G1/right_ankle_roll_link.STL"/>
+      </geometry>
+    </collision> -->
+    <!-- <collision>
+      <origin xyz="-0.06 0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="-0.06 -0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="0.13 0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="0.13 -0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision> -->
+  </link>
+  <joint name="right_ankle_roll_joint" type="revolute">
+    <origin xyz="0 0 -0.017558" rpy="0 0 0"/>
+    <parent link="right_ankle_pitch_link"/>
+    <child link="right_ankle_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-0.2618" upper="0.2618" effort="40" velocity="53"/>
+  </joint>
+
+  <!-- Torso -->
+  <link name="torso_link">
+    <inertial>
+      <origin xyz="0.00197122283 0.00266902037 0.16936510960" rpy="0 0 0"/>
+      <mass value="6.33959811"/>
+      <inertia ixx="0.06037082109027" ixy="-0.00010168407055" ixz="0.00069845018092" iyy="0.04355153796551" iyz="0.00034366422728" izz="0.03339981866010"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/torso_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/torso_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="torso_joint" type="fixed">
+    <origin xyz="0 0 0" rpy="0 0 0"/>
+    <parent link="pelvis"/>
+    <child link="torso_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.618" upper="2.618" effort="88" velocity="32"/>
+  </joint>
+
+  <!-- Head -->
+  <link name="head_link">
+    <inertial>
+      <origin xyz="0.00138066852 0.00028430950 0.42034187824" rpy="0 0 0"/>
+      <mass value="1.17976522"/>
+      <inertia ixx="0.00543236042361" ixy="0.00000140137425" ixz="0.00034554752228" iyy="0.00552885306699" iyz="0.00001501216392" izz="0.00165378108136"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/head_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/head_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="head_joint" type="fixed">
+    <origin xyz="0 0 0" rpy="0 0 0"/>
+    <parent link="torso_link"/>
+    <child link="head_link"/>
+  </joint>
+
+  <!-- L ARM -->
+  <link name="left_shoulder_pitch_link">
+    <inertial>
+      <origin xyz="-0.001431 0.048811 0.001304" rpy="0 0 0"/>
+      <mass value="0.713"/>
+      <inertia ixx="0.0004614" ixy="-9.3E-06" ixz="1E-05" iyy="0.0004146" iyz="5.5E-06" izz="0.0004416"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_shoulder_pitch_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_shoulder_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_shoulder_pitch_joint" type="fixed">
+    <origin xyz="-0.0025 0.10396 0.25928" rpy="0.27925 0 0"/>
+    <parent link="torso_link"/>
+    <child link="left_shoulder_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-2.9671" upper="2.7925" effort="21" velocity="53"/>
+  </joint>
+  <link name="left_shoulder_roll_link">
+    <inertial>
+      <origin xyz="-0.003415 0.006955 -0.064598" rpy="0 0 0"/>
+      <mass value="0.642"/>
+      <inertia ixx="0.0006159" ixy="0" ixz="-5.6E-06" iyy="0.0006835" iyz="2.1E-06" izz="0.000373"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_shoulder_roll_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="-0.004 0.006 -0.053" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.03" length="0.03"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_shoulder_roll_joint" type="fixed">
+    <origin xyz="0 0.052 0" rpy="-0.27925 0 0"/>
+    <parent link="left_shoulder_pitch_link"/>
+    <child link="left_shoulder_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-1.5882" upper="2.2515" effort="21" velocity="53"/>
+  </joint>
+  <link name="left_shoulder_yaw_link">
+    <inertial>
+      <origin xyz="0.000375 -0.00444 -0.072374" rpy="0 0 0"/>
+      <mass value="0.713"/>
+      <inertia ixx="0.0009699" ixy="7.1E-06" ixz="7.8E-06" iyy="0.0009691" iyz="-4.49E-05" izz="0.0003826"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_shoulder_yaw_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_shoulder_yaw_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_shoulder_yaw_joint" type="fixed">
+    <origin xyz="-0.00354 0.0062424 -0.1032" rpy="0 0 0"/>
+    <parent link="left_shoulder_roll_link"/>
+    <child link="left_shoulder_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.618" upper="2.618" effort="21" velocity="53"/>
+  </joint>
+  <link name="left_elbow_pitch_link">
+    <inertial>
+      <origin xyz="0.064497 0.002873 0" rpy="0 0 0"/>
+      <mass value="0.601"/>
+      <inertia ixx="0.0002845" ixy="8.06E-05" ixz="9E-09" iyy="0.0004647" iyz="5E-09" izz="0.0004712"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_elbow_pitch_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_elbow_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_elbow_pitch_joint" type="fixed">
+    <origin xyz="0 0.00189 -0.0855" rpy="0 0 0"/>
+    <parent link="left_shoulder_yaw_link"/>
+    <child link="left_elbow_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.2268" upper="3.4208" effort="21" velocity="53"/>
+  </joint>
+  <link name="left_elbow_roll_link">
+    <inertial>
+      <origin xyz="0.081811 0.001454 0.001" rpy="0 0 0"/>
+      <mass value="0.227"/>
+      <inertia ixx="0.0001093" ixy="-1.55E-05" ixz="-2.91E-05" iyy="0.000977" iyz="1E-06" izz="0.0009676"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_elbow_roll_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_elbow_roll_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_elbow_roll_joint" type="fixed">
+    <origin xyz="0.1 0 0" rpy="0 0 0"/>
+    <parent link="left_elbow_pitch_link"/>
+    <child link="left_elbow_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-2.0943" upper="2.0943" effort="21" velocity="53"/>
+  </joint>
+
+  <!-- R ARM -->
+  <link name="right_shoulder_pitch_link">
+    <inertial>
+      <origin xyz="-0.001431 -0.048811 0.001304" rpy="0 0 0"/>
+      <mass value="0.713"/>
+      <inertia ixx="0.0004614" ixy="9.3E-06" ixz="1E-05" iyy="0.0004146" iyz="-5.5E-06" izz="0.0004416"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_shoulder_pitch_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_shoulder_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_shoulder_pitch_joint" type="fixed">
+    <origin xyz="-0.0025 -0.10396 0.25928" rpy="-0.27925 0 0"/>
+    <parent link="torso_link"/>
+    <child link="right_shoulder_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-2.9671" upper="2.7925" effort="21" velocity="53"/>
+  </joint>
+  <link name="right_shoulder_roll_link">
+    <inertial>
+      <origin xyz="-0.003415 -0.006955 -0.064598" rpy="0 0 0"/>
+      <mass value="0.642"/>
+      <inertia ixx="0.0006159" ixy="0" ixz="-5.6E-06" iyy="0.0006835" iyz="-2.1E-06" izz="0.000373"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_shoulder_roll_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="-0.004 -0.006 -0.053" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.03" length="0.03"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_shoulder_roll_joint" type="fixed">
+    <origin xyz="0 -0.052 0" rpy="0.27925 0 0"/>
+    <parent link="right_shoulder_pitch_link"/>
+    <child link="right_shoulder_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-2.2515" upper="1.5882" effort="21" velocity="53"/>
+  </joint>
+  <link name="right_shoulder_yaw_link">
+    <inertial>
+      <origin xyz="0.000375 0.00444 -0.072374" rpy="0 0 0"/>
+      <mass value="0.713"/>
+      <inertia ixx="0.0009699" ixy="-7.1E-06" ixz="7.8E-06" iyy="0.0009691" iyz="4.49E-05" izz="0.0003826"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_shoulder_yaw_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_shoulder_yaw_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_shoulder_yaw_joint" type="fixed">
+    <origin xyz="-0.00354 -0.0062424 -0.1032" rpy="0 0 0"/>
+    <parent link="right_shoulder_roll_link"/>
+    <child link="right_shoulder_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.618" upper="2.618" effort="21" velocity="53"/>
+  </joint>
+  <link name="right_elbow_pitch_link">
+    <inertial>
+      <origin xyz="0.064497 -0.002873 0" rpy="0 0 0"/>
+      <mass value="0.601"/>
+      <inertia ixx="0.0002845" ixy="-8.06E-05" ixz="9E-09" iyy="0.0004647" iyz="-5E-09" izz="0.0004712"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_elbow_pitch_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_elbow_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_elbow_pitch_joint" type="fixed">
+    <origin xyz="0 -0.00189 -0.0855" rpy="0 0 0"/>
+    <parent link="right_shoulder_yaw_link"/>
+    <child link="right_elbow_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.2268" upper="3.4208" effort="21" velocity="53"/>
+  </joint>
+  <link name="right_elbow_roll_link">
+    <inertial>
+      <origin xyz="0.081811 -0.001454 0.001" rpy="0 0 0"/>
+      <mass value="0.227"/>
+      <inertia ixx="0.0001093" ixy="1.55E-05" ixz="-2.91E-05" iyy="0.000977" iyz="-1E-06" izz="0.0009676"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_elbow_roll_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_elbow_roll_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_elbow_roll_joint" type="fixed">
+    <origin xyz="0.1 0 0" rpy="0 0 0"/>
+    <parent link="right_elbow_pitch_link"/>
+    <child link="right_elbow_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-2.0943" upper="2.0943" effort="21" velocity="53"/>
+  </joint>
+
+  <!-- LOGO -->
+  <joint name="logo_joint" type="fixed">
+    <origin xyz="0 0 0" rpy="0 0 0"/>
+    <parent link="torso_link"/>
+    <child link="logo_link"/>
+  </joint>
+  <link name="logo_link">
+    <inertial>
+      <origin xyz="-0.00306556489 0.00005625032 0.24526459523" rpy="0 0 0"/>
+      <mass value="0.001"/>
+      <inertia ixx="0.00000089650105" ixy="0.00000007028325" ixz="0.00000002735427" iyy="0.00000644846110" iyz="0.00000000029738" izz="0.00000688611155"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/logo_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/logo_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+
+  <!-- IMU -->
+  <link name="imu_link"></link>
+  <joint name="imu_joint" type="fixed">
+    <origin xyz="-0.04233868314 0.00166 0.152067" rpy="0 0 0"/>
+    <parent link="torso_link"/>
+    <child link="imu_link"/>
+  </joint>
+
+  <!-- Gripper -->
+  <joint name="left_palm_joint" type="fixed">
+    <origin xyz="0.120 0 0" rpy="0 0 0"/>
+    <parent link="left_elbow_roll_link"/>
+    <child link="left_palm_link"/>
+    <axis xyz="0 0 0"/>
+  </joint>
+  <link name="left_palm_link">
+    <inertial>
+      <origin xyz="0.05578422502 0.00148509387 -0.00032670143" rpy="0 0 0"/>
+      <mass value="0.28126043"/>
+      <inertia ixx="0.00017638101047" ixy="-0.00000652171140" ixz="0.00000021301775" iyy="0.00031084250366" iyz="0.00000000708018" izz="0.00018939350451"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_palm_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_palm_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_zero_joint" type="fixed">
+    <origin xyz="0.050 0 0" rpy="0 0 0"/>
+    <parent link="left_palm_link"/>
+    <child link="left_zero_link"/>
+    <axis xyz="0 1 0"/>
+    <limit effort="0.7" velocity="23" lower="-0.523598" upper="0.523598"/>
+  </joint>
+  <link name="left_zero_link">
+    <inertial>
+      <origin xyz="-0.02604664462 -0.00877511180 0.00001086048" rpy="0 0 0"/>
+      <mass value="0.05098927"/>
+      <inertia ixx="0.00000679821860" ixy="-0.00000098281498" ixz="-0.00000000406407" iyy="0.00000501631802" iyz="0.00000002833450" izz="0.00000868017156"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_zero_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <!-- <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_zero_link.STL"/>
+      </geometry>
+    </collision> -->
+  </link>
+  <joint name="left_one_joint" type="fixed">
+    <origin xyz="-0.02652502957 -0.0188 -0.00005" rpy="0 0 0"/>
+    <parent link="left_zero_link"/>
+    <child link="left_one_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1" upper="1.2"/>
+  </joint>
+  <link name="left_one_link">
+    <inertial>
+      <origin xyz="-0.00181538145 -0.03270571529 -0.00018587117" rpy="0 0 0"/>
+      <mass value="0.04776204"/>
+      <inertia ixx="0.00000800634866" ixy="-0.00000020811350" ixz="-0.00000004687654" iyy="0.00000366333777" iyz="0.00000005133915" izz="0.00000785913677"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_one_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="-0.001 -0.032 0" rpy="0 0 0"/>
+      <geometry>
+        <box size="0.02 0.03 0.02"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_two_joint" type="fixed">
+    <origin xyz="0 -0.0431 0" rpy="0 0 0"/>
+    <parent link="left_one_link"/>
+    <child link="left_two_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="0" upper="1.84"/>
+  </joint>
+  <link name="left_two_link">
+    <inertial>
+      <origin xyz="-0.00135257289 -0.02375910756 -0.00011966865" rpy="0 0 0"/>
+      <mass value="0.01385835"/>
+      <inertia ixx="0.00000254698095" ixy="-0.00000010086310" ixz="-0.00000000325710" iyy="0.00000094744556" iyz="0.00000003605469" izz="0.00000204288618"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_two_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_two_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_three_joint" type="fixed">
+    <origin xyz="0.085 0.004 0.02395" rpy="0 0 0"/>
+    <parent link="left_palm_link"/>
+    <child link="left_three_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1.84" upper="0.3"/>
+  </joint>
+  <link name="left_three_link">
+    <inertial>
+      <origin xyz="0.03625612321 -0.00157250459 0.00026989918" rpy="0 0 0"/>
+      <mass value="0.04884689"/>
+      <inertia ixx="0.00000383218469" ixy="0.00000003239984" ixz="0.00000001923377" iyy="0.00000890974599" iyz="-0.00000004429800" izz="0.00000865612355"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_three_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_three_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_four_joint" type="fixed">
+    <origin xyz="0.0471 -0.0036 0" rpy="0 0 0"/>
+    <parent link="left_three_link"/>
+    <child link="left_four_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1.84" upper="0"/>
+  </joint>
+  <link name="left_four_link">
+    <inertial>
+      <origin xyz="0.02375910756 0.00135257289 0.00011966865" rpy="0 0 0"/>
+      <mass value="0.01385835"/>
+      <inertia ixx="0.00000094744556" ixy="-0.00000010086310" ixz="0.00000003605469" iyy="0.00000254698095" iyz="-0.00000000325710" izz="0.00000204288618"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_four_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_four_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_five_joint" type="fixed">
+    <origin xyz="0.085 0.004 -0.02395" rpy="0 0 0"/>
+    <parent link="left_palm_link"/>
+    <child link="left_five_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1.84" upper="0.3"/>
+  </joint>
+  <link name="left_five_link">
+    <inertial>
+      <origin xyz="0.03625612321 -0.00157250459 0.00026989918" rpy="0 0 0"/>
+      <mass value="0.04884689"/>
+      <inertia ixx="0.00000383218469" ixy="0.00000003239984" ixz="0.00000001923377" iyy="0.00000890974599" iyz="-0.00000004429800" izz="0.00000865612355"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_five_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_five_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_six_joint" type="fixed">
+    <origin xyz="0.0471 -0.0036 0" rpy="0 0 0"/>
+    <parent link="left_five_link"/>
+    <child link="left_six_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1.84" upper="0"/>
+  </joint>
+  <link name="left_six_link">
+    <inertial>
+      <origin xyz="0.02375910756 0.00135257289 0.00011966865" rpy="0 0 0"/>
+      <mass value="0.01385835"/>
+      <inertia ixx="0.00000094744556" ixy="-0.00000010086310" ixz="0.00000003605469" iyy="0.00000254698095" iyz="-0.00000000325710" izz="0.00000204288618"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_six_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_six_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+
+
+  <joint name="right_palm_joint" type="fixed">
+    <origin xyz="0.120 0 0" rpy="0 0 0"/>
+    <parent link="right_elbow_roll_link"/>
+    <child link="right_palm_link"/>
+    <axis xyz="0 0 0"/>
+  </joint>
+  <link name="right_palm_link">
+    <inertial>
+      <origin xyz="0.05578422502 -0.00148509387 -0.00032670143" rpy="0 0 0"/>
+      <mass value="0.28126043"/>
+      <inertia ixx="0.00017638101047" ixy="0.00000652171140" ixz="0.00000021301775" iyy="0.00031084250366" iyz="-0.00000000708018" izz="0.00018939350451"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_palm_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_palm_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_zero_joint" type="fixed">
+    <origin xyz="0.050 0 0" rpy="0 0 0"/>
+    <parent link="right_palm_link"/>
+    <child link="right_zero_link"/>
+    <axis xyz="0 1 0"/>
+    <limit effort="0.7" velocity="23" lower="-0.523598" upper="0.523598"/>
+  </joint>
+  <link name="right_zero_link">
+    <inertial>
+      <origin xyz="-0.02604664462 0.00877511180 0.00001086048" rpy="0 0 0"/>
+      <mass value="0.05098927"/>
+      <inertia ixx="0.00000679821860" ixy="0.00000098281498" ixz="-0.00000000406407" iyy="0.00000501631802" iyz="-0.00000002833450" izz="0.00000868017156"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_zero_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <!-- <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_zero_link.STL"/>
+      </geometry>
+    </collision> -->
+  </link>
+  <joint name="right_one_joint" type="fixed">
+    <origin xyz="-0.02652502957 0.0188 -0.00005" rpy="0 0 0"/>
+    <parent link="right_zero_link"/>
+    <child link="right_one_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1.2" upper="1"/>
+  </joint>
+  <link name="right_one_link">
+    <inertial>
+      <origin xyz="-0.00181538145 0.03270571529 -0.00018587117" rpy="0 0 0"/>
+      <mass value="0.04776204"/>
+      <inertia ixx="0.00000800634866" ixy="0.00000020811350" ixz="-0.00000004687654" iyy="0.00000366333777" iyz="-0.00000005133915" izz="0.00000785913677"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_one_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="-0.001 0.032 0" rpy="0 0 0"/>
+      <geometry>
+        <box size="0.02 0.03 0.02"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_two_joint" type="fixed">
+    <origin xyz="0 0.0431 0" rpy="0 0 0"/>
+    <parent link="right_one_link"/>
+    <child link="right_two_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1.84" upper="0"/>
+  </joint>
+  <link name="right_two_link">
+    <inertial>
+      <origin xyz="-0.00135257289 0.02375910756 -0.00011966865" rpy="0 0 0"/>
+      <mass value="0.01385835"/>
+      <inertia ixx="0.00000254698095" ixy="0.00000010086310" ixz="-0.00000000325710" iyy="0.00000094744556" iyz="-0.00000003605469" izz="0.00000204288618"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_two_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_two_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_three_joint" type="fixed">
+    <origin xyz="0.085 -0.004 0.02395" rpy="0 0 0"/>
+    <parent link="right_palm_link"/>
+    <child link="right_three_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-0.3" upper="1.84"/>
+  </joint>
+  <link name="right_three_link">
+    <inertial>
+      <origin xyz="0.03625612321 0.00157250459 0.00026989918" rpy="0 0 0"/>
+      <mass value="0.04884689"/>
+      <inertia ixx="0.00000383218469" ixy="-0.00000003239984" ixz="0.00000001923377" iyy="0.00000890974599" iyz="0.00000004429800" izz="0.00000865612355"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_three_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_three_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_four_joint" type="fixed">
+    <origin xyz="0.0471 0.0036 0" rpy="0 0 0"/>
+    <parent link="right_three_link"/>
+    <child link="right_four_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="0" upper="1.84"/>
+  </joint>
+  <link name="right_four_link">
+    <inertial>
+      <origin xyz="0.02375910756 -0.00135257289 0.00011966865" rpy="0 0 0"/>
+      <mass value="0.01385835"/>
+      <inertia ixx="0.00000094744556" ixy="0.00000010086310" ixz="0.00000003605469" iyy="0.00000254698095" iyz="0.00000000325710" izz="0.00000204288618"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_four_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_four_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_five_joint" type="fixed">
+    <origin xyz="0.085 -0.004 -0.02395" rpy="0 0 0"/>
+    <parent link="right_palm_link"/>
+    <child link="right_five_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-0.3" upper="1.84"/>
+  </joint>
+  <link name="right_five_link">
+    <inertial>
+      <origin xyz="0.03625612321 0.00157250459 0.00026989918" rpy="0 0 0"/>
+      <mass value="0.04884689"/>
+      <inertia ixx="0.00000383218469" ixy="-0.00000003239984" ixz="0.00000001923377" iyy="0.00000890974599" iyz="0.00000004429800" izz="0.00000865612355"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_five_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_five_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_six_joint" type="fixed">
+    <origin xyz="0.0471 0.0036 0" rpy="0 0 0"/>
+    <parent link="right_five_link"/>
+    <child link="right_six_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="0" upper="1.84"/>
+  </joint>
+  <link name="right_six_link">
+    <inertial>
+      <origin xyz="0.02375910756 -0.00135257289 0.00011966865" rpy="0 0 0"/>
+      <mass value="0.01385835"/>
+      <inertia ixx="0.00000094744556" ixy="0.00000010086310" ixz="0.00000003605469" iyy="0.00000254698095" iyz="0.00000000325710" izz="0.00000204288618"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_six_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_six_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+</robot>
\ No newline at end of file
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/g1_12dof.urdf b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/g1_12dof.urdf
new file mode 100644
index 0000000000000000000000000000000000000000..5eaeba7203543090ea4a39d97a59c53ef98ee668
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/g1_12dof.urdf
@@ -0,0 +1,1295 @@
+<robot name="g1">
+
+  <link name="pelvis">
+    <inertial>
+      <origin xyz="0 0 -0.07605" rpy="0 0 0"/>
+      <mass value="2.86"/>
+      <inertia ixx="0.0079143" ixy="0" ixz="1.6E-06" iyy="0.0069837" iyz="0" izz="0.0059404"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/pelvis.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/pelvis.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="pelvis_contour_link">
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/pelvis_contour_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+  </link>
+  <joint name="pelvis_contour_joint" type="fixed">
+    <parent link="pelvis"/>
+    <child link="pelvis_contour_link"/>
+  </joint>
+
+  <!-- Leg -->
+  <link name="left_hip_pitch_link">
+    <inertial>
+      <origin xyz="0.001962 0.049392 -0.000941" rpy="0 0 0"/>
+      <mass value="1.299"/>
+      <inertia ixx="0.0013873" ixy="-1.63E-05" ixz="-1E-06" iyy="0.0009059" iyz="-4.24E-05" izz="0.0009196"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_hip_pitch_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_hip_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_hip_pitch_joint" type="revolute">
+    <origin xyz="0 0.06445 -0.1027" rpy="0 -0.34907 0"/>
+    <parent link="pelvis"/>
+    <child link="left_hip_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-2.35" upper="3.05" effort="88" velocity="32"/>
+  </joint>
+  <link name="left_hip_roll_link">
+    <inertial>
+      <origin xyz="0.024757 -0.001036 -0.086323" rpy="0 0 0"/>
+      <mass value="1.446"/>
+      <inertia ixx="0.0022702" ixy="-3.7E-06" ixz="-0.0003789" iyy="0.002304" iyz="-1.8E-05" izz="0.0016001"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_hip_roll_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_hip_roll_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_hip_roll_joint" type="revolute">
+    <origin xyz="0 0.0523 0" rpy="0 0 0"/>
+    <parent link="left_hip_pitch_link"/>
+    <child link="left_hip_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-0.26" upper="2.53" effort="88" velocity="32"/>
+  </joint>
+  <link name="left_hip_yaw_link">
+    <inertial>
+      <origin xyz="-0.053554 -0.011477 -0.14067" rpy="0 0 0"/>
+      <mass value="2.052"/>
+      <inertia ixx="0.0087264" ixy="-0.0004402" ixz="-0.0036676" iyy="0.011374" iyz="-0.0006654" izz="0.004279"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_hip_yaw_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_hip_yaw_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_hip_yaw_joint" type="revolute">
+    <origin xyz="0.01966 -0.0012139 -0.1241" rpy="0 0 0"/>
+    <parent link="left_hip_roll_link"/>
+    <child link="left_hip_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.75" upper="2.75" effort="88" velocity="32"/>
+  </joint>
+  <link name="left_knee_link">
+    <inertial>
+      <origin xyz="0.005505 0.006534 -0.116629" rpy="0 0 0"/>
+      <mass value="2.252"/>
+      <inertia ixx="0.012443837" ixy="0.000053496" ixz="-0.000437641" iyy="0.012674902" iyz="-0.000682499" izz="0.001986501"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_knee_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0.007 0.005 -0.15" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.015" length="0.15"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_knee_joint" type="revolute">
+    <origin xyz="-0.078292 -0.0017335 -0.177225" rpy="0 0.5096 0"/>
+    <parent link="left_hip_yaw_link"/>
+    <child link="left_knee_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.33489" upper="2.5449" effort="139" velocity="20"/>
+  </joint>
+  <link name="left_ankle_pitch_link">
+    <inertial>
+      <origin xyz="-0.007269 0 0.011137" rpy="0 0 0"/>
+      <mass value="0.074"/>
+      <inertia ixx="8.4E-06" ixy="0" ixz="-2.9E-06" iyy="1.89E-05" iyz="0" izz="1.26E-05"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_ankle_pitch_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_ankle_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_ankle_pitch_joint" type="revolute">
+    <origin xyz="0 0.0040687 -0.30007" rpy="0 -0.16053 0"/>
+    <parent link="left_knee_link"/>
+    <child link="left_ankle_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.68" upper="0.73" effort="40" velocity="53"/>
+  </joint>
+  <link name="left_ankle_roll_link">
+    <inertial>
+      <origin xyz="0.024762 2E-05 -0.012526" rpy="0 0 0"/>
+      <mass value="0.391"/>
+      <inertia ixx="0.0001552" ixy="-1E-07" ixz="7.51E-05" iyy="0.0010657" iyz="1E-07" izz="0.001098"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_ankle_roll_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_ankle_roll_link.STL"/>
+      </geometry>
+    </collision>
+    <!-- <collision>
+      <origin xyz="-0.06 0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="-0.06 -0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="0.13 0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="0.13 -0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision> -->
+  </link>
+  <joint name="left_ankle_roll_joint" type="revolute">
+    <origin xyz="0 0 -0.017558" rpy="0 0 0"/>
+    <parent link="left_ankle_pitch_link"/>
+    <child link="left_ankle_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-0.2618" upper="0.2618" effort="40" velocity="53"/>
+  </joint>
+  <link name="right_hip_pitch_link">
+    <inertial>
+      <origin xyz="0.001962 -0.049392 -0.000941" rpy="0 0 0"/>
+      <mass value="1.299"/>
+      <inertia ixx="0.0013873" ixy="1.63E-05" ixz="-1E-06" iyy="0.0009059" iyz="4.24E-05" izz="0.0009196"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_hip_pitch_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_hip_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_hip_pitch_joint" type="revolute">
+    <origin xyz="0 -0.06445 -0.1027" rpy="0 -0.34907 0"/>
+    <parent link="pelvis"/>
+    <child link="right_hip_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-2.35" upper="3.05" effort="88" velocity="32"/>
+  </joint>
+  <link name="right_hip_roll_link">
+    <inertial>
+      <origin xyz="0.024757 0.001036 -0.086323" rpy="0 0 0"/>
+      <mass value="1.446"/>
+      <inertia ixx="0.0022702" ixy="3.7E-06" ixz="-0.0003789" iyy="0.002304" iyz="1.8E-05" izz="0.0016001"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_hip_roll_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_hip_roll_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_hip_roll_joint" type="revolute">
+    <origin xyz="0 -0.0523 0" rpy="0 0 0"/>
+    <parent link="right_hip_pitch_link"/>
+    <child link="right_hip_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-2.53" upper="0.26" effort="88" velocity="32"/>
+  </joint>
+  <link name="right_hip_yaw_link">
+    <inertial>
+      <origin xyz="-0.053554 0.011477 -0.14067" rpy="0 0 0"/>
+      <mass value="2.052"/>
+      <inertia ixx="0.0087264" ixy="0.0004402" ixz="-0.0036676" iyy="0.011374" iyz="0.0006654" izz="0.004279"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_hip_yaw_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_hip_yaw_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_hip_yaw_joint" type="revolute">
+    <origin xyz="0.01966 0.0012139 -0.1241" rpy="0 0 0"/>
+    <parent link="right_hip_roll_link"/>
+    <child link="right_hip_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.75" upper="2.75" effort="88" velocity="32"/>
+  </joint>
+  <link name="right_knee_link">
+    <inertial>
+      <origin xyz="0.005505 -0.006534 -0.116629" rpy="0 0 0"/>
+      <mass value="2.252"/>
+      <inertia ixx="0.012443837" ixy="-0.000053496" ixz="-0.000437641" iyy="0.012674902" iyz="0.000682499" izz="0.001986501"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_knee_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0.007 -0.005 -0.15" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.015" length="0.15"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_knee_joint" type="revolute">
+    <origin xyz="-0.078292 0.0017335 -0.177225" rpy="0 0.5096 0"/>
+    <parent link="right_hip_yaw_link"/>
+    <child link="right_knee_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.33489" upper="2.5449" effort="139" velocity="20"/>
+  </joint>
+  <link name="right_ankle_pitch_link">
+    <inertial>
+      <origin xyz="-0.007269 0 0.011137" rpy="0 0 0"/>
+      <mass value="0.074"/>
+      <inertia ixx="8.4E-06" ixy="0" ixz="-2.9E-06" iyy="1.89E-05" iyz="0" izz="1.26E-05"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_ankle_pitch_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_ankle_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_ankle_pitch_joint" type="revolute">
+    <origin xyz="0 -0.0040687 -0.30007" rpy="0 -0.16053 0"/>
+    <parent link="right_knee_link"/>
+    <child link="right_ankle_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.68" upper="0.73" effort="40" velocity="53"/>
+  </joint>
+  <link name="right_ankle_roll_link">
+    <inertial>
+      <origin xyz="0.024762 -2E-05 -0.012526" rpy="0 0 0"/>
+      <mass value="0.391"/>
+      <inertia ixx="0.0001552" ixy="1E-07" ixz="7.51E-05" iyy="0.0010657" iyz="-1E-07" izz="0.001098"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_ankle_roll_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_ankle_roll_link.STL"/>
+      </geometry>
+    </collision>
+    <!-- <collision>
+      <origin xyz="-0.06 0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="-0.06 -0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="0.13 0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="0.13 -0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision> -->
+  </link>
+  <joint name="right_ankle_roll_joint" type="revolute">
+    <origin xyz="0 0 -0.017558" rpy="0 0 0"/>
+    <parent link="right_ankle_pitch_link"/>
+    <child link="right_ankle_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-0.2618" upper="0.2618" effort="40" velocity="53"/>
+  </joint>
+
+  <!-- Torso -->
+  <link name="torso_link">
+    <inertial>
+      <origin xyz="0.00197122283 0.00266902037 0.16936510960" rpy="0 0 0"/>
+      <mass value="6.33959811"/>
+      <inertia ixx="0.06037082109027" ixy="-0.00010168407055" ixz="0.00069845018092" iyy="0.04355153796551" iyz="0.00034366422728" izz="0.03339981866010"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/torso_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/torso_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="torso_joint" type="fixed">
+    <origin xyz="0 0 0" rpy="0 0 0"/>
+    <parent link="pelvis"/>
+    <child link="torso_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.618" upper="2.618" effort="88" velocity="32"/>
+  </joint>
+
+  <!-- Head -->
+  <link name="head_link">
+    <inertial>
+      <origin xyz="0.00138066852 0.00028430950 0.42034187824" rpy="0 0 0"/>
+      <mass value="1.17976522"/>
+      <inertia ixx="0.00543236042361" ixy="0.00000140137425" ixz="0.00034554752228" iyy="0.00552885306699" iyz="0.00001501216392" izz="0.00165378108136"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/head_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/head_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="head_joint" type="fixed">
+    <origin xyz="0 0 0" rpy="0 0 0"/>
+    <parent link="torso_link"/>
+    <child link="head_link"/>
+  </joint>
+
+  <!-- L ARM -->
+  <link name="left_shoulder_pitch_link">
+    <inertial>
+      <origin xyz="-0.001431 0.048811 0.001304" rpy="0 0 0"/>
+      <mass value="0.713"/>
+      <inertia ixx="0.0004614" ixy="-9.3E-06" ixz="1E-05" iyy="0.0004146" iyz="5.5E-06" izz="0.0004416"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_shoulder_pitch_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_shoulder_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_shoulder_pitch_joint" type="fixed">
+    <origin xyz="-0.0025 0.10396 0.25928" rpy="0.27925 0 0"/>
+    <parent link="torso_link"/>
+    <child link="left_shoulder_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-2.9671" upper="2.7925" effort="21" velocity="53"/>
+  </joint>
+  <link name="left_shoulder_roll_link">
+    <inertial>
+      <origin xyz="-0.003415 0.006955 -0.064598" rpy="0 0 0"/>
+      <mass value="0.642"/>
+      <inertia ixx="0.0006159" ixy="0" ixz="-5.6E-06" iyy="0.0006835" iyz="2.1E-06" izz="0.000373"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_shoulder_roll_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="-0.004 0.006 -0.053" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.03" length="0.03"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_shoulder_roll_joint" type="fixed">
+    <origin xyz="0 0.052 0" rpy="-0.27925 0 0"/>
+    <parent link="left_shoulder_pitch_link"/>
+    <child link="left_shoulder_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-1.5882" upper="2.2515" effort="21" velocity="53"/>
+  </joint>
+  <link name="left_shoulder_yaw_link">
+    <inertial>
+      <origin xyz="0.000375 -0.00444 -0.072374" rpy="0 0 0"/>
+      <mass value="0.713"/>
+      <inertia ixx="0.0009699" ixy="7.1E-06" ixz="7.8E-06" iyy="0.0009691" iyz="-4.49E-05" izz="0.0003826"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_shoulder_yaw_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_shoulder_yaw_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_shoulder_yaw_joint" type="fixed">
+    <origin xyz="-0.00354 0.0062424 -0.1032" rpy="0 0 0"/>
+    <parent link="left_shoulder_roll_link"/>
+    <child link="left_shoulder_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.618" upper="2.618" effort="21" velocity="53"/>
+  </joint>
+  <link name="left_elbow_pitch_link">
+    <inertial>
+      <origin xyz="0.064497 0.002873 0" rpy="0 0 0"/>
+      <mass value="0.601"/>
+      <inertia ixx="0.0002845" ixy="8.06E-05" ixz="9E-09" iyy="0.0004647" iyz="5E-09" izz="0.0004712"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_elbow_pitch_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_elbow_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_elbow_pitch_joint" type="fixed">
+    <origin xyz="0 0.00189 -0.0855" rpy="0 0 0"/>
+    <parent link="left_shoulder_yaw_link"/>
+    <child link="left_elbow_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.2268" upper="3.4208" effort="21" velocity="53"/>
+  </joint>
+  <link name="left_elbow_roll_link">
+    <inertial>
+      <origin xyz="0.081811 0.001454 0.001" rpy="0 0 0"/>
+      <mass value="0.227"/>
+      <inertia ixx="0.0001093" ixy="-1.55E-05" ixz="-2.91E-05" iyy="0.000977" iyz="1E-06" izz="0.0009676"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_elbow_roll_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_elbow_roll_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_elbow_roll_joint" type="fixed">
+    <origin xyz="0.1 0 0" rpy="0 0 0"/>
+    <parent link="left_elbow_pitch_link"/>
+    <child link="left_elbow_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-2.0943" upper="2.0943" effort="21" velocity="53"/>
+  </joint>
+
+  <!-- R ARM -->
+  <link name="right_shoulder_pitch_link">
+    <inertial>
+      <origin xyz="-0.001431 -0.048811 0.001304" rpy="0 0 0"/>
+      <mass value="0.713"/>
+      <inertia ixx="0.0004614" ixy="9.3E-06" ixz="1E-05" iyy="0.0004146" iyz="-5.5E-06" izz="0.0004416"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_shoulder_pitch_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_shoulder_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_shoulder_pitch_joint" type="fixed">
+    <origin xyz="-0.0025 -0.10396 0.25928" rpy="-0.27925 0 0"/>
+    <parent link="torso_link"/>
+    <child link="right_shoulder_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-2.9671" upper="2.7925" effort="21" velocity="53"/>
+  </joint>
+  <link name="right_shoulder_roll_link">
+    <inertial>
+      <origin xyz="-0.003415 -0.006955 -0.064598" rpy="0 0 0"/>
+      <mass value="0.642"/>
+      <inertia ixx="0.0006159" ixy="0" ixz="-5.6E-06" iyy="0.0006835" iyz="-2.1E-06" izz="0.000373"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_shoulder_roll_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="-0.004 -0.006 -0.053" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.03" length="0.03"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_shoulder_roll_joint" type="fixed">
+    <origin xyz="0 -0.052 0" rpy="0.27925 0 0"/>
+    <parent link="right_shoulder_pitch_link"/>
+    <child link="right_shoulder_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-2.2515" upper="1.5882" effort="21" velocity="53"/>
+  </joint>
+  <link name="right_shoulder_yaw_link">
+    <inertial>
+      <origin xyz="0.000375 0.00444 -0.072374" rpy="0 0 0"/>
+      <mass value="0.713"/>
+      <inertia ixx="0.0009699" ixy="-7.1E-06" ixz="7.8E-06" iyy="0.0009691" iyz="4.49E-05" izz="0.0003826"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_shoulder_yaw_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_shoulder_yaw_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_shoulder_yaw_joint" type="fixed">
+    <origin xyz="-0.00354 -0.0062424 -0.1032" rpy="0 0 0"/>
+    <parent link="right_shoulder_roll_link"/>
+    <child link="right_shoulder_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.618" upper="2.618" effort="21" velocity="53"/>
+  </joint>
+  <link name="right_elbow_pitch_link">
+    <inertial>
+      <origin xyz="0.064497 -0.002873 0" rpy="0 0 0"/>
+      <mass value="0.601"/>
+      <inertia ixx="0.0002845" ixy="-8.06E-05" ixz="9E-09" iyy="0.0004647" iyz="-5E-09" izz="0.0004712"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_elbow_pitch_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_elbow_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_elbow_pitch_joint" type="fixed">
+    <origin xyz="0 -0.00189 -0.0855" rpy="0 0 0"/>
+    <parent link="right_shoulder_yaw_link"/>
+    <child link="right_elbow_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.2268" upper="3.4208" effort="21" velocity="53"/>
+  </joint>
+  <link name="right_elbow_roll_link">
+    <inertial>
+      <origin xyz="0.081811 -0.001454 0.001" rpy="0 0 0"/>
+      <mass value="0.227"/>
+      <inertia ixx="0.0001093" ixy="1.55E-05" ixz="-2.91E-05" iyy="0.000977" iyz="-1E-06" izz="0.0009676"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_elbow_roll_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_elbow_roll_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_elbow_roll_joint" type="fixed">
+    <origin xyz="0.1 0 0" rpy="0 0 0"/>
+    <parent link="right_elbow_pitch_link"/>
+    <child link="right_elbow_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-2.0943" upper="2.0943" effort="21" velocity="53"/>
+  </joint>
+
+  <!-- LOGO -->
+  <joint name="logo_joint" type="fixed">
+    <origin xyz="0 0 0" rpy="0 0 0"/>
+    <parent link="torso_link"/>
+    <child link="logo_link"/>
+  </joint>
+  <link name="logo_link">
+    <inertial>
+      <origin xyz="-0.00306556489 0.00005625032 0.24526459523" rpy="0 0 0"/>
+      <mass value="0.001"/>
+      <inertia ixx="0.00000089650105" ixy="0.00000007028325" ixz="0.00000002735427" iyy="0.00000644846110" iyz="0.00000000029738" izz="0.00000688611155"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/logo_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/logo_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+
+  <!-- IMU -->
+  <link name="imu_link"></link>
+  <joint name="imu_joint" type="fixed">
+    <origin xyz="-0.04233868314 0.00166 0.152067" rpy="0 0 0"/>
+    <parent link="torso_link"/>
+    <child link="imu_link"/>
+  </joint>
+
+  <!-- Gripper -->
+  <joint name="left_palm_joint" type="fixed">
+    <origin xyz="0.120 0 0" rpy="0 0 0"/>
+    <parent link="left_elbow_roll_link"/>
+    <child link="left_palm_link"/>
+    <axis xyz="0 0 0"/>
+  </joint>
+  <link name="left_palm_link">
+    <inertial>
+      <origin xyz="0.05578422502 0.00148509387 -0.00032670143" rpy="0 0 0"/>
+      <mass value="0.28126043"/>
+      <inertia ixx="0.00017638101047" ixy="-0.00000652171140" ixz="0.00000021301775" iyy="0.00031084250366" iyz="0.00000000708018" izz="0.00018939350451"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_palm_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_palm_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_zero_joint" type="fixed">
+    <origin xyz="0.050 0 0" rpy="0 0 0"/>
+    <parent link="left_palm_link"/>
+    <child link="left_zero_link"/>
+    <axis xyz="0 1 0"/>
+    <limit effort="0.7" velocity="23" lower="-0.523598" upper="0.523598"/>
+  </joint>
+  <link name="left_zero_link">
+    <inertial>
+      <origin xyz="-0.02604664462 -0.00877511180 0.00001086048" rpy="0 0 0"/>
+      <mass value="0.05098927"/>
+      <inertia ixx="0.00000679821860" ixy="-0.00000098281498" ixz="-0.00000000406407" iyy="0.00000501631802" iyz="0.00000002833450" izz="0.00000868017156"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_zero_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <!-- <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_zero_link.STL"/>
+      </geometry>
+    </collision> -->
+  </link>
+  <joint name="left_one_joint" type="fixed">
+    <origin xyz="-0.02652502957 -0.0188 -0.00005" rpy="0 0 0"/>
+    <parent link="left_zero_link"/>
+    <child link="left_one_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1" upper="1.2"/>
+  </joint>
+  <link name="left_one_link">
+    <inertial>
+      <origin xyz="-0.00181538145 -0.03270571529 -0.00018587117" rpy="0 0 0"/>
+      <mass value="0.04776204"/>
+      <inertia ixx="0.00000800634866" ixy="-0.00000020811350" ixz="-0.00000004687654" iyy="0.00000366333777" iyz="0.00000005133915" izz="0.00000785913677"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_one_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="-0.001 -0.032 0" rpy="0 0 0"/>
+      <geometry>
+        <box size="0.02 0.03 0.02"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_two_joint" type="fixed">
+    <origin xyz="0 -0.0431 0" rpy="0 0 0"/>
+    <parent link="left_one_link"/>
+    <child link="left_two_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="0" upper="1.84"/>
+  </joint>
+  <link name="left_two_link">
+    <inertial>
+      <origin xyz="-0.00135257289 -0.02375910756 -0.00011966865" rpy="0 0 0"/>
+      <mass value="0.01385835"/>
+      <inertia ixx="0.00000254698095" ixy="-0.00000010086310" ixz="-0.00000000325710" iyy="0.00000094744556" iyz="0.00000003605469" izz="0.00000204288618"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_two_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_two_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_three_joint" type="fixed">
+    <origin xyz="0.085 0.004 0.02395" rpy="0 0 0"/>
+    <parent link="left_palm_link"/>
+    <child link="left_three_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1.84" upper="0.3"/>
+  </joint>
+  <link name="left_three_link">
+    <inertial>
+      <origin xyz="0.03625612321 -0.00157250459 0.00026989918" rpy="0 0 0"/>
+      <mass value="0.04884689"/>
+      <inertia ixx="0.00000383218469" ixy="0.00000003239984" ixz="0.00000001923377" iyy="0.00000890974599" iyz="-0.00000004429800" izz="0.00000865612355"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_three_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_three_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_four_joint" type="fixed">
+    <origin xyz="0.0471 -0.0036 0" rpy="0 0 0"/>
+    <parent link="left_three_link"/>
+    <child link="left_four_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1.84" upper="0"/>
+  </joint>
+  <link name="left_four_link">
+    <inertial>
+      <origin xyz="0.02375910756 0.00135257289 0.00011966865" rpy="0 0 0"/>
+      <mass value="0.01385835"/>
+      <inertia ixx="0.00000094744556" ixy="-0.00000010086310" ixz="0.00000003605469" iyy="0.00000254698095" iyz="-0.00000000325710" izz="0.00000204288618"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_four_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_four_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_five_joint" type="fixed">
+    <origin xyz="0.085 0.004 -0.02395" rpy="0 0 0"/>
+    <parent link="left_palm_link"/>
+    <child link="left_five_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1.84" upper="0.3"/>
+  </joint>
+  <link name="left_five_link">
+    <inertial>
+      <origin xyz="0.03625612321 -0.00157250459 0.00026989918" rpy="0 0 0"/>
+      <mass value="0.04884689"/>
+      <inertia ixx="0.00000383218469" ixy="0.00000003239984" ixz="0.00000001923377" iyy="0.00000890974599" iyz="-0.00000004429800" izz="0.00000865612355"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_five_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_five_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_six_joint" type="fixed">
+    <origin xyz="0.0471 -0.0036 0" rpy="0 0 0"/>
+    <parent link="left_five_link"/>
+    <child link="left_six_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1.84" upper="0"/>
+  </joint>
+  <link name="left_six_link">
+    <inertial>
+      <origin xyz="0.02375910756 0.00135257289 0.00011966865" rpy="0 0 0"/>
+      <mass value="0.01385835"/>
+      <inertia ixx="0.00000094744556" ixy="-0.00000010086310" ixz="0.00000003605469" iyy="0.00000254698095" iyz="-0.00000000325710" izz="0.00000204288618"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_six_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/left_six_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+
+
+  <joint name="right_palm_joint" type="fixed">
+    <origin xyz="0.120 0 0" rpy="0 0 0"/>
+    <parent link="right_elbow_roll_link"/>
+    <child link="right_palm_link"/>
+    <axis xyz="0 0 0"/>
+  </joint>
+  <link name="right_palm_link">
+    <inertial>
+      <origin xyz="0.05578422502 -0.00148509387 -0.00032670143" rpy="0 0 0"/>
+      <mass value="0.28126043"/>
+      <inertia ixx="0.00017638101047" ixy="0.00000652171140" ixz="0.00000021301775" iyy="0.00031084250366" iyz="-0.00000000708018" izz="0.00018939350451"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_palm_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_palm_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_zero_joint" type="fixed">
+    <origin xyz="0.050 0 0" rpy="0 0 0"/>
+    <parent link="right_palm_link"/>
+    <child link="right_zero_link"/>
+    <axis xyz="0 1 0"/>
+    <limit effort="0.7" velocity="23" lower="-0.523598" upper="0.523598"/>
+  </joint>
+  <link name="right_zero_link">
+    <inertial>
+      <origin xyz="-0.02604664462 0.00877511180 0.00001086048" rpy="0 0 0"/>
+      <mass value="0.05098927"/>
+      <inertia ixx="0.00000679821860" ixy="0.00000098281498" ixz="-0.00000000406407" iyy="0.00000501631802" iyz="-0.00000002833450" izz="0.00000868017156"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_zero_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <!-- <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_zero_link.STL"/>
+      </geometry>
+    </collision> -->
+  </link>
+  <joint name="right_one_joint" type="fixed">
+    <origin xyz="-0.02652502957 0.0188 -0.00005" rpy="0 0 0"/>
+    <parent link="right_zero_link"/>
+    <child link="right_one_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1.2" upper="1"/>
+  </joint>
+  <link name="right_one_link">
+    <inertial>
+      <origin xyz="-0.00181538145 0.03270571529 -0.00018587117" rpy="0 0 0"/>
+      <mass value="0.04776204"/>
+      <inertia ixx="0.00000800634866" ixy="0.00000020811350" ixz="-0.00000004687654" iyy="0.00000366333777" iyz="-0.00000005133915" izz="0.00000785913677"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_one_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="-0.001 0.032 0" rpy="0 0 0"/>
+      <geometry>
+        <box size="0.02 0.03 0.02"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_two_joint" type="fixed">
+    <origin xyz="0 0.0431 0" rpy="0 0 0"/>
+    <parent link="right_one_link"/>
+    <child link="right_two_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1.84" upper="0"/>
+  </joint>
+  <link name="right_two_link">
+    <inertial>
+      <origin xyz="-0.00135257289 0.02375910756 -0.00011966865" rpy="0 0 0"/>
+      <mass value="0.01385835"/>
+      <inertia ixx="0.00000254698095" ixy="0.00000010086310" ixz="-0.00000000325710" iyy="0.00000094744556" iyz="-0.00000003605469" izz="0.00000204288618"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_two_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_two_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_three_joint" type="fixed">
+    <origin xyz="0.085 -0.004 0.02395" rpy="0 0 0"/>
+    <parent link="right_palm_link"/>
+    <child link="right_three_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-0.3" upper="1.84"/>
+  </joint>
+  <link name="right_three_link">
+    <inertial>
+      <origin xyz="0.03625612321 0.00157250459 0.00026989918" rpy="0 0 0"/>
+      <mass value="0.04884689"/>
+      <inertia ixx="0.00000383218469" ixy="-0.00000003239984" ixz="0.00000001923377" iyy="0.00000890974599" iyz="0.00000004429800" izz="0.00000865612355"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_three_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_three_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_four_joint" type="fixed">
+    <origin xyz="0.0471 0.0036 0" rpy="0 0 0"/>
+    <parent link="right_three_link"/>
+    <child link="right_four_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="0" upper="1.84"/>
+  </joint>
+  <link name="right_four_link">
+    <inertial>
+      <origin xyz="0.02375910756 -0.00135257289 0.00011966865" rpy="0 0 0"/>
+      <mass value="0.01385835"/>
+      <inertia ixx="0.00000094744556" ixy="0.00000010086310" ixz="0.00000003605469" iyy="0.00000254698095" iyz="0.00000000325710" izz="0.00000204288618"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_four_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_four_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_five_joint" type="fixed">
+    <origin xyz="0.085 -0.004 -0.02395" rpy="0 0 0"/>
+    <parent link="right_palm_link"/>
+    <child link="right_five_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-0.3" upper="1.84"/>
+  </joint>
+  <link name="right_five_link">
+    <inertial>
+      <origin xyz="0.03625612321 0.00157250459 0.00026989918" rpy="0 0 0"/>
+      <mass value="0.04884689"/>
+      <inertia ixx="0.00000383218469" ixy="-0.00000003239984" ixz="0.00000001923377" iyy="0.00000890974599" iyz="0.00000004429800" izz="0.00000865612355"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_five_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_five_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_six_joint" type="fixed">
+    <origin xyz="0.0471 0.0036 0" rpy="0 0 0"/>
+    <parent link="right_five_link"/>
+    <child link="right_six_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="0" upper="1.84"/>
+  </joint>
+  <link name="right_six_link">
+    <inertial>
+      <origin xyz="0.02375910756 -0.00135257289 0.00011966865" rpy="0 0 0"/>
+      <mass value="0.01385835"/>
+      <inertia ixx="0.00000094744556" ixy="0.00000010086310" ixz="0.00000003605469" iyy="0.00000254698095" iyz="0.00000000325710" izz="0.00000204288618"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_six_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="../mesh/G1/right_six_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+</robot>
\ No newline at end of file
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/g1_23dof.urdf b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/g1_23dof.urdf
new file mode 100644
index 0000000000000000000000000000000000000000..35562b0148ac88a6fcbd8a9f35e514d051b61f73
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/g1_23dof.urdf
@@ -0,0 +1,893 @@
+<robot name="g1_23dof">
+
+  <!-- [CAUTION] uncomment when convert to mujoco -->
+  <!-- <link name="world"></link>
+  <joint name="floating_base_joint" type="floating">
+    <parent link="world"/>
+    <child link="pelvis"/>
+  </joint> -->
+
+  <link name="pelvis">
+    <inertial>
+      <origin xyz="0 0 -0.07605" rpy="0 0 0"/>
+      <mass value="3.813"/>
+      <inertia ixx="0.010549" ixy="0" ixz="2.1E-06" iyy="0.0093089" iyz="0" izz="0.0079184"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/pelvis.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+  </link>
+  <link name="pelvis_contour_link">
+    <inertial>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <mass value="0.001"/>
+      <inertia ixx="1e-7" ixy="0" ixz="0" iyy="1e-7" iyz="0" izz="1e-7"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/pelvis_contour_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/pelvis_contour_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="pelvis_contour_joint" type="fixed">
+    <parent link="pelvis"/>
+    <child link="pelvis_contour_link"/>
+  </joint>
+
+  <!-- Legs -->
+  <link name="left_hip_pitch_link">
+    <inertial>
+      <origin xyz="0.002741 0.047791 -0.02606" rpy="0 0 0"/>
+      <mass value="1.35"/>
+      <inertia ixx="0.001811" ixy="3.68E-05" ixz="-3.44E-05" iyy="0.0014193" iyz="0.000171" izz="0.0012812"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_hip_pitch_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_hip_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_hip_pitch_joint" type="revolute">
+    <origin xyz="0 0.064452 -0.1027" rpy="0 0 0"/>
+    <parent link="pelvis"/>
+    <child link="left_hip_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-2.5307" upper="2.8798" effort="88" velocity="32"/>
+  </joint>
+  <link name="left_hip_roll_link">
+    <inertial>
+      <origin xyz="0.029812 -0.001045 -0.087934" rpy="0 0 0"/>
+      <mass value="1.52"/>
+      <inertia ixx="0.0023773" ixy="-3.8E-06" ixz="-0.0003908" iyy="0.0024123" iyz="1.84E-05" izz="0.0016595"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_hip_roll_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_hip_roll_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_hip_roll_joint" type="revolute">
+    <origin xyz="0 0.052 -0.030465" rpy="0 -0.1749 0"/>
+    <parent link="left_hip_pitch_link"/>
+    <child link="left_hip_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-0.5236" upper="2.9671" effort="88" velocity="32"/>
+  </joint>
+  <link name="left_hip_yaw_link">
+    <inertial>
+      <origin xyz="-0.057709 -0.010981 -0.15078" rpy="0 0 0"/>
+      <mass value="1.702"/>
+      <inertia ixx="0.0057774" ixy="-0.0005411" ixz="-0.0023948" iyy="0.0076124" iyz="-0.0007072" izz="0.003149"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_hip_yaw_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_hip_yaw_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_hip_yaw_joint" type="revolute">
+    <origin xyz="0.025001 0 -0.12412" rpy="0 0 0"/>
+    <parent link="left_hip_roll_link"/>
+    <child link="left_hip_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.7576" upper="2.7576" effort="88" velocity="32"/>
+  </joint>
+  <link name="left_knee_link">
+    <inertial>
+      <origin xyz="0.005457 0.003964 -0.12074" rpy="0 0 0"/>
+      <mass value="1.932"/>
+      <inertia ixx="0.011329" ixy="4.82E-05" ixz="-4.49E-05" iyy="0.011277" iyz="-0.0007146" izz="0.0015168"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_knee_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_knee_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_knee_joint" type="revolute">
+    <origin xyz="-0.078273 0.0021489 -0.17734" rpy="0 0.1749 0"/>
+    <parent link="left_hip_yaw_link"/>
+    <child link="left_knee_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.087267" upper="2.8798" effort="139" velocity="20"/>
+  </joint>
+  <link name="left_ankle_pitch_link">
+    <inertial>
+      <origin xyz="-0.007269 0 0.011137" rpy="0 0 0"/>
+      <mass value="0.074"/>
+      <inertia ixx="8.4E-06" ixy="0" ixz="-2.9E-06" iyy="1.89E-05" iyz="0" izz="1.26E-05"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_ankle_pitch_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_ankle_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_ankle_pitch_joint" type="revolute">
+    <origin xyz="0 -9.4445E-05 -0.30001" rpy="0 0 0"/>
+    <parent link="left_knee_link"/>
+    <child link="left_ankle_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.87267" upper="0.5236" effort="50" velocity="37"/>
+  </joint>
+  <link name="left_ankle_roll_link">
+    <inertial>
+      <origin xyz="0.026505 0 -0.016425" rpy="0 0 0"/>
+      <mass value="0.608"/>
+      <inertia ixx="0.0002231" ixy="2E-07" ixz="8.91E-05" iyy="0.0016161" iyz="-1E-07" izz="0.0016667"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_ankle_roll_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="-0.05 0.025 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.005"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="-0.05 -0.025 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.005"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="0.12 0.03 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.005"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="0.12 -0.03 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.005"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_ankle_roll_joint" type="revolute">
+    <origin xyz="0 0 -0.017558" rpy="0 0 0"/>
+    <parent link="left_ankle_pitch_link"/>
+    <child link="left_ankle_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-0.2618" upper="0.2618" effort="50" velocity="37"/>
+  </joint>
+  <link name="right_hip_pitch_link">
+    <inertial>
+      <origin xyz="0.002741 -0.047791 -0.02606" rpy="0 0 0"/>
+      <mass value="1.35"/>
+      <inertia ixx="0.001811" ixy="-3.68E-05" ixz="-3.44E-05" iyy="0.0014193" iyz="-0.000171" izz="0.0012812"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_hip_pitch_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_hip_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_hip_pitch_joint" type="revolute">
+    <origin xyz="0 -0.064452 -0.1027" rpy="0 0 0"/>
+    <parent link="pelvis"/>
+    <child link="right_hip_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-2.5307" upper="2.8798" effort="88" velocity="32"/>
+  </joint>
+  <link name="right_hip_roll_link">
+    <inertial>
+      <origin xyz="0.029812 0.001045 -0.087934" rpy="0 0 0"/>
+      <mass value="1.52"/>
+      <inertia ixx="0.0023773" ixy="3.8E-06" ixz="-0.0003908" iyy="0.0024123" iyz="-1.84E-05" izz="0.0016595"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_hip_roll_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_hip_roll_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_hip_roll_joint" type="revolute">
+    <origin xyz="0 -0.052 -0.030465" rpy="0 -0.1749 0"/>
+    <parent link="right_hip_pitch_link"/>
+    <child link="right_hip_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-2.9671" upper="0.5236" effort="88" velocity="32"/>
+  </joint>
+  <link name="right_hip_yaw_link">
+    <inertial>
+      <origin xyz="-0.057709 0.010981 -0.15078" rpy="0 0 0"/>
+      <mass value="1.702"/>
+      <inertia ixx="0.0057774" ixy="0.0005411" ixz="-0.0023948" iyy="0.0076124" iyz="0.0007072" izz="0.003149"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_hip_yaw_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_hip_yaw_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_hip_yaw_joint" type="revolute">
+    <origin xyz="0.025001 0 -0.12412" rpy="0 0 0"/>
+    <parent link="right_hip_roll_link"/>
+    <child link="right_hip_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.7576" upper="2.7576" effort="88" velocity="32"/>
+  </joint>
+  <link name="right_knee_link">
+    <inertial>
+      <origin xyz="0.005457 -0.003964 -0.12074" rpy="0 0 0"/>
+      <mass value="1.932"/>
+      <inertia ixx="0.011329" ixy="-4.82E-05" ixz="4.49E-05" iyy="0.011277" iyz="0.0007146" izz="0.0015168"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_knee_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_knee_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_knee_joint" type="revolute">
+    <origin xyz="-0.078273 -0.0021489 -0.17734" rpy="0 0.1749 0"/>
+    <parent link="right_hip_yaw_link"/>
+    <child link="right_knee_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.087267" upper="2.8798" effort="139" velocity="20"/>
+  </joint>
+  <link name="right_ankle_pitch_link">
+    <inertial>
+      <origin xyz="-0.007269 0 0.011137" rpy="0 0 0"/>
+      <mass value="0.074"/>
+      <inertia ixx="8.4E-06" ixy="0" ixz="-2.9E-06" iyy="1.89E-05" iyz="0" izz="1.26E-05"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_ankle_pitch_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_ankle_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_ankle_pitch_joint" type="revolute">
+    <origin xyz="0 9.4445E-05 -0.30001" rpy="0 0 0"/>
+    <parent link="right_knee_link"/>
+    <child link="right_ankle_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.87267" upper="0.5236" effort="50" velocity="37"/>
+  </joint>
+  <link name="right_ankle_roll_link">
+    <inertial>
+      <origin xyz="0.026505 0 -0.016425" rpy="0 0 0"/>
+      <mass value="0.608"/>
+      <inertia ixx="0.0002231" ixy="-2E-07" ixz="8.91E-05" iyy="0.0016161" iyz="1E-07" izz="0.0016667"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_ankle_roll_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="-0.05 0.025 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.005"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="-0.05 -0.025 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.005"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="0.12 0.03 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.005"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="0.12 -0.03 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.005"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_ankle_roll_joint" type="revolute">
+    <origin xyz="0 0 -0.017558" rpy="0 0 0"/>
+    <parent link="right_ankle_pitch_link"/>
+    <child link="right_ankle_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-0.2618" upper="0.2618" effort="50" velocity="37"/>
+  </joint>
+
+  <!-- Torso -->
+  <link name="waist_yaw_fixed_link">
+    <inertial>
+      <origin xyz="0.003964 0 0.018769" rpy="0 0 0"/>
+      <mass value="0.244"/>
+      <inertia ixx="9.9587E-05" ixy="-1.833E-06" ixz="-1.2617E-05" iyy="0.00012411" iyz="-1.18E-07" izz="0.00015586"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/waist_yaw_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+  </link>
+  <joint name="waist_yaw_fixed_joint" type="fixed">
+    <origin xyz="0.0039635 0 -0.054" rpy="0 0 0"/>
+    <parent link="torso_link"/>
+    <child link="waist_yaw_fixed_link"/>
+  </joint>
+  <joint name="waist_yaw_joint" type="fixed">
+    <origin xyz="-0.0039635 0 0.054" rpy="0 0 0"/>
+    <parent link="pelvis"/>
+    <child link="torso_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.618" upper="2.618" effort="88" velocity="32"/>
+  </joint>
+  <link name="torso_link">
+    <inertial>
+      <origin xyz="0.002601 0.000257 0.153719" rpy="0 0 0"/>
+      <mass value="8.562"/>
+      <inertia ixx="0.065674966" ixy="-8.597E-05" ixz="-0.001737252" iyy="0.053535188" iyz="8.6899E-05" izz="0.030808125"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/torso_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/torso_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+
+  <!-- LOGO -->
+  <joint name="logo_joint" type="fixed">
+    <origin xyz="0.0039635 0 -0.054" rpy="0 0 0"/>
+    <parent link="torso_link"/>
+    <child link="logo_link"/>
+  </joint>
+  <link name="logo_link">
+    <inertial>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <mass value="0.001"/>
+      <inertia ixx="1e-7" ixy="0" ixz="0" iyy="1e-7" iyz="0" izz="1e-7"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/logo_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/logo_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+
+  <!-- Head -->
+  <link name="head_link">
+    <inertial>
+      <origin xyz="0.005267 0.000299 0.449869" rpy="0 0 0"/>
+      <mass value="1.036"/>
+      <inertia ixx="0.004085051" ixy="-2.543E-06" ixz="-6.9455E-05" iyy="0.004185212" iyz="-3.726E-06" izz="0.001807911"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/head_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/head_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="head_joint" type="fixed">
+    <origin xyz="0.0039635 0 -0.054" rpy="0 0 0"/>
+    <parent link="torso_link"/>
+    <child link="head_link"/>
+  </joint>
+
+  <!-- Waist Support -->
+  <link name="waist_support_link">
+    <inertial>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <mass value="0.001"/>
+      <inertia ixx="1e-7" ixy="0" ixz="0" iyy="1e-7" iyz="0" izz="1e-7"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/waist_support_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/waist_support_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="waist_support_joint" type="fixed">
+    <origin xyz="0.0039635 0 -0.054" rpy="0 0 0"/>
+    <parent link="torso_link"/>
+    <child link="waist_support_link"/>
+  </joint>
+
+  <!-- IMU -->
+  <link name="imu_link"></link>
+  <joint name="imu_joint" type="fixed">
+    <origin xyz="-0.03959 -0.00224 0.13792" rpy="0 0 0"/>
+    <parent link="torso_link"/>
+    <child link="imu_link"/>
+  </joint>
+
+  <!-- d435 -->
+  <link name="d435_link"></link>
+  <joint name="d435_joint" type="fixed">
+    <origin xyz="0.0576235 0.01753 0.41987" rpy="0 0.8307767239493009 0"/>
+    <parent link="torso_link"/>
+    <child link="d435_link"/>
+  </joint>
+
+  <!-- mid360 -->
+  <link name="mid360_link"></link>
+  <joint name="mid360_joint" type="fixed">
+    <origin xyz="0.0002835 0.00003 0.40618" rpy="0 0.04014257279586953 0"/>
+    <parent link="torso_link"/>
+    <child link="mid360_link"/>
+  </joint>
+
+  <!-- Arm -->
+  <link name="left_shoulder_pitch_link">
+    <inertial>
+      <origin xyz="0 0.035892 -0.011628" rpy="0 0 0"/>
+      <mass value="0.718"/>
+      <inertia ixx="0.0004291" ixy="-9.2E-06" ixz="6.4E-06" iyy="0.000453" iyz="2.26E-05" izz="0.000423"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_shoulder_pitch_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0.04 -0.01" rpy="0 1.5707963267948966 0"/>
+      <geometry>
+        <cylinder radius="0.03" length="0.05"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_shoulder_pitch_joint" type="fixed">
+    <origin xyz="0.0039563 0.10022 0.23778" rpy="0.27931 5.4949E-05 -0.00019159"/>
+    <parent link="torso_link"/>
+    <child link="left_shoulder_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-3.0892" upper="2.6704" effort="25" velocity="37"/>
+  </joint>
+  <link name="left_shoulder_roll_link">
+    <inertial>
+      <origin xyz="-0.000227 0.00727 -0.063243" rpy="0 0 0"/>
+      <mass value="0.643"/>
+      <inertia ixx="0.0006177" ixy="-1E-06" ixz="8.7E-06" iyy="0.0006912" iyz="-5.3E-06" izz="0.0003894"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_shoulder_roll_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="-0.004 0.006 -0.053" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.03" length="0.03"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_shoulder_roll_joint" type="fixed">
+    <origin xyz="0 0.038 -0.013831" rpy="-0.27925 0 0"/>
+    <parent link="left_shoulder_pitch_link"/>
+    <child link="left_shoulder_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-1.5882" upper="2.2515" effort="25" velocity="37"/>
+  </joint>
+  <link name="left_shoulder_yaw_link">
+    <inertial>
+      <origin xyz="0.010773 -0.002949 -0.072009" rpy="0 0 0"/>
+      <mass value="0.734"/>
+      <inertia ixx="0.0009988" ixy="7.9E-06" ixz="0.0001412" iyy="0.0010605" iyz="-2.86E-05" izz="0.0004354"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_shoulder_yaw_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_shoulder_yaw_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_shoulder_yaw_joint" type="fixed">
+    <origin xyz="0 0.00624 -0.1032" rpy="0 0 0"/>
+    <parent link="left_shoulder_roll_link"/>
+    <child link="left_shoulder_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.618" upper="2.618" effort="25" velocity="37"/>
+  </joint>
+  <link name="left_elbow_link">
+    <inertial>
+      <origin xyz="0.064956 0.004454 -0.010062" rpy="0 0 0"/>
+      <mass value="0.6"/>
+      <inertia ixx="0.0002891" ixy="6.53E-05" ixz="1.72E-05" iyy="0.0004152" iyz="-5.6E-06" izz="0.0004197"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_elbow_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_elbow_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_elbow_joint" type="fixed">
+    <origin xyz="0.015783 0 -0.080518" rpy="0 0 0"/>
+    <parent link="left_shoulder_yaw_link"/>
+    <child link="left_elbow_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-1.0472" upper="2.0944" effort="25" velocity="37"/>
+  </joint>
+  <joint name="left_wrist_roll_joint" type="fixed">
+    <origin xyz="0.100 0.00188791 -0.010" rpy="0 0 0"/>
+    <axis xyz="1 0 0"/>
+    <parent link="left_elbow_link"/>
+    <child link="left_wrist_roll_rubber_hand"/>
+    <limit effort="25" velocity="37" lower="-1.972222054" upper="1.972222054"/>
+  </joint>
+  <link name="left_wrist_roll_rubber_hand">
+    <inertial>
+      <origin xyz="0.10794656650 0.00163511945 0.00202244863" rpy="0 0 0"/>
+      <mass value="0.35692864"/>
+      <inertia ixx="0.00019613494735" ixy="-0.00000419816908" ixz="-0.00003950860580" iyy="0.00200280358206" iyz="0.00000249774203" izz="0.00194181412808"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_wrist_roll_rubber_hand.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/left_wrist_roll_rubber_hand.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="right_shoulder_pitch_link">
+    <inertial>
+      <origin xyz="0 -0.035892 -0.011628" rpy="0 0 0"/>
+      <mass value="0.718"/>
+      <inertia ixx="0.0004291" ixy="9.2E-06" ixz="6.4E-06" iyy="0.000453" iyz="-2.26E-05" izz="0.000423"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_shoulder_pitch_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 -0.04 -0.01" rpy="0 1.5707963267948966 0"/>
+      <geometry>
+        <cylinder radius="0.03" length="0.05"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_shoulder_pitch_joint" type="fixed">
+    <origin xyz="0.0039563 -0.10021 0.23778" rpy="-0.27931 5.4949E-05 0.00019159"/>
+    <parent link="torso_link"/>
+    <child link="right_shoulder_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-3.0892" upper="2.6704" effort="25" velocity="37"/>
+  </joint>
+  <link name="right_shoulder_roll_link">
+    <inertial>
+      <origin xyz="-0.000227 -0.00727 -0.063243" rpy="0 0 0"/>
+      <mass value="0.643"/>
+      <inertia ixx="0.0006177" ixy="1E-06" ixz="8.7E-06" iyy="0.0006912" iyz="5.3E-06" izz="0.0003894"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_shoulder_roll_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="-0.004 -0.006 -0.053" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.03" length="0.03"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_shoulder_roll_joint" type="fixed">
+    <origin xyz="0 -0.038 -0.013831" rpy="0.27925 0 0"/>
+    <parent link="right_shoulder_pitch_link"/>
+    <child link="right_shoulder_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-2.2515" upper="1.5882" effort="25" velocity="37"/>
+  </joint>
+  <link name="right_shoulder_yaw_link">
+    <inertial>
+      <origin xyz="0.010773 0.002949 -0.072009" rpy="0 0 0"/>
+      <mass value="0.734"/>
+      <inertia ixx="0.0009988" ixy="-7.9E-06" ixz="0.0001412" iyy="0.0010605" iyz="2.86E-05" izz="0.0004354"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_shoulder_yaw_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_shoulder_yaw_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_shoulder_yaw_joint" type="fixed">
+    <origin xyz="0 -0.00624 -0.1032" rpy="0 0 0"/>
+    <parent link="right_shoulder_roll_link"/>
+    <child link="right_shoulder_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.618" upper="2.618" effort="25" velocity="37"/>
+  </joint>
+  <link name="right_elbow_link">
+    <inertial>
+      <origin xyz="0.064956 -0.004454 -0.010062" rpy="0 0 0"/>
+      <mass value="0.6"/>
+      <inertia ixx="0.0002891" ixy="-6.53E-05" ixz="1.72E-05" iyy="0.0004152" iyz="5.6E-06" izz="0.0004197"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_elbow_link.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_elbow_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_elbow_joint" type="fixed">
+    <origin xyz="0.015783 0 -0.080518" rpy="0 0 0"/>
+    <parent link="right_shoulder_yaw_link"/>
+    <child link="right_elbow_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-1.0472" upper="2.0944" effort="25" velocity="37"/>
+  </joint>
+  <joint name="right_wrist_roll_joint" type="fixed">
+    <origin xyz="0.100 -0.00188791 -0.010" rpy="0 0 0"/>
+    <axis xyz="1 0 0"/>
+    <parent link="right_elbow_link"/>
+    <child link="right_wrist_roll_rubber_hand"/>
+    <limit effort="25" velocity="37" lower="-1.972222054" upper="1.972222054"/>
+  </joint>
+  <link name="right_wrist_roll_rubber_hand">
+    <inertial>
+      <origin xyz="0.10794656650 -0.00163511945 0.00202244863" rpy="0 0 0"/>
+      <mass value="0.35692864"/>
+      <inertia ixx="0.00019613494735" ixy="0.00000419816908" ixz="-0.00003950860580" iyy="0.00200280358206" iyz="-0.00000249774203" izz="0.00194181412808"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_wrist_roll_rubber_hand.STL"/>
+      </geometry>
+      <material name="white">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1_23DoF/right_wrist_roll_rubber_hand.STL"/>
+      </geometry>
+    </collision>
+  </link>
+</robot>
\ No newline at end of file
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/g1_unitree.urdf b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/g1_unitree.urdf
new file mode 100644
index 0000000000000000000000000000000000000000..8648c790da963ec5cb7861749147c76e7de45555
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/g1_unitree.urdf
@@ -0,0 +1,1295 @@
+<robot name="g1">
+
+  <link name="pelvis">
+    <inertial>
+      <origin xyz="0 0 -0.07605" rpy="0 0 0"/>
+      <mass value="2.86"/>
+      <inertia ixx="0.0079143" ixy="0" ixz="1.6E-06" iyy="0.0069837" iyz="0" izz="0.0059404"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/pelvis.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/pelvis.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="pelvis_contour_link">
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/pelvis_contour_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+  </link>
+  <joint name="pelvis_contour_joint" type="fixed">
+    <parent link="pelvis"/>
+    <child link="pelvis_contour_link"/>
+  </joint>
+
+  <!-- Leg -->
+  <link name="left_hip_pitch_link">
+    <inertial>
+      <origin xyz="0.001962 0.049392 -0.000941" rpy="0 0 0"/>
+      <mass value="1.299"/>
+      <inertia ixx="0.0013873" ixy="-1.63E-05" ixz="-1E-06" iyy="0.0009059" iyz="-4.24E-05" izz="0.0009196"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_hip_pitch_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_hip_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_hip_pitch_joint" type="revolute">
+    <origin xyz="0 0.06445 -0.1027" rpy="0 -0.34907 0"/>
+    <parent link="pelvis"/>
+    <child link="left_hip_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-2.35" upper="3.05" effort="88" velocity="32"/>
+  </joint>
+  <link name="left_hip_roll_link">
+    <inertial>
+      <origin xyz="0.024757 -0.001036 -0.086323" rpy="0 0 0"/>
+      <mass value="1.446"/>
+      <inertia ixx="0.0022702" ixy="-3.7E-06" ixz="-0.0003789" iyy="0.002304" iyz="-1.8E-05" izz="0.0016001"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_hip_roll_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_hip_roll_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_hip_roll_joint" type="revolute">
+    <origin xyz="0 0.0523 0" rpy="0 0 0"/>
+    <parent link="left_hip_pitch_link"/>
+    <child link="left_hip_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-0.26" upper="2.53" effort="88" velocity="32"/>
+  </joint>
+  <link name="left_hip_yaw_link">
+    <inertial>
+      <origin xyz="-0.053554 -0.011477 -0.14067" rpy="0 0 0"/>
+      <mass value="2.052"/>
+      <inertia ixx="0.0087264" ixy="-0.0004402" ixz="-0.0036676" iyy="0.011374" iyz="-0.0006654" izz="0.004279"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_hip_yaw_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_hip_yaw_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_hip_yaw_joint" type="revolute">
+    <origin xyz="0.01966 -0.0012139 -0.1241" rpy="0 0 0"/>
+    <parent link="left_hip_roll_link"/>
+    <child link="left_hip_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.75" upper="2.75" effort="88" velocity="32"/>
+  </joint>
+  <link name="left_knee_link">
+    <inertial>
+      <origin xyz="0.005505 0.006534 -0.116629" rpy="0 0 0"/>
+      <mass value="2.252"/>
+      <inertia ixx="0.012443837" ixy="0.000053496" ixz="-0.000437641" iyy="0.012674902" iyz="-0.000682499" izz="0.001986501"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_knee_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0.007 0.005 -0.15" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.015" length="0.15"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_knee_joint" type="revolute">
+    <origin xyz="-0.078292 -0.0017335 -0.177225" rpy="0 0.5096 0"/>
+    <parent link="left_hip_yaw_link"/>
+    <child link="left_knee_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.33489" upper="2.5449" effort="139" velocity="20"/>
+  </joint>
+  <link name="left_ankle_pitch_link">
+    <inertial>
+      <origin xyz="-0.007269 0 0.011137" rpy="0 0 0"/>
+      <mass value="0.074"/>
+      <inertia ixx="8.4E-06" ixy="0" ixz="-2.9E-06" iyy="1.89E-05" iyz="0" izz="1.26E-05"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_ankle_pitch_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_ankle_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_ankle_pitch_joint" type="revolute">
+    <origin xyz="0 0.0040687 -0.30007" rpy="0 -0.16053 0"/>
+    <parent link="left_knee_link"/>
+    <child link="left_ankle_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.68" upper="0.73" effort="40" velocity="53"/>
+  </joint>
+  <link name="left_ankle_roll_link">
+    <inertial>
+      <origin xyz="0.024762 2E-05 -0.012526" rpy="0 0 0"/>
+      <mass value="0.391"/>
+      <inertia ixx="0.0001552" ixy="-1E-07" ixz="7.51E-05" iyy="0.0010657" iyz="1E-07" izz="0.001098"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_ankle_roll_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_ankle_roll_link.STL"/>
+      </geometry>
+    </collision>
+    <!-- <collision>
+      <origin xyz="-0.06 0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="-0.06 -0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="0.13 0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="0.13 -0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision> -->
+  </link>
+  <joint name="left_ankle_roll_joint" type="revolute">
+    <origin xyz="0 0 -0.017558" rpy="0 0 0"/>
+    <parent link="left_ankle_pitch_link"/>
+    <child link="left_ankle_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-0.2618" upper="0.2618" effort="40" velocity="53"/>
+  </joint>
+  <link name="right_hip_pitch_link">
+    <inertial>
+      <origin xyz="0.001962 -0.049392 -0.000941" rpy="0 0 0"/>
+      <mass value="1.299"/>
+      <inertia ixx="0.0013873" ixy="1.63E-05" ixz="-1E-06" iyy="0.0009059" iyz="4.24E-05" izz="0.0009196"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_hip_pitch_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_hip_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_hip_pitch_joint" type="revolute">
+    <origin xyz="0 -0.06445 -0.1027" rpy="0 -0.34907 0"/>
+    <parent link="pelvis"/>
+    <child link="right_hip_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-2.35" upper="3.05" effort="88" velocity="32"/>
+  </joint>
+  <link name="right_hip_roll_link">
+    <inertial>
+      <origin xyz="0.024757 0.001036 -0.086323" rpy="0 0 0"/>
+      <mass value="1.446"/>
+      <inertia ixx="0.0022702" ixy="3.7E-06" ixz="-0.0003789" iyy="0.002304" iyz="1.8E-05" izz="0.0016001"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_hip_roll_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_hip_roll_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_hip_roll_joint" type="revolute">
+    <origin xyz="0 -0.0523 0" rpy="0 0 0"/>
+    <parent link="right_hip_pitch_link"/>
+    <child link="right_hip_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-2.53" upper="0.26" effort="88" velocity="32"/>
+  </joint>
+  <link name="right_hip_yaw_link">
+    <inertial>
+      <origin xyz="-0.053554 0.011477 -0.14067" rpy="0 0 0"/>
+      <mass value="2.052"/>
+      <inertia ixx="0.0087264" ixy="0.0004402" ixz="-0.0036676" iyy="0.011374" iyz="0.0006654" izz="0.004279"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_hip_yaw_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_hip_yaw_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_hip_yaw_joint" type="revolute">
+    <origin xyz="0.01966 0.0012139 -0.1241" rpy="0 0 0"/>
+    <parent link="right_hip_roll_link"/>
+    <child link="right_hip_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.75" upper="2.75" effort="88" velocity="32"/>
+  </joint>
+  <link name="right_knee_link">
+    <inertial>
+      <origin xyz="0.005505 -0.006534 -0.116629" rpy="0 0 0"/>
+      <mass value="2.252"/>
+      <inertia ixx="0.012443837" ixy="-0.000053496" ixz="-0.000437641" iyy="0.012674902" iyz="0.000682499" izz="0.001986501"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_knee_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0.007 -0.005 -0.15" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.015" length="0.15"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_knee_joint" type="revolute">
+    <origin xyz="-0.078292 0.0017335 -0.177225" rpy="0 0.5096 0"/>
+    <parent link="right_hip_yaw_link"/>
+    <child link="right_knee_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.33489" upper="2.5449" effort="139" velocity="20"/>
+  </joint>
+  <link name="right_ankle_pitch_link">
+    <inertial>
+      <origin xyz="-0.007269 0 0.011137" rpy="0 0 0"/>
+      <mass value="0.074"/>
+      <inertia ixx="8.4E-06" ixy="0" ixz="-2.9E-06" iyy="1.89E-05" iyz="0" izz="1.26E-05"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_ankle_pitch_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_ankle_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_ankle_pitch_joint" type="revolute">
+    <origin xyz="0 -0.0040687 -0.30007" rpy="0 -0.16053 0"/>
+    <parent link="right_knee_link"/>
+    <child link="right_ankle_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.68" upper="0.73" effort="40" velocity="53"/>
+  </joint>
+  <link name="right_ankle_roll_link">
+    <inertial>
+      <origin xyz="0.024762 -2E-05 -0.012526" rpy="0 0 0"/>
+      <mass value="0.391"/>
+      <inertia ixx="0.0001552" ixy="1E-07" ixz="7.51E-05" iyy="0.0010657" iyz="-1E-07" izz="0.001098"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_ankle_roll_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_ankle_roll_link.STL"/>
+      </geometry>
+    </collision>
+    <!-- <collision>
+      <origin xyz="-0.06 0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="-0.06 -0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="0.13 0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision>
+    <collision>
+      <origin xyz="0.13 -0.02 -0.03" rpy="0 0 0"/>
+      <geometry>
+        <sphere radius="0.001"/>
+      </geometry>
+    </collision> -->
+  </link>
+  <joint name="right_ankle_roll_joint" type="revolute">
+    <origin xyz="0 0 -0.017558" rpy="0 0 0"/>
+    <parent link="right_ankle_pitch_link"/>
+    <child link="right_ankle_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-0.2618" upper="0.2618" effort="40" velocity="53"/>
+  </joint>
+
+  <!-- Torso -->
+  <link name="torso_link">
+    <inertial>
+      <origin xyz="0.00197122283 0.00266902037 0.16936510960" rpy="0 0 0"/>
+      <mass value="6.33959811"/>
+      <inertia ixx="0.06037082109027" ixy="-0.00010168407055" ixz="0.00069845018092" iyy="0.04355153796551" iyz="0.00034366422728" izz="0.03339981866010"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/torso_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/torso_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="torso_joint" type="fixed">
+    <origin xyz="0 0 0" rpy="0 0 0"/>
+    <parent link="pelvis"/>
+    <child link="torso_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.618" upper="2.618" effort="88" velocity="32"/>
+  </joint>
+
+  <!-- Head -->
+  <link name="head_link">
+    <inertial>
+      <origin xyz="0.00138066852 0.00028430950 0.42034187824" rpy="0 0 0"/>
+      <mass value="1.17976522"/>
+      <inertia ixx="0.00543236042361" ixy="0.00000140137425" ixz="0.00034554752228" iyy="0.00552885306699" iyz="0.00001501216392" izz="0.00165378108136"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/head_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/head_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="head_joint" type="fixed">
+    <origin xyz="0 0 0" rpy="0 0 0"/>
+    <parent link="torso_link"/>
+    <child link="head_link"/>
+  </joint>
+
+  <!-- L ARM -->
+  <link name="left_shoulder_pitch_link">
+    <inertial>
+      <origin xyz="-0.001431 0.048811 0.001304" rpy="0 0 0"/>
+      <mass value="0.713"/>
+      <inertia ixx="0.0004614" ixy="-9.3E-06" ixz="1E-05" iyy="0.0004146" iyz="5.5E-06" izz="0.0004416"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_shoulder_pitch_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_shoulder_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_shoulder_pitch_joint" type="fixed">
+    <origin xyz="-0.0025 0.10396 0.25928" rpy="0.27925 0 0"/>
+    <parent link="torso_link"/>
+    <child link="left_shoulder_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-2.9671" upper="2.7925" effort="21" velocity="53"/>
+  </joint>
+  <link name="left_shoulder_roll_link">
+    <inertial>
+      <origin xyz="-0.003415 0.006955 -0.064598" rpy="0 0 0"/>
+      <mass value="0.642"/>
+      <inertia ixx="0.0006159" ixy="0" ixz="-5.6E-06" iyy="0.0006835" iyz="2.1E-06" izz="0.000373"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_shoulder_roll_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="-0.004 0.006 -0.053" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.03" length="0.03"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_shoulder_roll_joint" type="fixed">
+    <origin xyz="0 0.052 0" rpy="-0.27925 0 0"/>
+    <parent link="left_shoulder_pitch_link"/>
+    <child link="left_shoulder_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-1.5882" upper="2.2515" effort="21" velocity="53"/>
+  </joint>
+  <link name="left_shoulder_yaw_link">
+    <inertial>
+      <origin xyz="0.000375 -0.00444 -0.072374" rpy="0 0 0"/>
+      <mass value="0.713"/>
+      <inertia ixx="0.0009699" ixy="7.1E-06" ixz="7.8E-06" iyy="0.0009691" iyz="-4.49E-05" izz="0.0003826"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_shoulder_yaw_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_shoulder_yaw_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_shoulder_yaw_joint" type="fixed">
+    <origin xyz="-0.00354 0.0062424 -0.1032" rpy="0 0 0"/>
+    <parent link="left_shoulder_roll_link"/>
+    <child link="left_shoulder_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.618" upper="2.618" effort="21" velocity="53"/>
+  </joint>
+  <link name="left_elbow_pitch_link">
+    <inertial>
+      <origin xyz="0.064497 0.002873 0" rpy="0 0 0"/>
+      <mass value="0.601"/>
+      <inertia ixx="0.0002845" ixy="8.06E-05" ixz="9E-09" iyy="0.0004647" iyz="5E-09" izz="0.0004712"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_elbow_pitch_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_elbow_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_elbow_pitch_joint" type="fixed">
+    <origin xyz="0 0.00189 -0.0855" rpy="0 0 0"/>
+    <parent link="left_shoulder_yaw_link"/>
+    <child link="left_elbow_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.2268" upper="3.4208" effort="21" velocity="53"/>
+  </joint>
+  <link name="left_elbow_roll_link">
+    <inertial>
+      <origin xyz="0.081811 0.001454 0.001" rpy="0 0 0"/>
+      <mass value="0.227"/>
+      <inertia ixx="0.0001093" ixy="-1.55E-05" ixz="-2.91E-05" iyy="0.000977" iyz="1E-06" izz="0.0009676"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_elbow_roll_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_elbow_roll_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_elbow_roll_joint" type="fixed">
+    <origin xyz="0.1 0 0" rpy="0 0 0"/>
+    <parent link="left_elbow_pitch_link"/>
+    <child link="left_elbow_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-2.0943" upper="2.0943" effort="21" velocity="53"/>
+  </joint>
+
+  <!-- R ARM -->
+  <link name="right_shoulder_pitch_link">
+    <inertial>
+      <origin xyz="-0.001431 -0.048811 0.001304" rpy="0 0 0"/>
+      <mass value="0.713"/>
+      <inertia ixx="0.0004614" ixy="9.3E-06" ixz="1E-05" iyy="0.0004146" iyz="-5.5E-06" izz="0.0004416"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_shoulder_pitch_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_shoulder_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_shoulder_pitch_joint" type="fixed">
+    <origin xyz="-0.0025 -0.10396 0.25928" rpy="-0.27925 0 0"/>
+    <parent link="torso_link"/>
+    <child link="right_shoulder_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-2.9671" upper="2.7925" effort="21" velocity="53"/>
+  </joint>
+  <link name="right_shoulder_roll_link">
+    <inertial>
+      <origin xyz="-0.003415 -0.006955 -0.064598" rpy="0 0 0"/>
+      <mass value="0.642"/>
+      <inertia ixx="0.0006159" ixy="0" ixz="-5.6E-06" iyy="0.0006835" iyz="-2.1E-06" izz="0.000373"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_shoulder_roll_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="-0.004 -0.006 -0.053" rpy="0 0 0"/>
+      <geometry>
+        <cylinder radius="0.03" length="0.03"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_shoulder_roll_joint" type="fixed">
+    <origin xyz="0 -0.052 0" rpy="0.27925 0 0"/>
+    <parent link="right_shoulder_pitch_link"/>
+    <child link="right_shoulder_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-2.2515" upper="1.5882" effort="21" velocity="53"/>
+  </joint>
+  <link name="right_shoulder_yaw_link">
+    <inertial>
+      <origin xyz="0.000375 0.00444 -0.072374" rpy="0 0 0"/>
+      <mass value="0.713"/>
+      <inertia ixx="0.0009699" ixy="-7.1E-06" ixz="7.8E-06" iyy="0.0009691" iyz="4.49E-05" izz="0.0003826"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_shoulder_yaw_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_shoulder_yaw_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_shoulder_yaw_joint" type="fixed">
+    <origin xyz="-0.00354 -0.0062424 -0.1032" rpy="0 0 0"/>
+    <parent link="right_shoulder_roll_link"/>
+    <child link="right_shoulder_yaw_link"/>
+    <axis xyz="0 0 1"/>
+    <limit lower="-2.618" upper="2.618" effort="21" velocity="53"/>
+  </joint>
+  <link name="right_elbow_pitch_link">
+    <inertial>
+      <origin xyz="0.064497 -0.002873 0" rpy="0 0 0"/>
+      <mass value="0.601"/>
+      <inertia ixx="0.0002845" ixy="-8.06E-05" ixz="9E-09" iyy="0.0004647" iyz="-5E-09" izz="0.0004712"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_elbow_pitch_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_elbow_pitch_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_elbow_pitch_joint" type="fixed">
+    <origin xyz="0 -0.00189 -0.0855" rpy="0 0 0"/>
+    <parent link="right_shoulder_yaw_link"/>
+    <child link="right_elbow_pitch_link"/>
+    <axis xyz="0 1 0"/>
+    <limit lower="-0.2268" upper="3.4208" effort="21" velocity="53"/>
+  </joint>
+  <link name="right_elbow_roll_link">
+    <inertial>
+      <origin xyz="0.081811 -0.001454 0.001" rpy="0 0 0"/>
+      <mass value="0.227"/>
+      <inertia ixx="0.0001093" ixy="1.55E-05" ixz="-2.91E-05" iyy="0.000977" iyz="-1E-06" izz="0.0009676"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_elbow_roll_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_elbow_roll_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_elbow_roll_joint" type="fixed">
+    <origin xyz="0.1 0 0" rpy="0 0 0"/>
+    <parent link="right_elbow_pitch_link"/>
+    <child link="right_elbow_roll_link"/>
+    <axis xyz="1 0 0"/>
+    <limit lower="-2.0943" upper="2.0943" effort="21" velocity="53"/>
+  </joint>
+
+  <!-- LOGO -->
+  <joint name="logo_joint" type="fixed">
+    <origin xyz="0 0 0" rpy="0 0 0"/>
+    <parent link="torso_link"/>
+    <child link="logo_link"/>
+  </joint>
+  <link name="logo_link">
+    <inertial>
+      <origin xyz="-0.00306556489 0.00005625032 0.24526459523" rpy="0 0 0"/>
+      <mass value="0.001"/>
+      <inertia ixx="0.00000089650105" ixy="0.00000007028325" ixz="0.00000002735427" iyy="0.00000644846110" iyz="0.00000000029738" izz="0.00000688611155"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/logo_link.STL"/>
+      </geometry>
+      <material name="dark">
+        <color rgba="0.2 0.2 0.2 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/logo_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+
+  <!-- IMU -->
+  <link name="imu_link"></link>
+  <joint name="imu_joint" type="fixed">
+    <origin xyz="-0.04233868314 0.00166 0.152067" rpy="0 0 0"/>
+    <parent link="torso_link"/>
+    <child link="imu_link"/>
+  </joint>
+
+  <!-- Gripper -->
+  <joint name="left_palm_joint" type="fixed">
+    <origin xyz="0.120 0 0" rpy="0 0 0"/>
+    <parent link="left_elbow_roll_link"/>
+    <child link="left_palm_link"/>
+    <axis xyz="0 0 0"/>
+  </joint>
+  <link name="left_palm_link">
+    <inertial>
+      <origin xyz="0.05578422502 0.00148509387 -0.00032670143" rpy="0 0 0"/>
+      <mass value="0.28126043"/>
+      <inertia ixx="0.00017638101047" ixy="-0.00000652171140" ixz="0.00000021301775" iyy="0.00031084250366" iyz="0.00000000708018" izz="0.00018939350451"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_palm_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_palm_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_zero_joint" type="fixed">
+    <origin xyz="0.050 0 0" rpy="0 0 0"/>
+    <parent link="left_palm_link"/>
+    <child link="left_zero_link"/>
+    <axis xyz="0 1 0"/>
+    <limit effort="0.7" velocity="23" lower="-0.523598" upper="0.523598"/>
+  </joint>
+  <link name="left_zero_link">
+    <inertial>
+      <origin xyz="-0.02604664462 -0.00877511180 0.00001086048" rpy="0 0 0"/>
+      <mass value="0.05098927"/>
+      <inertia ixx="0.00000679821860" ixy="-0.00000098281498" ixz="-0.00000000406407" iyy="0.00000501631802" iyz="0.00000002833450" izz="0.00000868017156"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_zero_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <!-- <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_zero_link.STL"/>
+      </geometry>
+    </collision> -->
+  </link>
+  <joint name="left_one_joint" type="fixed">
+    <origin xyz="-0.02652502957 -0.0188 -0.00005" rpy="0 0 0"/>
+    <parent link="left_zero_link"/>
+    <child link="left_one_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1" upper="1.2"/>
+  </joint>
+  <link name="left_one_link">
+    <inertial>
+      <origin xyz="-0.00181538145 -0.03270571529 -0.00018587117" rpy="0 0 0"/>
+      <mass value="0.04776204"/>
+      <inertia ixx="0.00000800634866" ixy="-0.00000020811350" ixz="-0.00000004687654" iyy="0.00000366333777" iyz="0.00000005133915" izz="0.00000785913677"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_one_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="-0.001 -0.032 0" rpy="0 0 0"/>
+      <geometry>
+        <box size="0.02 0.03 0.02"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_two_joint" type="fixed">
+    <origin xyz="0 -0.0431 0" rpy="0 0 0"/>
+    <parent link="left_one_link"/>
+    <child link="left_two_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="0" upper="1.84"/>
+  </joint>
+  <link name="left_two_link">
+    <inertial>
+      <origin xyz="-0.00135257289 -0.02375910756 -0.00011966865" rpy="0 0 0"/>
+      <mass value="0.01385835"/>
+      <inertia ixx="0.00000254698095" ixy="-0.00000010086310" ixz="-0.00000000325710" iyy="0.00000094744556" iyz="0.00000003605469" izz="0.00000204288618"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_two_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_two_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_three_joint" type="fixed">
+    <origin xyz="0.085 0.004 0.02395" rpy="0 0 0"/>
+    <parent link="left_palm_link"/>
+    <child link="left_three_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1.84" upper="0.3"/>
+  </joint>
+  <link name="left_three_link">
+    <inertial>
+      <origin xyz="0.03625612321 -0.00157250459 0.00026989918" rpy="0 0 0"/>
+      <mass value="0.04884689"/>
+      <inertia ixx="0.00000383218469" ixy="0.00000003239984" ixz="0.00000001923377" iyy="0.00000890974599" iyz="-0.00000004429800" izz="0.00000865612355"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_three_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_three_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_four_joint" type="fixed">
+    <origin xyz="0.0471 -0.0036 0" rpy="0 0 0"/>
+    <parent link="left_three_link"/>
+    <child link="left_four_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1.84" upper="0"/>
+  </joint>
+  <link name="left_four_link">
+    <inertial>
+      <origin xyz="0.02375910756 0.00135257289 0.00011966865" rpy="0 0 0"/>
+      <mass value="0.01385835"/>
+      <inertia ixx="0.00000094744556" ixy="-0.00000010086310" ixz="0.00000003605469" iyy="0.00000254698095" iyz="-0.00000000325710" izz="0.00000204288618"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_four_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_four_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_five_joint" type="fixed">
+    <origin xyz="0.085 0.004 -0.02395" rpy="0 0 0"/>
+    <parent link="left_palm_link"/>
+    <child link="left_five_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1.84" upper="0.3"/>
+  </joint>
+  <link name="left_five_link">
+    <inertial>
+      <origin xyz="0.03625612321 -0.00157250459 0.00026989918" rpy="0 0 0"/>
+      <mass value="0.04884689"/>
+      <inertia ixx="0.00000383218469" ixy="0.00000003239984" ixz="0.00000001923377" iyy="0.00000890974599" iyz="-0.00000004429800" izz="0.00000865612355"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_five_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_five_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="left_six_joint" type="fixed">
+    <origin xyz="0.0471 -0.0036 0" rpy="0 0 0"/>
+    <parent link="left_five_link"/>
+    <child link="left_six_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1.84" upper="0"/>
+  </joint>
+  <link name="left_six_link">
+    <inertial>
+      <origin xyz="0.02375910756 0.00135257289 0.00011966865" rpy="0 0 0"/>
+      <mass value="0.01385835"/>
+      <inertia ixx="0.00000094744556" ixy="-0.00000010086310" ixz="0.00000003605469" iyy="0.00000254698095" iyz="-0.00000000325710" izz="0.00000204288618"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_six_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/left_six_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+
+
+  <joint name="right_palm_joint" type="fixed">
+    <origin xyz="0.120 0 0" rpy="0 0 0"/>
+    <parent link="right_elbow_roll_link"/>
+    <child link="right_palm_link"/>
+    <axis xyz="0 0 0"/>
+  </joint>
+  <link name="right_palm_link">
+    <inertial>
+      <origin xyz="0.05578422502 -0.00148509387 -0.00032670143" rpy="0 0 0"/>
+      <mass value="0.28126043"/>
+      <inertia ixx="0.00017638101047" ixy="0.00000652171140" ixz="0.00000021301775" iyy="0.00031084250366" iyz="-0.00000000708018" izz="0.00018939350451"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_palm_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_palm_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_zero_joint" type="fixed">
+    <origin xyz="0.050 0 0" rpy="0 0 0"/>
+    <parent link="right_palm_link"/>
+    <child link="right_zero_link"/>
+    <axis xyz="0 1 0"/>
+    <limit effort="0.7" velocity="23" lower="-0.523598" upper="0.523598"/>
+  </joint>
+  <link name="right_zero_link">
+    <inertial>
+      <origin xyz="-0.02604664462 0.00877511180 0.00001086048" rpy="0 0 0"/>
+      <mass value="0.05098927"/>
+      <inertia ixx="0.00000679821860" ixy="0.00000098281498" ixz="-0.00000000406407" iyy="0.00000501631802" iyz="-0.00000002833450" izz="0.00000868017156"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_zero_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <!-- <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_zero_link.STL"/>
+      </geometry>
+    </collision> -->
+  </link>
+  <joint name="right_one_joint" type="fixed">
+    <origin xyz="-0.02652502957 0.0188 -0.00005" rpy="0 0 0"/>
+    <parent link="right_zero_link"/>
+    <child link="right_one_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1.2" upper="1"/>
+  </joint>
+  <link name="right_one_link">
+    <inertial>
+      <origin xyz="-0.00181538145 0.03270571529 -0.00018587117" rpy="0 0 0"/>
+      <mass value="0.04776204"/>
+      <inertia ixx="0.00000800634866" ixy="0.00000020811350" ixz="-0.00000004687654" iyy="0.00000366333777" iyz="-0.00000005133915" izz="0.00000785913677"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_one_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="-0.001 0.032 0" rpy="0 0 0"/>
+      <geometry>
+        <box size="0.02 0.03 0.02"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_two_joint" type="fixed">
+    <origin xyz="0 0.0431 0" rpy="0 0 0"/>
+    <parent link="right_one_link"/>
+    <child link="right_two_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-1.84" upper="0"/>
+  </joint>
+  <link name="right_two_link">
+    <inertial>
+      <origin xyz="-0.00135257289 0.02375910756 -0.00011966865" rpy="0 0 0"/>
+      <mass value="0.01385835"/>
+      <inertia ixx="0.00000254698095" ixy="0.00000010086310" ixz="-0.00000000325710" iyy="0.00000094744556" iyz="-0.00000003605469" izz="0.00000204288618"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_two_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_two_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_three_joint" type="fixed">
+    <origin xyz="0.085 -0.004 0.02395" rpy="0 0 0"/>
+    <parent link="right_palm_link"/>
+    <child link="right_three_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-0.3" upper="1.84"/>
+  </joint>
+  <link name="right_three_link">
+    <inertial>
+      <origin xyz="0.03625612321 0.00157250459 0.00026989918" rpy="0 0 0"/>
+      <mass value="0.04884689"/>
+      <inertia ixx="0.00000383218469" ixy="-0.00000003239984" ixz="0.00000001923377" iyy="0.00000890974599" iyz="0.00000004429800" izz="0.00000865612355"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_three_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_three_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_four_joint" type="fixed">
+    <origin xyz="0.0471 0.0036 0" rpy="0 0 0"/>
+    <parent link="right_three_link"/>
+    <child link="right_four_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="0" upper="1.84"/>
+  </joint>
+  <link name="right_four_link">
+    <inertial>
+      <origin xyz="0.02375910756 -0.00135257289 0.00011966865" rpy="0 0 0"/>
+      <mass value="0.01385835"/>
+      <inertia ixx="0.00000094744556" ixy="0.00000010086310" ixz="0.00000003605469" iyy="0.00000254698095" iyz="0.00000000325710" izz="0.00000204288618"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_four_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_four_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_five_joint" type="fixed">
+    <origin xyz="0.085 -0.004 -0.02395" rpy="0 0 0"/>
+    <parent link="right_palm_link"/>
+    <child link="right_five_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="-0.3" upper="1.84"/>
+  </joint>
+  <link name="right_five_link">
+    <inertial>
+      <origin xyz="0.03625612321 0.00157250459 0.00026989918" rpy="0 0 0"/>
+      <mass value="0.04884689"/>
+      <inertia ixx="0.00000383218469" ixy="-0.00000003239984" ixz="0.00000001923377" iyy="0.00000890974599" iyz="0.00000004429800" izz="0.00000865612355"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_five_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_five_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+  <joint name="right_six_joint" type="fixed">
+    <origin xyz="0.0471 0.0036 0" rpy="0 0 0"/>
+    <parent link="right_five_link"/>
+    <child link="right_six_link"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="0.7" velocity="23" lower="0" upper="1.84"/>
+  </joint>
+  <link name="right_six_link">
+    <inertial>
+      <origin xyz="0.02375910756 -0.00135257289 0.00011966865" rpy="0 0 0"/>
+      <mass value="0.01385835"/>
+      <inertia ixx="0.00000094744556" ixy="0.00000010086310" ixz="0.00000003605469" iyy="0.00000254698095" iyz="0.00000000325710" izz="0.00000204288618"/>
+    </inertial>
+    <visual>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_six_link.STL"/>
+      </geometry>
+      <material name="gray">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0" rpy="0 0 0"/>
+      <geometry>
+        <mesh filename="mesh/G1/right_six_link.STL"/>
+      </geometry>
+    </collision>
+  </link>
+</robot>
\ No newline at end of file
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1/left_ankle_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1/left_ankle_pitch_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..c291f607f5476a63412fad53dc622412b633cded
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1/left_ankle_pitch_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1/left_zero_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1/left_zero_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..f50147ac89f6090d8544730d0b56792d8213ba5a
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1/left_zero_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1/right_ankle_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1/right_ankle_pitch_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..36e3a64aa65c4baa37fda12157d28774445845a5
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1/right_ankle_pitch_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1/right_zero_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1/right_zero_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..4fd999bdbd23e0ca77364bd195430e1f1eb4631a
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1/right_zero_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/left_ankle_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/left_ankle_pitch_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..69de8490184afc698f633e34bb74e65351bd4689
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/left_ankle_pitch_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/left_elbow_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/left_elbow_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..1a96d99ba469960173129084ab6dd3bf8a732a71
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/left_elbow_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/left_hand_thumb_0_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/left_hand_thumb_0_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..3028bb4d6e1ae3d30d2504259c08e4106bbacf63
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/left_hand_thumb_0_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/left_wrist_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/left_wrist_pitch_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..82cc224a8e41251d879502f9809e31d0988ec7f9
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/left_wrist_pitch_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/right_ankle_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/right_ankle_pitch_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..e77d8a2fe1e5d56fac049833d254d6ffa4f6b350
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/right_ankle_pitch_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/right_elbow_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/right_elbow_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..f259e3812efb9985d6463c04d3e8a4b53793a699
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/right_elbow_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/right_hand_thumb_0_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/right_hand_thumb_0_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..1cae7f18e16605cb9d6b1d1a0cf6e5c5c360a344
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/right_hand_thumb_0_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/right_wrist_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/right_wrist_pitch_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..da194543c40df9d492abb0553c06cc614a78caae
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/right_wrist_pitch_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/torso_constraint_L_rod_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/torso_constraint_L_rod_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..6747f3f9341bd72b3803385c135c3ce750322e9d
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/torso_constraint_L_rod_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/torso_constraint_R_rod_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/torso_constraint_R_rod_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..95cf415f72f1679a5867ca21a4af774f0b217cad
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/torso_constraint_R_rod_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/waist_roll_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/waist_roll_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..65831abd2a64bc8c36e31016964413a3d2116725
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/meshes/old/mesh/G1_23DoF/waist_roll_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1/left_ankle_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1/left_ankle_pitch_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..c291f607f5476a63412fad53dc622412b633cded
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1/left_ankle_pitch_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1/left_zero_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1/left_zero_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..f50147ac89f6090d8544730d0b56792d8213ba5a
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1/left_zero_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1/right_ankle_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1/right_ankle_pitch_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..36e3a64aa65c4baa37fda12157d28774445845a5
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1/right_ankle_pitch_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/left_ankle_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/left_ankle_pitch_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..69de8490184afc698f633e34bb74e65351bd4689
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/left_ankle_pitch_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/left_elbow_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/left_elbow_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..1a96d99ba469960173129084ab6dd3bf8a732a71
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/left_elbow_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/left_hand_thumb_0_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/left_hand_thumb_0_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..3028bb4d6e1ae3d30d2504259c08e4106bbacf63
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/left_hand_thumb_0_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/left_wrist_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/left_wrist_pitch_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..82cc224a8e41251d879502f9809e31d0988ec7f9
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/left_wrist_pitch_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/right_ankle_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/right_ankle_pitch_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..e77d8a2fe1e5d56fac049833d254d6ffa4f6b350
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/right_ankle_pitch_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/right_elbow_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/right_elbow_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..f259e3812efb9985d6463c04d3e8a4b53793a699
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/right_elbow_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/right_hand_thumb_0_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/right_hand_thumb_0_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..1cae7f18e16605cb9d6b1d1a0cf6e5c5c360a344
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/right_hand_thumb_0_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/right_wrist_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/right_wrist_pitch_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..da194543c40df9d492abb0553c06cc614a78caae
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/right_wrist_pitch_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/torso_constraint_L_rod_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/torso_constraint_L_rod_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..6747f3f9341bd72b3803385c135c3ce750322e9d
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/torso_constraint_L_rod_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/torso_constraint_R_rod_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/torso_constraint_R_rod_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..95cf415f72f1679a5867ca21a4af774f0b217cad
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/torso_constraint_R_rod_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/waist_roll_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/waist_roll_link.STL
new file mode 100644
index 0000000000000000000000000000000000000000..65831abd2a64bc8c36e31016964413a3d2116725
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robots/g1/old/mesh/G1_23DoF/waist_roll_link.STL differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/__pycache__/__init__.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d3dba20fe283986f79d30fe718c1b94ce4ce5cab
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/__pycache__/__init__.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/__pycache__/__init__.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6f578ed8d7bb868697a6f9996e182e2f521fdd24
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/__pycache__/__init__.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/env_utils/__init__.py b/GR00T-WholeBodyControl/gear_sonic/envs/env_utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/env_utils/__pycache__/__init__.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/env_utils/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..417b8c11ee237bac3c627f4b255236dd4ea61b30
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/env_utils/__pycache__/__init__.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/env_utils/__pycache__/__init__.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/env_utils/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..dadcb85dd9781533b506059585e423a811d53a7c
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/env_utils/__pycache__/__init__.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/env_utils/__pycache__/joint_utils.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/env_utils/__pycache__/joint_utils.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..11e86929959cfd96a737bd716a002d51921632c5
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/env_utils/__pycache__/joint_utils.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/env_utils/__pycache__/joint_utils.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/env_utils/__pycache__/joint_utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9d4797c97424bf4cc854651584667c79a37a6779
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/env_utils/__pycache__/joint_utils.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/env_utils/joint_utils.py b/GR00T-WholeBodyControl/gear_sonic/envs/env_utils/joint_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..b3a1c6a40cf147c017da165fd65b335d89c0c52d
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/envs/env_utils/joint_utils.py
@@ -0,0 +1,85 @@
+"""Joint utility functions and constants for G1 robot.
+
+This module provides joint ordering constants and helper functions for mapping
+between motion library data and robot joints.
+"""
+
+import torch
+
+# G1 body joint names in IsaacLab order (29 DOF)
+G1_ISAACLab_ORDER = [
+    "left_hip_pitch_joint",
+    "right_hip_pitch_joint",
+    "waist_yaw_joint",
+    "left_hip_roll_joint",
+    "right_hip_roll_joint",
+    "waist_roll_joint",
+    "left_hip_yaw_joint",
+    "right_hip_yaw_joint",
+    "waist_pitch_joint",
+    "left_knee_joint",
+    "right_knee_joint",
+    "left_shoulder_pitch_joint",
+    "right_shoulder_pitch_joint",
+    "left_ankle_pitch_joint",
+    "right_ankle_pitch_joint",
+    "left_shoulder_roll_joint",
+    "right_shoulder_roll_joint",
+    "left_ankle_roll_joint",
+    "right_ankle_roll_joint",
+    "left_shoulder_yaw_joint",
+    "right_shoulder_yaw_joint",
+    "left_elbow_joint",
+    "right_elbow_joint",
+    "left_wrist_roll_joint",
+    "right_wrist_roll_joint",
+    "left_wrist_pitch_joint",
+    "right_wrist_pitch_joint",
+    "left_wrist_yaw_joint",
+    "right_wrist_yaw_joint",
+]
+
+# G1 hand joint names (14 DOF) - order from g1_43dof.yaml
+G1_HAND_JOINTS = [
+    "left_hand_index_0_joint",
+    "left_hand_index_1_joint",
+    "left_hand_middle_0_joint",
+    "left_hand_middle_1_joint",
+    "left_hand_thumb_0_joint",
+    "left_hand_thumb_1_joint",
+    "left_hand_thumb_2_joint",
+    "right_hand_index_0_joint",
+    "right_hand_index_1_joint",
+    "right_hand_middle_0_joint",
+    "right_hand_middle_1_joint",
+    "right_hand_thumb_0_joint",
+    "right_hand_thumb_1_joint",
+    "right_hand_thumb_2_joint",
+]
+
+# Caches for joint indices
+_body_joint_indices_cache = {}
+_hand_joint_indices_cache = {}
+
+
+def _get_joint_indices_by_names(asset, joint_names: list, cache: dict) -> torch.Tensor:
+    """Get indices of specified joints in the robot's joint list."""
+    cache_key = (id(asset), tuple(joint_names))
+    if cache_key in cache:
+        return cache[cache_key]
+
+    robot_joint_names = asset.joint_names
+    indices = [robot_joint_names.index(n) for n in joint_names if n in robot_joint_names]
+    indices_tensor = torch.tensor(indices, dtype=torch.long, device=asset.device)
+    cache[cache_key] = indices_tensor
+    return indices_tensor
+
+
+def get_body_joint_indices(asset) -> torch.Tensor:
+    """Get indices of body joints (29 DOF) using G1_ISAACLab_ORDER."""
+    return _get_joint_indices_by_names(asset, G1_ISAACLab_ORDER, _body_joint_indices_cache)
+
+
+def get_hand_joint_indices(asset) -> torch.Tensor:
+    """Get indices of hand joints (14 DOF) using G1_HAND_JOINTS."""
+    return _get_joint_indices_by_names(asset, G1_HAND_JOINTS, _hand_joint_indices_cache)
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/__init__.py b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/__pycache__/__init__.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3d75a534a9dc26aec923cb68947ba04ff2bbd7fb
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/__pycache__/__init__.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/__pycache__/__init__.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..65ccee45578f53f8963f346ddc2027d70334dbac
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/__pycache__/__init__.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/__pycache__/modular_tracking_env_cfg.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/__pycache__/modular_tracking_env_cfg.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..704b03091d4a93128ac87447d0175942847d3ee8
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/__pycache__/modular_tracking_env_cfg.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/__pycache__/modular_tracking_env_cfg.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/__pycache__/modular_tracking_env_cfg.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..84649e24753be2319f218ac4b1e2d05859df77dc
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/__pycache__/modular_tracking_env_cfg.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__init__.py b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..4344d19fe2439b2f951dc9a523ffa40f25cee770
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__init__.py
@@ -0,0 +1,12 @@
+"""This sub-module contains the functions that are specific to the locomotion environments."""
+
+from isaaclab.envs.mdp import *  # noqa: F401, F403
+
+from gear_sonic.envs.manager_env.mdp.actions import *  # noqa: F401, F403
+from gear_sonic.envs.manager_env.mdp.commands import *  # noqa: F401, F403
+from gear_sonic.envs.manager_env.mdp.curriculum import *  # noqa: F401, F403
+from gear_sonic.envs.manager_env.mdp.events import *  # noqa: F401, F403
+from gear_sonic.envs.manager_env.mdp.observations import *  # noqa: F401, F403
+from gear_sonic.envs.manager_env.mdp.recorders import *  # noqa: F401, F403
+from gear_sonic.envs.manager_env.mdp.rewards import *  # noqa: F401, F403
+from gear_sonic.envs.manager_env.mdp.terminations import *  # noqa: F401, F403
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/__init__.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7712c34a5856c8430c5bdb941d88964f34ccd529
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/__init__.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/__init__.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..329967b714dafa6301fff2a8778c89c86f1cc8b0
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/__init__.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/actions.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/actions.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f564b2b2f46d6a4f1adf8a003a34eb882be005b6
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/actions.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/actions.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/actions.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d90a9b6289c7ab1776642a5b44f8152d4e6cccae
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/actions.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/curriculum.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/curriculum.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0d4702b315d50d962fdfa90d248ca62ce7d434a2
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/curriculum.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/curriculum.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/curriculum.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..22f7f687aa9f5bc355ad26fb9f25817322daf6de
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/curriculum.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/events.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/events.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3669b9146e91d06fe7e77d78ca09532adfdabf28
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/events.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/events.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/events.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..014b418cb72254d3ec1d381642d35fd2bb89eec2
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/events.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/observations.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/observations.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..81468eae592dab89b4d847704f7b1a7ba7f8c710
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/observations.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/observations.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/observations.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..20b523d447e9b06533d90a78accdd56a57db2775
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/observations.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/recorders.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/recorders.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..97cad030a70d2da31442facbd4d5cd0ed9acd382
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/recorders.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/recorders.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/recorders.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..719b09471dbd8d46f5d704391cf867b7f641bee3
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/recorders.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/rewards.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/rewards.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cda577aa418d1d42fa0ba60c29a6289c8216fa13
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/rewards.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/rewards.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/rewards.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c174f2744b2a319e6e31e9115e0a134a03ca1277
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/rewards.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/terminations.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/terminations.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..39c2ee68f5b73206ddb014f49059d72a7ebe197a
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/terminations.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/terminations.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/terminations.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b846a258f608f8b08a8c6c291ce0bc55186a2cc8
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/terminations.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/terrain.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/terrain.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..dd102912666e5d29dd0a3e3ab5bb4eb054cf2344
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/terrain.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/terrain.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/terrain.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ad4f85ad352d6894f232890476484d96d851363d
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/terrain.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/utils.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/utils.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7670735eb6d993ca70e62685f131b9b8b17dddce
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/utils.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/utils.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..12ebab39db9dea4cdb15da7fad832e7309def7ba
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/__pycache__/utils.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/actions.py b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/actions.py
new file mode 100644
index 0000000000000000000000000000000000000000..3a637df4b8cc5e7a92be2d5231d39f55f0d4abee
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/actions.py
@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from isaaclab.utils import configclass
+
+# Joint ordering constants
+G1_MUJOCO_ORDER = [
+    "left_hip_pitch_joint",
+    "left_hip_roll_joint",
+    "left_hip_yaw_joint",
+    "left_knee_joint",
+    "left_ankle_pitch_joint",
+    "left_ankle_roll_joint",
+    "right_hip_pitch_joint",
+    "right_hip_roll_joint",
+    "right_hip_yaw_joint",
+    "right_knee_joint",
+    "right_ankle_pitch_joint",
+    "right_ankle_roll_joint",
+    "waist_yaw_joint",
+    "waist_roll_joint",
+    "waist_pitch_joint",
+    "left_shoulder_pitch_joint",
+    "left_shoulder_roll_joint",
+    "left_shoulder_yaw_joint",
+    "left_elbow_joint",
+    "left_wrist_roll_joint",
+    "left_wrist_pitch_joint",
+    "left_wrist_yaw_joint",
+    "right_shoulder_pitch_joint",
+    "right_shoulder_roll_joint",
+    "right_shoulder_yaw_joint",
+    "right_elbow_joint",
+    "right_wrist_roll_joint",
+    "right_wrist_pitch_joint",
+    "right_wrist_yaw_joint",
+]
+
+
+@configclass
+class ActionsCfg:
+    """Action specifications for the MDP."""
+
+    joint_pos = None
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/actuators.py b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/actuators.py
new file mode 100644
index 0000000000000000000000000000000000000000..a3be306d8307c0cb1f869aca34fc80f5330b37ae
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/actuators.py
@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from isaaclab.actuators import ImplicitActuator, ImplicitActuatorCfg
+from isaaclab.utils import DelayBuffer, configclass
+from isaaclab.utils.types import ArticulationActions
+import torch
+
+
+class DelayedImplicitActuator(ImplicitActuator):
+    """Ideal PD actuator with delayed command application.
+
+    This class extends the :class:`IdealPDActuator` class by adding a delay to the actuator commands. The delay
+    is implemented using a circular buffer that stores the actuator commands for a certain number of physics steps.
+    The most recent actuation value is pushed to the buffer at every physics step, but the final actuation value
+    applied to the simulation is lagged by a certain number of physics steps.
+
+    The amount of time lag is configurable and can be set to a random value between the minimum and maximum time
+    lag bounds at every reset. The minimum and maximum time lag values are set in the configuration instance passed
+    to the class.
+    """
+
+    cfg: DelayedImplicitActuatorCfg
+    """The configuration for the actuator model."""
+
+    def __init__(self, cfg: DelayedImplicitActuatorCfg, *args, **kwargs):
+        super().__init__(cfg, *args, **kwargs)
+        # instantiate the delay buffers
+        self.positions_delay_buffer = DelayBuffer(
+            cfg.max_delay, self._num_envs, device=self._device
+        )
+        self.velocities_delay_buffer = DelayBuffer(
+            cfg.max_delay, self._num_envs, device=self._device
+        )
+        self.efforts_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device)
+        # all of the envs
+        self._ALL_INDICES = torch.arange(self._num_envs, dtype=torch.long, device=self._device)
+
+    def reset(self, env_ids: Sequence[int]):
+        super().reset(env_ids)
+        # number of environments (since env_ids can be a slice)
+        if env_ids is None or env_ids == slice(None):
+            num_envs = self._num_envs
+        else:
+            num_envs = len(env_ids)
+        # set a new random delay for environments in env_ids
+        time_lags = torch.randint(
+            low=self.cfg.min_delay,
+            high=self.cfg.max_delay + 1,
+            size=(num_envs,),
+            dtype=torch.int,
+            device=self._device,
+        )
+        # set delays
+        self.positions_delay_buffer.set_time_lag(time_lags, env_ids)
+        self.velocities_delay_buffer.set_time_lag(time_lags, env_ids)
+        self.efforts_delay_buffer.set_time_lag(time_lags, env_ids)
+        # reset buffers
+        self.positions_delay_buffer.reset(env_ids)
+        self.velocities_delay_buffer.reset(env_ids)
+        self.efforts_delay_buffer.reset(env_ids)
+
+    def compute(
+        self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor
+    ) -> ArticulationActions:
+        # apply delay based on the delay the model for all the setpoints
+        control_action.joint_positions = self.positions_delay_buffer.compute(
+            control_action.joint_positions
+        )
+        control_action.joint_velocities = self.velocities_delay_buffer.compute(
+            control_action.joint_velocities
+        )
+        control_action.joint_efforts = self.efforts_delay_buffer.compute(
+            control_action.joint_efforts
+        )
+        # compte actuator model
+        return super().compute(control_action, joint_pos, joint_vel)
+
+
+@configclass
+class DelayedImplicitActuatorCfg(ImplicitActuatorCfg):
+    """Configuration for a delayed PD actuator."""
+
+    class_type: type = DelayedImplicitActuator
+
+    min_delay: int = 0
+    """Minimum number of physics time-steps with which the actuator command may be delayed. Defaults to 0."""
+
+    max_delay: int = 0
+    """Maximum number of physics time-steps with which the actuator command may be delayed. Defaults to 0."""
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/commands.py b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/commands.py
new file mode 100644
index 0000000000000000000000000000000000000000..529ff3ce9c334018614fed89c75354387a8f26eb
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/commands.py
@@ -0,0 +1,4292 @@
+"""Motion tracking command terms for humanoid whole-body control RL environments."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+import copy
+import dataclasses
+import glob
+import os
+from typing import TYPE_CHECKING
+
+import easydict
+from isaaclab.assets import Articulation
+from isaaclab.managers import CommandTerm, CommandTermCfg
+from isaaclab.markers import VisualizationMarkers, VisualizationMarkersCfg
+from isaaclab.markers.config import DEFORMABLE_TARGET_MARKER_CFG
+import isaaclab.sim as sim_utils
+from isaaclab.utils import configclass
+from isaaclab.utils.math import (
+    matrix_from_quat,
+    quat_apply,
+    quat_apply_yaw,
+    quat_error_magnitude,
+    quat_from_euler_xyz,
+    quat_inv,
+    quat_mul,
+    sample_uniform,
+)
+import numpy as np
+import torch
+
+from gear_sonic.envs.env_utils import joint_utils
+from gear_sonic.isaac_utils import rotations
+from gear_sonic.trl.utils import common, order_converter, torch_transform
+from gear_sonic.utils.motion_lib import motion_lib_robot
+
+if TYPE_CHECKING:
+    from isaaclab.envs import ManagerBasedRLEnv
+
+# Constants for multi-object mode: inactive objects placed far away
+# Objects spread vertically (Z) since envs only vary in X,Y - much simpler!
+INACTIVE_OBJECT_BASE_OFFSET = torch.tensor([1000.0, 0.0, -50.0])  # 1km away in X, 50m underground
+INACTIVE_OBJECT_Z_SPACING = 10.0  # 10m vertical spacing between objects (must be > chair height)
+
+
+def _init_variable_frames(
+    enabled: bool,
+    min_frames: int,
+    num_future_frames: int,
+    step: int,
+    num_envs: int,
+    device: torch.device,
+) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    """Initialize variable frame support tensors.
+
+    Returns (per_env_num_frames, frame_choices) or (None, None) if disabled.
+    """
+    if not enabled:
+        return None, None
+    assert step > 0, f"variable_frames_step must be positive, got {step}"
+    per_env_num_frames = torch.full((num_envs,), num_future_frames, device=device, dtype=torch.long)
+    frame_choices = torch.arange(min_frames, num_future_frames + 1, step, device=device)
+    assert len(frame_choices) > 0, (
+        f"No valid frame choices: variable_frames_min={min_frames} "
+        f"> num_future_frames={num_future_frames}"
+    )
+    return per_env_num_frames, frame_choices
+
+
+@configclass
+class CommandsCfg:
+    """Command specifications for the MDP."""
+
+    motion = None
+    force = None
+
+
+class TrackingCommand(CommandTerm):
+    """Provide reference motion trajectories for motion-tracking RL.
+
+    This is the primary command term for SONIC-style humanoid control. It manages
+    a motion library of pre-recorded motion clips (robot joint trajectories, SMPL
+    poses, and optionally object trajectories) and serves reference frames to the
+    policy and reward system at each simulation step.
+
+    Key responsibilities:
+        - Load and index motion clips from the motion library (MotionLibRobot).
+        - Sample motion IDs and start times for each environment at episode reset.
+        - Advance the time cursor each step and provide current + multi-future
+          reference frames (positions, orientations, velocities, joint states).
+        - Transform reference quantities into robot-local, egocentric, or
+          heading-canonicalized coordinate frames for observation terms.
+        - Handle DOF mismatch between motion data and robot (e.g., extra finger
+          joints not present in the motion library).
+        - Support adaptive sampling, contact-based initialization, variable future
+          frame counts, encoder mode sampling (G1/SMPL/teleop), and multi-object
+          scene management.
+
+    The command exposes ~60 property-based accessors that observation and reward
+    terms query each step. Properties are named by the pattern::
+
+        {quantity}[_multi_future][_dif][_{frame}]
+
+    where ``quantity`` is body_pos, joint_pos, smpl_pose, etc.; ``multi_future``
+    means values for all ``num_future_frames`` reference frames stacked;
+    ``dif`` is the difference relative to the robot; and ``frame`` is ``w``
+    (world), ``l`` (robot-local / de-headed), or ``b`` (body).
+    """
+
+    cfg: TrackingCommandCfg
+
+    def __init__(
+        self, cfg: TrackingCommandCfg, env: ManagerBasedRLEnv, max_num_load_motions: int = None  # noqa: RUF013
+    ):
+        """Initialize the tracking command with motion library and body mappings.
+
+        Loads motion data, builds DOF/body index mappings between the motion
+        library (MuJoCo ordering) and the IsaacLab robot, pre-allocates per-env
+        tensors for motion IDs, time steps, and future frame offsets, and
+        optionally sets up height-map raycasting.
+
+        Args:
+            cfg: Configuration dataclass specifying motion files, body names,
+                future frame settings, encoder sampling probabilities, etc.
+            env: The manager-based RL environment that owns this command.
+            max_num_load_motions: Cap on how many unique motion clips to load
+                into GPU memory. Defaults to ``min(num_envs, 1024)`` or the
+                full library when ``use_paired_motions`` is enabled.
+        """
+        super().__init__(cfg, env)
+
+        self.is_evaluating = False
+        self.cmd_body_names = self.cfg.body_names
+        self.robot: Articulation = env.scene[cfg.asset_name]
+        self.robot_anchor_body_index = self.robot.body_names.index(self.cfg.anchor_body)
+        self.motion_anchor_body_index = self.cfg.body_names.index(self.cfg.anchor_body)
+        self.vr_3point_body_indices = [
+            self.robot.body_names.index(name) for name in self.cfg.vr_3point_body
+        ]
+        self.vr_3point_body_indices_motion = [
+            self.cfg.body_names.index(name) for name in self.cfg.vr_3point_body
+        ]
+        self.vr_3point_body_offsets = (
+            torch.tensor(self.cfg.vr_3point_body_offset, dtype=torch.float32, device=self.device)
+            .view(1, -1, 3)
+            .repeat(self.num_envs, 1, 1)
+        )
+
+        self.reward_point_body_indices = [
+            self.robot.body_names.index(name) for name in self.cfg.reward_point_body
+        ]
+        self.reward_point_body_offsets = (
+            torch.tensor(self.cfg.reward_point_body_offset, dtype=torch.float32, device=self.device)
+            .view(1, -1, 3)
+            .repeat(self.num_envs, 1, 1)
+        )
+        self.reward_point_body_indices_motion = [
+            self.cfg.body_names.index(name) for name in self.cfg.reward_point_body
+        ]
+
+        self.down_dir = (
+            torch.tensor([0.0, 0.0, -1.0], dtype=torch.float32, device=self.device)
+            .view(1, -1)
+            .repeat(self.num_envs, 1)
+        )
+        self.body_indexes = torch.tensor(
+            self.robot.find_bodies(self.cfg.body_names, preserve_order=True)[0],
+            dtype=torch.long,
+            device=self.device,
+        )
+
+        isaac_lab_joints = env.cfg.isaaclab_to_mujoco_mapping["isaaclab_joints"]
+
+        self.isaaclab_to_mujoco_dof = env.cfg.isaaclab_to_mujoco_mapping["isaaclab_to_mujoco_dof"]
+        self.mujoco_to_isaaclab_dof = env.cfg.isaaclab_to_mujoco_mapping["mujoco_to_isaaclab_dof"]
+        self.lower_joint_indices_mujoco = list(range(12))
+        self.lower_joint_isaaclab_indices = [
+            self.isaaclab_to_mujoco_dof[i] for i in self.lower_joint_indices_mujoco
+        ]
+        self.isaaclab_to_mujoco_body = env.cfg.isaaclab_to_mujoco_mapping["isaaclab_to_mujoco_body"]
+        self.mujoco_to_isaaclab_body = env.cfg.isaaclab_to_mujoco_mapping["mujoco_to_isaaclab_body"]
+        self.running_ref_root_height = torch.zeros(
+            self.num_envs, dtype=torch.float, device=self.device
+        )
+        self.body_indexes_data = [isaac_lab_joints.index(name) for name in self.cfg.body_names]
+        if self.cfg.motion_lib_cfg is not None:
+            motion_lib_cfg = easydict.EasyDict(self.cfg.motion_lib_cfg)
+        else:
+            motion_lib_cfg = easydict.EasyDict(
+                {
+                    "motion_file": self.cfg.motion_file,
+                    "smpl_motion_file": getattr(self.cfg, "smpl_motion_file", None),
+                    "asset": {
+                        "assetRoot": "gear_sonic/data/assets/robot_description/mjcf/",
+                        "assetFileName": "g1_29dof_rev_1_0.xml",
+                        "urdfFileName": "",
+                    },
+                    "extend_config": [],
+                    "target_fps": 50,
+                    "multi_thread": True,
+                    "filter_motion_keys": self.cfg.filter_motion_keys,
+                }
+            )
+        # Only override filter_motion_keys if explicitly set at command level
+        # (don't overwrite the value from motion_lib_cfg if it exists there)
+        filter_keys = self.cfg.filter_motion_keys
+        if filter_keys is None:
+            filter_keys = motion_lib_cfg.get("filter_motion_keys", None)
+
+        motion_lib_cfg.update(
+            {
+                "mujoco_to_isaaclab_dof": self.mujoco_to_isaaclab_dof,
+                "mujoco_to_isaaclab_body": self.mujoco_to_isaaclab_body,
+                "isaaclab_to_mujoco_dof": self.isaaclab_to_mujoco_dof,
+                "isaaclab_to_mujoco_body": self.isaaclab_to_mujoco_body,
+                "body_indexes": self.body_indexes,
+                "body_indexes_data": self.body_indexes_data,
+                "filter_motion_keys": filter_keys,
+                "lower_joint_indices_mujoco": self.lower_joint_indices_mujoco,
+                "cat_upper_body_poses": self.cfg.cat_upper_body_poses,
+                "cat_upper_body_poses_prob": self.cfg.cat_upper_body_poses_prob,
+                "randomize_heading": self.cfg.randomize_heading,
+                "freeze_frame_aug": self.cfg.freeze_frame_aug,
+                "freeze_frame_aug_prob": self.cfg.freeze_frame_aug_prob,
+                "randomize_wrist_poses": self.cfg.randomize_wrist_poses,
+                "randomize_wrist_prob": self.cfg.randomize_wrist_prob,
+                "randomize_wrist_std": self.cfg.randomize_wrist_std,
+            }
+        )
+
+        self.motion_lib = motion_lib_robot.MotionLibRobot(
+            motion_lib_cfg, self.num_envs, self.device
+        )
+        if max_num_load_motions is None:
+            if self.cfg.use_paired_motions:
+                self.max_num_load_motions = self.motion_lib._num_unique_motions  # noqa: SLF001
+            else:
+                self.max_num_load_motions = min(self.num_envs, 1024)
+        else:
+            self.max_num_load_motions = max_num_load_motions
+        self.motion_lib.load_motions_for_training(max_num_seqs=self.max_num_load_motions)
+        self.use_adaptive_sampling = self.motion_lib.use_adaptive_sampling
+
+        # Load contact data for contact-based initialization
+        self._load_contact_data()
+
+        # Setup DOF mapping for handling mismatch between motion library and robot
+
+        self.robot_num_dof = self.robot.num_joints
+        self.motion_lib_num_dof = self.cfg.motion_lib_num_dof
+        if self.motion_lib_num_dof is None:
+            self.motion_lib_num_dof = self.robot_num_dof
+
+        self.extra_num_dof = self.robot_num_dof - self.motion_lib_num_dof
+        self.has_dof_mismatch = self.extra_num_dof > 0
+
+        if self.has_dof_mismatch:
+            self.body_joint_indices = joint_utils.get_body_joint_indices(self.robot)
+            self.extra_joint_indices = joint_utils.get_hand_joint_indices(self.robot)
+            self.extra_default_positions = torch.tensor(
+                self.cfg.hand_default_positions or [0.0] * self.extra_num_dof,
+                dtype=torch.float32,
+                device=self.device,
+            )
+            self.extra_default_velocities = torch.tensor(
+                self.cfg.hand_default_velocities or [0.0] * self.extra_num_dof,
+                dtype=torch.float32,
+                device=self.device,
+            )
+
+        # Step 1: Select which motions to use
+        if self.cfg.use_paired_motions:
+            # Assign motion IDs sequentially (wraps around if more envs than motions)
+            self.motion_ids = (
+                torch.arange(self.num_envs, device=self.device)
+                % self.motion_lib._num_motions  # noqa: SLF001
+            )
+        elif getattr(self.cfg, "sample_unique_motions", False):
+            # Sample without replacement - each env gets a unique motion
+            num_available = len(self.motion_lib._curr_motion_ids)  # noqa: SLF001
+            if self.num_envs > num_available:
+                raise ValueError(
+                    f"sample_unique_motions=True requires num_envs ({self.num_envs}) <= "
+                    f"num_available_motions ({num_available})"
+                )
+            perm = torch.randperm(num_available, device=self.device)[: self.num_envs]
+            self.motion_ids = perm
+            print(  # noqa: T201
+                f"[TrackingCommand] Sampled {self.num_envs} unique motions (no duplicates)"
+            )
+        else:
+            # Random sampling (can have duplicates)
+            self.motion_ids = self.motion_lib.sample_motions(self.num_envs)
+
+        # Step 2: Sample start time steps for selected motions
+        self.motion_start_time_steps = self.motion_lib.sample_time_steps(
+            self.motion_ids, truncate_time=None
+        )
+
+        # # Debug: print motion assignments
+        # motion_keys = self.motion_lib._motion_data_keys
+        # for env_id in range(self.num_envs):
+        #     motion_id = int(self.motion_ids[env_id].item())
+        #     motion_key = motion_keys[motion_id]
+        #     print(f"env {env_id}: motion_id={motion_id}, motion_key={motion_key}")
+
+        # Step 3: Override start time steps if configured
+        if self.cfg.sample_from_n_initial_frames is not None:
+            # Sample uniformly from first N frames
+            n_frames = self.cfg.sample_from_n_initial_frames
+            self.motion_start_time_steps = torch.randint(
+                0,
+                n_frames,
+                (self.num_envs,),
+                dtype=self.motion_start_time_steps.dtype,
+                device=self.device,
+            )
+        elif self.cfg.start_from_first_frame:
+            self.motion_start_time_steps.zero_()
+        self.motion_num_steps = self.motion_lib.get_motion_num_steps(self.motion_ids)
+
+        self.time_steps = torch.zeros(self.num_envs, dtype=torch.long, device=self.device)
+
+        # Object position randomization offset (per-env, resampled at reset)
+        self._object_position_offset = torch.zeros(self.num_envs, 3, device=self.device)
+
+        self.body_pos_relative_w = torch.zeros(
+            self.num_envs, len(cfg.body_names), 3, device=self.device
+        )
+        self.body_quat_relative_w = torch.zeros(
+            self.num_envs, len(cfg.body_names), 4, device=self.device
+        )
+        self.body_quat_relative_w[:, :, 0] = 1.0
+
+        self.num_future_frames = self.cfg.num_future_frames
+        # Motion lib is at target_fps; ref frames are spaced by dt_future_ref_frames (seconds).
+        # frame_skips = number of motion-lib steps between consecutive ref frames (integer).
+        # Effective spacing equals dt_future_ref_frames only when (dt_future_ref_frames * target_fps) is an integer;  # noqa: E501
+        # otherwise integer division truncates and the velocity calculation will be incorrect.
+        self.frame_skips = self.cfg.dt_future_ref_frames // (1.0 / motion_lib_cfg.target_fps)
+        steps_exact = self.cfg.dt_future_ref_frames * motion_lib_cfg.target_fps
+        if abs(steps_exact - round(steps_exact)) > 1e-9:
+            import warnings
+
+            warnings.warn(
+                f"dt_future_ref_frames={self.cfg.dt_future_ref_frames} * target_fps={motion_lib_cfg.target_fps} "
+                f"= {steps_exact} is not an integer; "
+                f"using frame_skips={self.frame_skips} so effective ref-frame spacing is "
+                f"{self.frame_skips / motion_lib_cfg.target_fps:.4f}s (not {self.cfg.dt_future_ref_frames}s). "
+                "Velocity-based losses may use incorrect dt.",
+                stacklevel=2,
+            )
+
+        self.future_time_steps_init = (
+            (
+                torch.arange(self.num_future_frames, device=self.device, dtype=torch.long)
+                * self.frame_skips
+            )
+            .view(1, -1)
+            .repeat(self.num_envs, 1)
+        )
+
+        if self.cfg.smpl_num_future_frames is None:
+            self.smpl_num_future_frames = self.num_future_frames
+        else:
+            self.smpl_num_future_frames = self.cfg.smpl_num_future_frames
+        if self.cfg.smpl_dt_future_ref_frames is None:
+            self.smpl_dt_future_ref_frames = self.cfg.dt_future_ref_frames
+        else:
+            self.smpl_dt_future_ref_frames = self.cfg.smpl_dt_future_ref_frames
+
+        self.smpl_frame_skips = self.smpl_dt_future_ref_frames // (1.0 / motion_lib_cfg.target_fps)
+        self.smpl_future_time_steps_init = (
+            (
+                torch.arange(self.smpl_num_future_frames, device=self.device, dtype=torch.long)
+                * self.smpl_frame_skips
+            )
+            .view(1, -1)
+            .repeat(self.num_envs, 1)
+        )
+
+        self.future_motion_ids = self.motion_ids.repeat_interleave(self.num_future_frames)
+        self.smpl_future_motion_ids = self.motion_ids.repeat_interleave(self.smpl_num_future_frames)
+
+        # Variable frame support
+        self.variable_frames_enabled = getattr(self.cfg, "variable_frames_enabled", False)
+        self.per_env_num_frames, self._frame_choices = _init_variable_frames(
+            self.variable_frames_enabled,
+            getattr(self.cfg, "variable_frames_min", 16),
+            self.num_future_frames,
+            getattr(self.cfg, "variable_frames_step", 4),
+            self.num_envs,
+            self.device,
+        )
+
+        self.encoder_sample_probs_dict = self.cfg.encoder_sample_probs
+        self.optimize_encoders_ratio_for_CHIP = getattr(
+            self.cfg, "optimize_encoders_ratio_for_CHIP", False
+        )
+        self.encoder_sample_probs = None
+        if self.encoder_sample_probs_dict is not None:
+            self.encoder_sample_probs = torch.tensor(
+                list(self.encoder_sample_probs_dict.values()), device=self.device
+            )
+            self.encoder_sample_probs = self.encoder_sample_probs / self.encoder_sample_probs.sum()
+            self.encoder_sample_probs_no_smpl_dict = copy.deepcopy(self.encoder_sample_probs_dict)
+
+            if "smpl" in self.encoder_sample_probs_no_smpl_dict:
+                self.encoder_sample_probs_no_smpl_dict["smpl"] = 0.0
+            self.encoder_sample_probs_no_smpl = torch.tensor(
+                list(self.encoder_sample_probs_no_smpl_dict.values()), device=self.device
+            )
+            self.encoder_sample_probs_no_smpl = (
+                self.encoder_sample_probs_no_smpl / self.encoder_sample_probs_no_smpl.sum()
+            )
+            self.g1_encoder_index = list(self.encoder_sample_probs_dict.keys()).index("g1")
+            if "smpl" in self.encoder_sample_probs_dict:
+                self.smpl_encoder_index = list(self.encoder_sample_probs_dict.keys()).index("smpl")
+            else:
+                self.smpl_encoder_index = None
+
+            if "teleop" in self.encoder_sample_probs_dict:
+                self.teleop_encoder_index = list(self.encoder_sample_probs_dict.keys()).index(
+                    "teleop"
+                )
+            else:
+                self.teleop_encoder_index = None
+
+            if "soma" in self.encoder_sample_probs_dict:
+                self.soma_encoder_index = list(self.encoder_sample_probs_dict.keys()).index("soma")
+                encoder_sample_probs_no_soma_dict = copy.deepcopy(self.encoder_sample_probs_dict)
+                encoder_sample_probs_no_soma_dict["soma"] = 0.0
+                self.encoder_sample_probs_no_soma = torch.tensor(
+                    list(encoder_sample_probs_no_soma_dict.values()), device=self.device
+                )
+                no_soma_sum = self.encoder_sample_probs_no_soma.sum()
+                if no_soma_sum > 0:
+                    self.encoder_sample_probs_no_soma = (
+                        self.encoder_sample_probs_no_soma / no_soma_sum
+                    )
+                else:
+                    # All probs zero (e.g. use_encoder=soma forced) — fall back to G1
+                    self.encoder_sample_probs_no_soma[self.g1_encoder_index] = 1.0
+                encoder_sample_probs_no_smpl_no_soma_dict = copy.deepcopy(
+                    self.encoder_sample_probs_no_smpl_dict
+                )
+                if "soma" in encoder_sample_probs_no_smpl_no_soma_dict:
+                    encoder_sample_probs_no_smpl_no_soma_dict["soma"] = 0.0
+                self.encoder_sample_probs_no_smpl_no_soma = torch.tensor(
+                    list(encoder_sample_probs_no_smpl_no_soma_dict.values()), device=self.device
+                )
+                no_smpl_no_soma_sum = self.encoder_sample_probs_no_smpl_no_soma.sum()
+                if no_smpl_no_soma_sum > 0:
+                    self.encoder_sample_probs_no_smpl_no_soma = (
+                        self.encoder_sample_probs_no_smpl_no_soma / no_smpl_no_soma_sum
+                    )
+                else:
+                    # All probs zero — fall back to G1
+                    self.encoder_sample_probs_no_smpl_no_soma[self.g1_encoder_index] = 1.0
+            else:
+                self.soma_encoder_index = None
+                self.encoder_sample_probs_no_soma = None
+                self.encoder_sample_probs_no_smpl_no_soma = None
+
+            self.teleop_sample_prob_when_smpl = self.cfg.teleop_sample_prob_when_smpl
+
+            self.encoder_index = torch.zeros(
+                (self.num_envs, self.encoder_sample_probs.shape[0]),
+                dtype=torch.long,
+                device=self.device,
+            )
+
+            if "smpl" in self.encoder_sample_probs_dict:
+                assert (
+                    self.smpl_encoder_index > self.g1_encoder_index
+                ), f"SMPL encoder index {self.smpl_encoder_index} must be greater than G1 encoder index {self.g1_encoder_index} to when both exist!"  # noqa: E501
+
+        self.metrics["error_anchor_pos"] = torch.zeros(self.num_envs, device=self.device)
+        self.metrics["error_anchor_rot"] = torch.zeros(self.num_envs, device=self.device)
+        self.metrics["error_anchor_lin_vel"] = torch.zeros(self.num_envs, device=self.device)
+        self.metrics["error_anchor_ang_vel"] = torch.zeros(self.num_envs, device=self.device)
+        self.metrics["error_body_pos"] = torch.zeros(self.num_envs, device=self.device)
+        self.metrics["error_body_rot"] = torch.zeros(self.num_envs, device=self.device)
+        self.metrics["error_joint_pos"] = torch.zeros(self.num_envs, device=self.device)
+        self.metrics["error_joint_vel"] = torch.zeros(self.num_envs, device=self.device)
+
+        self.use_ref_motion_root_quat_w_as_anchor = False
+        self.ref_motion_root_rotation_noise = None
+        self.num_bodies = len(self.cfg.body_names)
+
+        # Multi-object mode detection: check for multiple object_* entries in scene
+        # (dynamically detected to avoid storing metadata on scene config)
+        self._multi_object_mode = False
+        self._object_names = []  # List of (safe_name, original_name) tuples
+        self._active_object_name = None
+
+        # Check if we have multiple objects (object_*) vs single object ("object")
+        if hasattr(self._env, "scene") and hasattr(self._env.scene, "rigid_objects"):
+            rigid_obj_keys = list(self._env.scene.rigid_objects.keys())
+            multi_obj_keys = [k for k in rigid_obj_keys if k.startswith("object_")]
+            if len(multi_obj_keys) > 0:
+                self._multi_object_mode = True
+                # Build object names list: (safe_name, original_name)
+                # safe_name has underscores, original_name has hyphens (for motion lookup)
+                for key in multi_obj_keys:
+                    safe_name = key[7:]  # Remove "object_" prefix
+                    original_name = safe_name.replace(  # noqa: F841
+                        "_", "-"
+                    )  # Convert back for motion lookup
+                    # But we need to be careful - only convert the ones that were originally hyphens
+                    # Actually we don't know which underscores were hyphens, so let's store just safe_name
+                    # and try both when looking up motions
+                    self._object_names.append(safe_name)
+                print(  # noqa: T201
+                    f"[Multi-Object Mode] Detected {len(self._object_names)} objects in scene"
+                )
+
+        # Cache for table metadata loaded from pkl files (loaded once per motion key)
+        self._table_meta_cache = {}  # motion_key -> {'table_pos': tensor, 'table_quat': tensor}
+
+        if self.cfg.use_height_map:
+            try:
+                from simple_raycaster.raycaster import MultiMeshRaycaster
+            except ImportError:
+                command_install = "pip install -e git+https://github.com/Agent-3154/simple-raycaster.git@197daa6dcb146c5ce3e675a173328e17df6b9777#egg=simple-raycaster"
+                raise ImportError(  # noqa: B904
+                    f"simple-raycaster is required for height map observation. Install with: {command_install}"
+                )
+            import omni
+
+            prim_path_patterns = [
+                # r"/World/ground",
+                r"/World/envs/env_\d+/Object",
+            ]
+            stage = omni.usd.get_context().get_stage()
+
+            self.height_map = MultiMeshRaycaster.from_prim_paths(
+                paths=prim_path_patterns,
+                stage=stage,
+                device=str(self.device),
+            )
+            mesh_filters = [[f"/World/envs/env_{i}/Object"] for i in range(self.num_envs)]
+            self.num_mesh_per_cam, self.mesh_ids, self.cam_ids = self.height_map.get_mesh_ids(
+                mesh_filters, device=self.device
+            )
+
+            lin = torch.linspace(
+                -0.5 * self.cfg.height_map_size,
+                0.5 * self.cfg.height_map_size,
+                int(self.cfg.height_map_size / self.cfg.height_map_resolution) + 1,
+                device=self.device,
+            )
+            grid_x, grid_y = torch.meshgrid(lin, lin)
+            zeros = torch.zeros_like(grid_x)
+            scan_offsets = torch.stack(
+                [grid_x, grid_y, zeros],
+                dim=-1,
+            ).view(-1, 3)
+            scan_offsets[:, 2] = -1.0
+            ray_dirs_local = torch.nn.functional.normalize(scan_offsets, dim=-1)
+            self.ray_dirs_local = ray_dirs_local.expand(self.num_envs, -1, -1)
+            self.num_rays = self.ray_dirs_local.shape[1]
+            self.num_rays_x, self.num_rays_y = grid_x.shape
+            self.scan_dot_pos_w = torch.zeros(
+                self.num_envs, self.num_rays_x, self.num_rays_y, 3, device=self.device
+            )
+
+    # =========================================================================
+    # Offline (kinematic-only) factory
+    # =========================================================================
+
+    @classmethod
+    def create_offline(
+        cls,
+        motion_lib_cfg,
+        device,
+    ):
+        """Create a MotionCommand without env/sim for offline kinematic use.
+
+        All @property methods that only use motion_lib data work unchanged.
+        Properties that require robot state (self.robot) will use the reference
+        motion root as the anchor (equivalent to use_ref_motion_root_quat_w_as_anchor=True).
+
+        Expected keys in motion_lib_cfg:
+            num_future_frames: Number of future reference frames.
+            dt_future_ref_frames: Time delta between future reference frames (seconds).
+        """
+        inst = object.__new__(cls)
+        inst._offline = True  # noqa: SLF001
+        inst._debug_vis_handle = None  # noqa: SLF001
+        num_envs = 1
+
+        # device / num_envs are read-only properties on ManagerTermBase that
+        # delegate to self._env, so we provide a lightweight stand-in.
+        class _MinimalEnv:
+            pass
+
+        _env = _MinimalEnv()
+        _env.device = device
+        _env.num_envs = num_envs
+        inst._env = _env  # noqa: SLF001
+
+        # Motion lib setup
+        motion_lib_cfg = (
+            easydict.EasyDict(motion_lib_cfg)
+            if not isinstance(motion_lib_cfg, easydict.EasyDict)
+            else motion_lib_cfg
+        )
+
+        num_future_frames = motion_lib_cfg.get("num_future_frames", 8)
+        dt_future_ref_frames = motion_lib_cfg.get("dt_future_ref_frames", 0.1)
+
+        # Inject body/DOF mapping into motion_lib_cfg so motion_lib handles
+        # body reordering and xyzw→wxyz quaternion conversion at load time.
+
+        isaaclab_to_mujoco_mapping = order_converter.G1Converter().get_isaaclab_to_mujoco_mapping()
+        motion_lib_cfg.update(
+            {
+                "mujoco_to_isaaclab_body": isaaclab_to_mujoco_mapping["mujoco_to_isaaclab_body"],
+                "mujoco_to_isaaclab_dof": isaaclab_to_mujoco_mapping["mujoco_to_isaaclab_dof"],
+                "isaaclab_to_mujoco_body": isaaclab_to_mujoco_mapping["isaaclab_to_mujoco_body"],
+                "isaaclab_to_mujoco_dof": isaaclab_to_mujoco_mapping["isaaclab_to_mujoco_dof"],
+            }
+        )
+        # body_indexes_data: same logic as __init__ — use body_names to select
+        # a subset when provided, otherwise default to all bodies.
+        isaaclab_joints = isaaclab_to_mujoco_mapping["isaaclab_joints"]
+        body_names = motion_lib_cfg.get("body_names", None)
+        if "body_indexes_data" not in motion_lib_cfg:
+            if body_names is not None:
+                motion_lib_cfg.body_indexes_data = [
+                    isaaclab_joints.index(name) for name in body_names
+                ]
+            else:
+                num_bodies_full = len(isaaclab_to_mujoco_mapping["mujoco_to_isaaclab_body"])
+                motion_lib_cfg.body_indexes_data = list(range(num_bodies_full))
+
+        inst.motion_lib = motion_lib_robot.MotionLibRobot(motion_lib_cfg, num_envs, device)
+        max_num_motions = motion_lib_cfg.get("max_num_motions", None)
+        inst.motion_lib.load_motions_for_training(max_num_seqs=max_num_motions)
+
+        # Timing
+        target_fps = motion_lib_cfg.get("target_fps", 50)
+        inst.num_future_frames = num_future_frames
+        inst.frame_skips = dt_future_ref_frames // (1.0 / target_fps)
+        steps_exact = dt_future_ref_frames * target_fps
+        if abs(steps_exact - round(steps_exact)) > 1e-9:
+            import warnings
+
+            warnings.warn(
+                f"dt_future_ref_frames={dt_future_ref_frames} * target_fps={target_fps} "
+                f"= {steps_exact} is not an integer; using frame_skips={inst.frame_skips}",
+                stacklevel=2,
+            )
+        inst.future_time_steps_init = (
+            (torch.arange(num_future_frames, device=device, dtype=torch.long) * inst.frame_skips)
+            .view(1, -1)
+            .repeat(num_envs, 1)
+        )
+        # Anchor body index (from the body_names list, matching __init__)
+        anchor_body = motion_lib_cfg.get("anchor_body", None)
+        if body_names is not None and anchor_body is not None:
+            inst.motion_anchor_body_index = body_names.index(anchor_body)
+        else:
+            inst.motion_anchor_body_index = 0  # Default to root
+        inst.num_bodies = len(motion_lib_cfg.body_indexes_data)
+
+        # DOF mapping
+        if isaaclab_to_mujoco_mapping is not None:
+            inst.isaaclab_to_mujoco_dof = isaaclab_to_mujoco_mapping["isaaclab_to_mujoco_dof"]
+            inst.mujoco_to_isaaclab_dof = isaaclab_to_mujoco_mapping["mujoco_to_isaaclab_dof"]
+        inst.lower_joint_indices_mujoco = list(range(12))
+        if hasattr(inst, "isaaclab_to_mujoco_dof"):
+            inst.lower_joint_isaaclab_indices = [
+                inst.isaaclab_to_mujoco_dof[i] for i in inst.lower_joint_indices_mujoco
+            ]
+
+        # Variable frame support (offline)
+        inst.variable_frames_enabled = motion_lib_cfg.get("variable_frames_enabled", False)
+        inst.per_env_num_frames, inst._frame_choices = _init_variable_frames(  # noqa: SLF001
+            inst.variable_frames_enabled,
+            motion_lib_cfg.get("variable_frames_min", 16),
+            num_future_frames,
+            motion_lib_cfg.get("variable_frames_step", 4),
+            num_envs,
+            device,
+        )
+
+        # Motion state (updated per-sample via set_motion_state)
+        inst.motion_ids = torch.zeros(num_envs, dtype=torch.long, device=device)
+        inst.time_steps = torch.zeros(num_envs, dtype=torch.long, device=device)
+        inst.motion_start_time_steps = torch.zeros(num_envs, dtype=torch.long, device=device)
+        inst.future_motion_ids = inst.motion_ids.repeat_interleave(num_future_frames)
+        inst.motion_num_steps = torch.zeros(num_envs, dtype=torch.long, device=device)
+
+        # Offline mode: use ref motion root as robot anchor
+        inst.use_ref_motion_root_quat_w_as_anchor = True
+        inst.ref_motion_root_rotation_noise = None
+
+        # SMPL future frames (may differ from regular future frames)
+        smpl_num_future_frames = motion_lib_cfg.get("smpl_num_future_frames", None)
+        smpl_dt_future_ref_frames = motion_lib_cfg.get("smpl_dt_future_ref_frames", None)
+        inst.smpl_num_future_frames = (
+            smpl_num_future_frames if smpl_num_future_frames is not None else num_future_frames
+        )
+        inst.smpl_dt_future_ref_frames = (
+            smpl_dt_future_ref_frames
+            if smpl_dt_future_ref_frames is not None
+            else dt_future_ref_frames
+        )
+        inst.smpl_frame_skips = inst.smpl_dt_future_ref_frames // (1.0 / target_fps)
+        smpl_steps_exact = inst.smpl_dt_future_ref_frames * target_fps
+        if abs(smpl_steps_exact - round(smpl_steps_exact)) > 1e-9:
+            import warnings
+
+            warnings.warn(
+                f"smpl_dt_future_ref_frames={inst.smpl_dt_future_ref_frames} * target_fps={target_fps} "
+                f"= {smpl_steps_exact} is not an integer; using smpl_frame_skips={inst.smpl_frame_skips}",
+                stacklevel=2,
+            )
+        inst.smpl_future_time_steps_init = (
+            (
+                torch.arange(inst.smpl_num_future_frames, device=device, dtype=torch.long)
+                * inst.smpl_frame_skips
+            )
+            .view(1, -1)
+            .repeat(num_envs, 1)
+        )
+        inst.smpl_future_motion_ids = inst.motion_ids.repeat_interleave(inst.smpl_num_future_frames)
+
+        # Encoder sampling (not used offline, but set for compat)
+        inst.encoder_sample_probs_dict = None
+        inst.encoder_sample_probs = None
+        inst.is_evaluating = False
+
+        # Metrics dict (normally set by CommandTerm.__init__)
+        inst.metrics = {}
+
+        return inst
+
+    def set_is_evaluating(self, is_evaluating: bool = True):
+        """Toggle evaluation mode, which disables reset randomizations."""
+        self.is_evaluating = is_evaluating
+
+    def forward_motion_samples(self, env_ids: Sequence[int]):
+        """Assign sequential motion IDs and reset time steps for given envs.
+
+        Used during paired/evaluation mode to deterministically cycle through
+        motions. Updates motion IDs, start times, motion lengths, and caches
+        the initial body pose for relative-frame computation.
+
+        Args:
+            env_ids: Environment indices to reassign motions for.
+        """
+        self.motion_ids[env_ids] = (
+            torch.arange(self.num_envs).to(self.device)
+            % self.motion_lib._num_motions  # noqa: SLF001
+        )[env_ids]
+        sampled_times = self.motion_lib.sample_time_steps(
+            self.motion_ids[env_ids], truncate_time=None
+        )
+        if self.cfg.sample_from_n_initial_frames is not None:
+            # Sample uniformly from first N frames
+            n_frames = self.cfg.sample_from_n_initial_frames
+            sampled_times = torch.randint(
+                0, n_frames, (len(env_ids),), dtype=sampled_times.dtype, device=self.device
+            )
+        elif self.cfg.start_from_first_frame:
+            sampled_times.zero_()
+        self.motion_start_time_steps[env_ids] = sampled_times
+        self.motion_num_steps[env_ids] = self.motion_lib.get_motion_num_steps(
+            self.motion_ids[env_ids]
+        )
+        self.time_steps[env_ids] = 0
+        self.body_pos_relative_w[env_ids] = self.motion_lib.get_body_pos_w(
+            self.motion_ids[env_ids], self.motion_start_time_steps[env_ids]
+        )
+        self.body_quat_relative_w[env_ids] = self.motion_lib.get_body_quat_w(
+            self.motion_ids[env_ids], self.motion_start_time_steps[env_ids]
+        )
+
+    @property
+    def command(self) -> torch.Tensor:  # TODO Consider again if this is the best observation  # noqa: TD002, TD003, TD004
+        """Return current-frame joint positions and velocities concatenated.
+
+        Returns:
+            Tensor of shape ``(num_envs, 2 * num_dof)``.
+        """
+        return torch.cat([self.joint_pos, self.joint_vel], dim=1)
+
+    @property
+    def command_z(self) -> torch.Tensor:
+        """Return reference root height (z) for the current frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, 1)``.
+        """
+        return self.root_z
+
+    @property
+    def command_z_multi_future(self) -> torch.Tensor:
+        """Return reference root height (z) for the first future frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, 1)``.
+        """
+        return self.root_z_multi_future
+
+    @property
+    def command_vel(self) -> torch.Tensor:  # TODO Consider again if this is the best observation  # noqa: TD002, TD003, TD004
+        """Return reference root velocity (2D linear + 1D angular) in body frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, 3)``.
+        """
+        return torch.cat([self.root_lin_vel_b_2d, self.root_ang_vel_b_1d], dim=1)
+
+    @property
+    def command_max(self) -> torch.Tensor:
+        """Return all body state (pos, quat, lin_vel, ang_vel) flattened.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_bodies * 13)``.
+        """
+        return torch.cat(
+            [self.body_pos_w, self.body_quat_w, self.body_lin_vel_w, self.body_ang_vel_w], dim=-1
+        ).view(self.num_envs, -1)
+
+    @property
+    def command_max(self) -> torch.Tensor:  # noqa: F811
+        """Return all body state (pos, quat, lin_vel, ang_vel) flattened.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_bodies * 13)``.
+        """
+        return torch.cat(
+            [self.body_pos_w, self.body_quat_w, self.body_lin_vel_w, self.body_ang_vel_w], dim=-1
+        ).view(self.num_envs, -1)
+
+    @property
+    def command_max_diff_l(self) -> torch.Tensor:
+        """Return body state differences in robot-local frame, flattened.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_bodies * 13)``.
+        """
+        return torch.cat(
+            [
+                self.body_pos_dif_l.view(self.num_envs, -1),
+                self.body_quat_dif_l.view(self.num_envs, -1),
+                self.body_lin_vel_l,
+                self.body_ang_vel_l,
+            ],
+            dim=-1,
+        ).view(self.num_envs, -1)
+
+    @property
+    def command_max_diff_l_multi_future(self) -> torch.Tensor:
+        """Return multi-future body state diffs and local poses concatenated.
+
+        Returns:
+            Tensor of shape ``(num_envs, <variable>)`` depending on body count and future frames.
+        """
+        # TODO: this is not done.  # noqa: TD002, TD003
+        return torch.cat(
+            [
+                self.body_pos_dif_l_multi_future.view(self.num_envs, -1),
+                self.body_quat_dif_l_multi_future.view(self.num_envs, -1),
+                self.body_lin_vel_dif_l_multi_future,
+                self.body_ang_vel_dif_l_multi_future,
+                self.bod_pos_local_multi_future,
+                self.body_quat_local_multi_future,
+                self.body_lin_vel_l_multi_future,
+                self.body_ang_vel_l_multi_future,
+            ],
+            dim=-1,
+        ).view(self.num_envs, -1)
+
+    @property
+    def command_max_multi_future(self) -> torch.Tensor:
+        """Return all body state for all future frames in world frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_future_frames * num_bodies * 13)``.
+        """
+        return torch.cat(
+            [
+                self.body_pos_w_multi_future,
+                self.body_quat_w_multi_future,
+                self.body_lin_vel_w_multi_future,
+                self.body_ang_vel_w_multi_future,
+            ],
+            dim=1,
+        )
+
+    @property
+    def command_multi_future(self) -> torch.Tensor:
+        """Return joint positions and velocities for all future frames, flattened.
+
+        Returns:
+            Tensor of shape ``(num_envs, 2 * num_future_frames * num_dof)``.
+        """
+        return torch.cat([self.joint_pos_multi_future, self.joint_vel_multi_future], dim=1)
+
+    @property
+    def command_multi_future_joint_pos(self) -> torch.Tensor:
+        """Return joint positions for all future frames, flattened.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_future_frames * num_dof)``.
+        """
+        return self.joint_pos_multi_future
+
+    @property
+    def command_multi_future_joint_body_pos(self) -> torch.Tensor:
+        """Return joint positions and anchor-relative body positions for all future frames.
+
+        Body positions are expressed relative to each frame's anchor position,
+        rotated into the anchor's yaw-inverse frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_future_frames * num_dof + num_future_frames * num_bodies * 3)``.
+        """
+        body_pos_multi_frame = self.motion_lib.get_body_pos_w(
+            self.future_motion_ids, self.future_time_steps
+        ).view(self.num_envs, self.num_future_frames, self.num_bodies, 3)
+        anchor_pos_multi_frame = (
+            self.motion_lib.get_body_pos_w(self.future_motion_ids, self.future_time_steps)[
+                :, self.motion_anchor_body_index
+            ]
+            .view(self.num_envs, self.num_future_frames, 1, 3)
+            .expand(self.num_envs, self.num_future_frames, self.num_bodies, 3)
+        )
+        body_pos_relative_to_anchor_multi_frame = body_pos_multi_frame - anchor_pos_multi_frame
+        anchor_quat_w_repeat = self.anchor_quat_w_multi_future.view(
+            self.num_envs, self.num_future_frames, 1, 4
+        ).expand(self.num_envs, self.num_future_frames, self.num_bodies, 4)
+        body_pos_multi_frame = quat_apply_yaw(
+            quat_inv(anchor_quat_w_repeat), body_pos_relative_to_anchor_multi_frame
+        ).reshape(self.num_envs, -1)
+        return torch.cat([self.joint_pos_multi_future, body_pos_multi_frame], dim=1)
+
+    @property
+    def command_multi_future_joint_body_abs_pos(self) -> torch.Tensor:
+        """Return joint positions and anchor-relative body positions (without rotation).
+
+        Unlike ``command_multi_future_joint_body_pos``, the body position deltas
+        are not rotated into the anchor's local frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_future_frames * num_dof + num_future_frames * num_bodies * 3)``.
+        """
+        body_pos_multi_frame = self.motion_lib.get_body_pos_w(
+            self.future_motion_ids, self.future_time_steps
+        ).reshape(self.num_envs, -1)
+        anchor_pos_multi_frame = (
+            self.motion_lib.get_body_pos_w(self.future_motion_ids, self.future_time_steps)[
+                :, self.motion_anchor_body_index
+            ]
+            .view(self.num_envs, self.num_future_frames, 1, 3)
+            .expand(self.num_envs, self.num_future_frames, self.num_bodies, 3)
+            .reshape(self.num_envs, -1)
+        )
+        body_pos_relative_to_anchor_multi_frame = body_pos_multi_frame - anchor_pos_multi_frame
+        return torch.cat(
+            [self.joint_pos_multi_future, body_pos_relative_to_anchor_multi_frame], dim=1
+        )
+
+    # @property
+    # def command_multi_future_joint_body_diff_pos(self) -> torch.Tensor:
+    #     body_pos_w = self.motion_lib.get_body_pos_w(self.future_motion_ids, self.future_time_steps).view(self.num_envs, self.num_future_frames, -1, 3)  # noqa: E501
+    #     body_pos_w_env = body_pos_w + self._env.scene.env_origins[:, None, None, :]
+    #     body_pos_dif = body_pos_w_env - self.robot_body_pos_w[:, None, :, :]
+    #     body_pos_relative_to_robot_anchor_multi_frame = quat_apply_yaw(quat_inv(self.robot_anchor_quat_w.repeat(1, self.num_future_frames, self.num_bodies, 1)), body_pos_dif).reshape(self.num_envs, -1)  # noqa: E501
+    #     return torch.cat([self.joint_pos_multi_future, body_pos_relative_to_robot_anchor_multi_frame], dim=1)
+
+    @property
+    def command_multi_future_joint_body_diff_pos(self) -> torch.Tensor:
+        """Return joint positions and robot-relative body position differences.
+
+        Computes body positions in robot-anchor-relative frame: positions are
+        translated so the XY comes from the robot root and Z from the reference
+        anchor, then rotated by the heading difference. The result is the
+        difference from the current robot body positions, expressed in the
+        robot's local frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_future_frames * num_dof + num_future_frames * num_bodies * 3)``.
+        """
+        N, F, B = self.num_envs, self.num_future_frames, self.num_bodies
+        anchor_pos_w_repeat = self.anchor_pos_w_multi_future.view(N, F, 1, 3).expand(N, F, B, 3)
+        anchor_quat_w_repeat = self.anchor_quat_w_multi_future.view(N, F, 1, 4).expand(N, F, B, 4)
+        robot_anchor_pos_w_repeat = self.robot_anchor_pos_w[:, None, None, :].expand(N, F, B, 3)
+        robot_anchor_quat_w_repeat = self.robot_anchor_quat_w[:, None, None, :].expand(N, F, B, 4)
+
+        delta_pos_w = robot_anchor_pos_w_repeat.clone()  # Root position of the robot
+        delta_pos_w[..., 2] = anchor_pos_w_repeat[..., 2]
+        delta_ori_w = torch_transform.get_heading_q(
+            quat_mul(robot_anchor_quat_w_repeat, quat_inv(anchor_quat_w_repeat))
+        )
+        body_pos_relative_w_multi_frame = delta_pos_w + quat_apply(
+            delta_ori_w, self.body_pos_w_multi_future.view(N, F, B, 3) - anchor_pos_w_repeat
+        )
+        body_pos_dif = body_pos_relative_w_multi_frame - self.robot_body_pos_w.view(
+            N, 1, B, 3
+        ).expand(N, F, B, 3)
+        body_pos_relative_to_robot_anchor_multi_frame = quat_apply_yaw(
+            quat_inv(robot_anchor_quat_w_repeat), body_pos_dif
+        ).reshape(N, -1)
+        return torch.cat(
+            [self.joint_pos_multi_future, body_pos_relative_to_robot_anchor_multi_frame], dim=1
+        )
+
+    @property
+    def command_multi_future_lower_body(self) -> torch.Tensor:
+        return torch.cat(
+            [self.joint_pos_lower_body_multi_future, self.joint_vel_lower_body_multi_future], dim=1
+        )
+
+    @property
+    def command_multi_future_lower_body_joint_pos(self) -> torch.Tensor:
+        return self.joint_pos_lower_body_multi_future
+
+    # =========================================================================
+    # Egocentric joint transforms (positions + rotations) for reference frames
+    # =========================================================================
+
+    @property
+    def num_bodies_full(self) -> int:
+        """Get the full number of bodies (all bodies, not just selected body_indexes)."""
+        return self.motion_lib.num_bodies_full
+
+    @property
+    def egocentric_joint_positions_multi_future(self) -> torch.Tensor:
+        """Get body positions (including pelvis) in egocentric frame for all future frames.
+
+        Egocentric frame = projected root frame (heading/yaw only rotation, z=0 projected).
+        For each future frame, positions are relative to that frame's projected root.
+
+        Returns:
+            torch.Tensor: [num_envs, num_future_frames, num_bodies_full, 3] joint positions
+                          in egocentric (projected root) frame
+        """
+        N, F, B = self.num_envs, self.num_future_frames, self.num_bodies_full
+
+        # Get full body positions in world frame for all future frames
+        body_pos_w = self.motion_lib.get_body_pos_w_full(
+            self.future_motion_ids, self.future_time_steps
+        ).view(N, F, B, 3)
+
+        # Get anchor (pelvis) positions for each future frame
+        anchor_body_idx_full = self.motion_lib.m_cfg.get("anchor_body_idx_full", 0)
+        anchor_pos_w = body_pos_w[:, :, anchor_body_idx_full, :]  # [N, F, 3]
+
+        # Project anchor to ground plane (z=0) for egocentric frame
+        anchor_pos_projected = anchor_pos_w.clone()
+        anchor_pos_projected[..., 2] = 0  # Set z to 0
+        anchor_pos_projected_expanded = anchor_pos_projected.view(N, F, 1, 3).expand(N, F, B, 3)
+
+        # Get anchor quaternions for each future frame
+        body_quat_w = self.motion_lib.get_body_quat_w_full(
+            self.future_motion_ids, self.future_time_steps
+        ).view(N, F, B, 4)
+        anchor_quat_w = body_quat_w[:, :, anchor_body_idx_full, :]  # [N, F, 4]
+
+        # Extract heading quaternion using get_heading_q (canonicalize)
+        anchor_heading_quat = torch_transform.get_heading_q(anchor_quat_w.reshape(-1, 4)).reshape(
+            N, F, 4
+        )
+        anchor_heading_quat_expanded = anchor_heading_quat.view(N, F, 1, 4).expand(N, F, B, 4)
+
+        # Compute egocentric positions:
+        # 1. Translate to projected anchor origin (z=0)
+        body_pos_relative = body_pos_w - anchor_pos_projected_expanded
+
+        # 2. Rotate by inverse of heading quaternion
+        body_pos_egocentric = quat_apply(
+            quat_inv(anchor_heading_quat_expanded.reshape(-1, 4)), body_pos_relative.reshape(-1, 3)
+        ).reshape(N, F, B, 3)
+
+        return body_pos_egocentric
+
+    @property
+    def egocentric_joint_rotations_multi_future(self) -> torch.Tensor:
+        """Get body rotations (including pelvis) in egocentric frame for all future frames.
+
+        Returns 6D rotation representation (first two columns of rotation matrix).
+        For each future frame, rotations are relative to that frame's projected root.
+
+        Returns:
+            torch.Tensor: [num_envs, num_future_frames, num_bodies_full, 6] joint rotations
+                          in 6D representation (first two columns of rotation matrix)
+        """
+        N, F, B = self.num_envs, self.num_future_frames, self.num_bodies_full
+
+        # Get full body quaternions in world frame for all future frames
+        body_quat_w = self.motion_lib.get_body_quat_w_full(
+            self.future_motion_ids, self.future_time_steps
+        ).view(N, F, B, 4)
+
+        # Get anchor quaternions for each future frame
+        anchor_body_idx_full = self.motion_lib.m_cfg.get(
+            "anchor_body_idx_full", 0
+        )  # pelvis is usually index 0
+        anchor_quat_w = body_quat_w[:, :, anchor_body_idx_full, :]  # [N, F, 4]
+
+        # Extract heading quaternion using get_heading_q (canonicalize)
+        anchor_heading_quat = torch_transform.get_heading_q(anchor_quat_w.reshape(-1, 4)).reshape(
+            N, F, 4
+        )
+        anchor_heading_quat_expanded = anchor_heading_quat.view(N, F, 1, 4).expand(N, F, B, 4)
+
+        # Compute relative rotation: q_relative = q_heading_inv * q_body
+        body_quat_egocentric = quat_mul(
+            quat_inv(anchor_heading_quat_expanded.reshape(-1, 4)), body_quat_w.reshape(-1, 4)
+        ).reshape(N, F, B, 4)
+
+        # Convert to 6D representation (first two columns of rotation matrix)
+        mat = matrix_from_quat(body_quat_egocentric)  # [N, F, B, 3, 3]
+        body_rot_6d = rotations.mat_to_rot6d_first_two_cols(mat).reshape(N, F, B, 6)  # [N, F, B, 6]
+
+        return body_rot_6d
+
+    @property
+    def egocentric_joint_transforms_multi_future(self) -> torch.Tensor:
+        """Combined egocentric joint transforms (positions + rotations) for all future frames.
+
+        Returns:
+            torch.Tensor: [num_envs, num_future_frames, num_bodies_full, 9]
+                          (3 for position + 6 for 6D rotation)
+        """
+        positions = self.egocentric_joint_positions_multi_future  # [N, F, B, 3]
+        rotations_ = self.egocentric_joint_rotations_multi_future  # [N, F, B, 6]
+        return torch.cat([positions, rotations_], dim=-1)
+
+    @property
+    def root_transforms_relative_to_first_frame(self) -> torch.Tensor:
+        """Get root transforms (position + rotation) relative to the first reference frame's
+        projected root (z=0, heading only).
+
+        Position: Delta from first frame's projected root position, rotated to first frame's heading frame.
+        Rotation: Relative rotation from first frame's heading quaternion (6D representation).
+
+        Returns:
+            torch.Tensor: [num_envs, num_future_frames, 9]
+                          (3 for relative position + 6 for 6D relative rotation)
+        """  # noqa: D205
+        N, F = self.num_envs, self.num_future_frames
+
+        # Get root positions for all future frames
+        root_pos_w = self.motion_lib.get_root_pos_w(
+            self.future_motion_ids, self.future_time_steps
+        ).view(N, F, 3)
+
+        # Get root quaternions for all future frames
+        root_quat_w = self.motion_lib.get_root_quat_w(
+            self.future_motion_ids, self.future_time_steps
+        ).view(N, F, 4)
+
+        # First frame as reference - use projected root (z=0)
+        first_frame_pos = root_pos_w[:, 0:1, :]  # [N, 1, 3]
+        first_frame_pos_projected = first_frame_pos.clone()
+        first_frame_pos_projected[..., 2] = 0  # Project to ground plane
+
+        # Extract heading quaternion using get_heading_q (canonicalize)
+        first_frame_quat = root_quat_w[:, 0:1, :]  # [N, 1, 4]
+        first_frame_heading = torch_transform.get_heading_q(
+            first_frame_quat.reshape(-1, 4)
+        ).reshape(N, 1, 4)
+
+        # Relative position: delta from projected first frame, rotated to heading frame
+        delta_pos_w = root_pos_w - first_frame_pos_projected.expand(N, F, 3)
+
+        # Rotate delta position by inverse of first frame's heading
+        delta_pos_local = quat_apply(
+            quat_inv(first_frame_heading.expand(N, F, 4).reshape(-1, 4)), delta_pos_w.reshape(-1, 3)
+        ).reshape(N, F, 3)
+
+        # Relative rotation: q_relative = q_heading_inv * q_current
+        relative_quat = quat_mul(
+            quat_inv(first_frame_heading.expand(N, F, 4).reshape(-1, 4)), root_quat_w.reshape(-1, 4)
+        ).reshape(N, F, 4)
+
+        # Convert to 6D representation (first two columns of rotation matrix)
+        mat = matrix_from_quat(relative_quat)  # [N, F, 3, 3]
+        relative_rot_6d = rotations.mat_to_rot6d_first_two_cols(mat).reshape(N, F, 6)  # [N, F, 6]
+
+        return torch.cat([delta_pos_local, relative_rot_6d], dim=-1)
+
+    @property
+    def smpl_joints(self) -> torch.Tensor:
+        """Return SMPL joint positions for the current frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_smpl_joints, 3)``.
+        """
+        return self.motion_lib.get_smpl_joints(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+
+    @property
+    def smpl_joints_multi_future(self) -> torch.Tensor:
+        """Return SMPL joint positions for all SMPL future frames.
+
+        Returns:
+            Tensor of shape ``(num_envs, smpl_num_future_frames, num_smpl_joints, 3)``.
+        """
+        smpl_joints_mf = self.motion_lib.get_smpl_joints(
+            self.smpl_future_motion_ids, self.smpl_future_time_steps
+        )
+        return smpl_joints_mf.view(
+            self.num_envs, self.smpl_num_future_frames, *smpl_joints_mf.shape[1:]
+        )
+
+    @property
+    def smpl_transl_multi_future(self) -> torch.Tensor:
+        smpl_transl_mf = self.motion_lib.get_smpl_transl(
+            self.smpl_future_motion_ids, self.smpl_future_time_steps
+        )
+        return smpl_transl_mf.view(
+            self.num_envs, self.smpl_num_future_frames, *smpl_transl_mf.shape[1:]
+        )
+
+    @property
+    def smpl_transl_z_multi_future(self) -> torch.Tensor:
+        smpl_transl_mf = self.motion_lib.get_smpl_transl(
+            self.smpl_future_motion_ids, self.smpl_future_time_steps
+        )
+        # Only return y dimension (height) --- y is the z for smpl
+        return smpl_transl_mf[..., 1:2].view(self.num_envs, self.smpl_num_future_frames, 1)
+
+    @property
+    def smpl_pose(self) -> torch.Tensor:
+        """Return full SMPL pose (root + body joints) in axis-angle for current frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, 72)`` (24 joints * 3 axis-angle).
+        """
+        return self.motion_lib.get_smpl_pose(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+
+    @property
+    def smpl_body_pose(self) -> torch.Tensor:
+        """Return SMPL body pose (excluding root) in axis-angle for current frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, 69)`` (23 body joints * 3 axis-angle).
+        """
+        return self.motion_lib.get_smpl_pose(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )[..., 3:]
+
+    @property
+    def smpl_body_pose_6d(self) -> torch.Tensor:
+        """Return SMPL body pose in 6D rotation representation for current frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, 23 * 6)``.
+        """
+        smpl_body_pose = self.smpl_body_pose.reshape(-1, 23, 3)
+        quat = torch_transform.angle_axis_to_quaternion(smpl_body_pose)
+        mat = matrix_from_quat(quat)
+        smpl_body_pose_6d = mat[..., :2].reshape(smpl_body_pose.shape[0], -1)
+        return smpl_body_pose_6d
+
+    def smpl_root_ytoz_up(self, root_quat_y_up: torch.Tensor) -> torch.Tensor:
+        """Convert SMPL root quaternion from Y-up to Z-up convention.
+
+        Args:
+            root_quat_y_up: Root quaternions in Y-up frame, ``(N, 4)``.
+
+        Returns:
+            Root quaternions rotated to Z-up frame, ``(N, 4)``.
+        """
+        base_rot = torch_transform.angle_axis_to_quaternion(
+            torch.tensor([[np.pi / 2, 0.0, 0.0]]).to(root_quat_y_up)
+        )
+        root_quat_z_up = rotations.quat_mul(
+            base_rot.repeat(root_quat_y_up.shape[0], 1), root_quat_y_up, w_last=False
+        )
+        return root_quat_z_up
+
+    @property
+    def smpl_pose_noheading(self) -> torch.Tensor:
+        """Return SMPL pose with root heading removed (preserving pitch/roll).
+
+        Converts SMPL root from Y-up if needed, removes the base rotation offset,
+        extracts and removes heading, then re-encodes as axis-angle.
+
+        Returns:
+            Tensor of shape ``(num_envs, 72)``.
+        """
+        smpl_pose = self.motion_lib.get_smpl_pose(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+        root_quat = torch_transform.angle_axis_to_quaternion(smpl_pose[..., :3]).view(-1, 4)
+        if self.motion_lib.smpl_y_up:
+            root_quat = self.smpl_root_ytoz_up(root_quat)
+        root_quat = rotations.remove_smpl_base_rot(root_quat, w_last=False)
+        root_heading_inv = rotations.calc_heading_quat_inv(root_quat, w_last=False)
+        root_quat_noheading = rotations.quat_mul(root_heading_inv, root_quat, w_last=False)
+        smpl_pose[..., :3] = torch_transform.quaternion_to_angle_axis(root_quat_noheading).view(
+            smpl_pose.shape[:-1] + (3,)
+        )
+        return smpl_pose
+
+    @property
+    def smpl_root_quat_w(self) -> torch.Tensor:
+        """Return SMPL root quaternion in Z-up world frame for current frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, 4)``.
+        """
+        smpl_pose = self.motion_lib.get_smpl_pose(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+        root_quat = torch_transform.angle_axis_to_quaternion(smpl_pose[..., :3]).view(-1, 4)
+        if self.motion_lib.smpl_y_up:
+            root_quat = self.smpl_root_ytoz_up(root_quat)
+        root_quat = rotations.remove_smpl_base_rot(root_quat, w_last=False)
+        return root_quat
+
+    @property
+    def smpl_root_quat_w_multi_future(self) -> torch.Tensor:
+        """Return SMPL root quaternions in Z-up world frame for all SMPL future frames.
+
+        Returns:
+            Tensor of shape ``(num_envs, smpl_num_future_frames, 4)``.
+        """
+        smpl_pose = self.motion_lib.get_smpl_pose(
+            self.smpl_future_motion_ids, self.smpl_future_time_steps
+        )
+        root_quat = torch_transform.angle_axis_to_quaternion(smpl_pose[..., :3]).view(-1, 4)
+        if self.motion_lib.smpl_y_up:
+            root_quat = self.smpl_root_ytoz_up(root_quat)
+        root_quat = rotations.remove_smpl_base_rot(root_quat, w_last=False).view(
+            self.num_envs, self.smpl_num_future_frames, 4
+        )
+        return root_quat
+
+    @property
+    def smpl_root_quat_w_dif_l_multi_future(self) -> torch.Tensor:
+        """Return SMPL root orientation relative to robot orientation for all SMPL future frames.
+
+        Computes ``quat_inv(robot_anchor) * smpl_root`` and returns 6D rotation
+        matrix representation (first 2 columns).
+
+        Returns:
+            Tensor of shape ``(num_envs, smpl_num_future_frames * 6)``.
+        """
+        smpl_pose = self.motion_lib.get_smpl_pose(
+            self.smpl_future_motion_ids, self.smpl_future_time_steps
+        )
+        root_quat = torch_transform.angle_axis_to_quaternion(smpl_pose[..., :3]).view(-1, 4)
+        if self.motion_lib.smpl_y_up:
+            root_quat = self.smpl_root_ytoz_up(root_quat)
+        root_quat = rotations.remove_smpl_base_rot(root_quat, w_last=False)
+        root_rot_dif = quat_mul(
+            quat_inv(
+                self.robot_anchor_quat_w.view(self.num_envs, 1, 4).repeat(
+                    1, self.smpl_num_future_frames, 1
+                )
+            ),
+            root_quat.view(self.num_envs, self.smpl_num_future_frames, 4),
+        )
+        mat = matrix_from_quat(root_rot_dif)
+        root_rot_dif_l_mat = mat[..., :2].reshape(mat.shape[0], -1)
+        return root_rot_dif_l_mat
+
+    @property
+    def smpl_root_quat_w_dif_refheading_multi_future(self) -> torch.Tensor:
+        """SMPL root orientation canonicalized by the first SMPL future frame's heading.
+
+        Same extraction as smpl_root_quat_w_dif_l_multi_future but uses the heading
+        of the first SMPL future frame instead of the robot's orientation.
+
+        Returns:
+            torch.Tensor: 6D rotation matrix representation,
+                shape (num_envs, smpl_num_future_frames * 6)
+        """
+        smpl_pose = self.motion_lib.get_smpl_pose(
+            self.smpl_future_motion_ids, self.smpl_future_time_steps
+        )
+        smpl_root_quat = torch_transform.angle_axis_to_quaternion(smpl_pose[..., :3]).view(-1, 4)
+        if self.motion_lib.smpl_y_up:
+            smpl_root_quat = self.smpl_root_ytoz_up(smpl_root_quat)
+        smpl_root_quat = rotations.remove_smpl_base_rot(smpl_root_quat, w_last=False)
+        smpl_root_quat = smpl_root_quat.view(self.num_envs, self.smpl_num_future_frames, 4)
+        ref_first_heading = torch_transform.get_heading_q(smpl_root_quat[:, 0, :])
+        root_rot_dif = quat_mul(
+            quat_inv(
+                ref_first_heading.view(self.num_envs, 1, 4).expand(
+                    -1, self.smpl_num_future_frames, -1
+                )
+            ),
+            smpl_root_quat,
+        )
+        mat = matrix_from_quat(root_rot_dif)
+        return mat[..., :2].reshape(mat.shape[0], -1)
+
+    @property
+    def smpl_root_quat_w_dif_heading_multi_future(self) -> torch.Tensor:
+        """SMPL root orientation canonicalized by robot heading (yaw) only.
+
+        Same extraction as smpl_root_quat_w_dif_l_multi_future but uses the
+        robot's heading (yaw) instead of full orientation for canonicalization.
+
+        Returns:
+            torch.Tensor: 6D rotation matrix representation,
+                shape (num_envs, smpl_num_future_frames * 6)
+        """
+        smpl_pose = self.motion_lib.get_smpl_pose(
+            self.smpl_future_motion_ids, self.smpl_future_time_steps
+        )
+        root_quat = torch_transform.angle_axis_to_quaternion(smpl_pose[..., :3]).view(-1, 4)
+        if self.motion_lib.smpl_y_up:
+            root_quat = self.smpl_root_ytoz_up(root_quat)
+        root_quat = rotations.remove_smpl_base_rot(root_quat, w_last=False)
+        root_rot_dif = quat_mul(
+            quat_inv(
+                self.anchor_heading_quat.view(self.num_envs, 1, 4).expand(
+                    -1, self.smpl_num_future_frames, -1
+                )
+            ),
+            root_quat.view(self.num_envs, self.smpl_num_future_frames, 4),
+        )
+        mat = matrix_from_quat(root_rot_dif)
+        return mat[..., :2].reshape(mat.shape[0], -1)
+
+    @property
+    def smpl_pose_multi_future(self) -> torch.Tensor:
+        return self.motion_lib.get_smpl_pose(
+            self.smpl_future_motion_ids, self.smpl_future_time_steps
+        ).reshape(self.num_envs, self.smpl_num_future_frames, -1)
+
+    @property
+    def smpl_body_pose_multi_future(self) -> torch.Tensor:
+        return self.motion_lib.get_smpl_pose(
+            self.smpl_future_motion_ids, self.smpl_future_time_steps
+        )[..., 3:].reshape(self.num_envs, -1)
+
+    @property
+    def smpl_body_pose_multi_future_6d(self) -> torch.Tensor:
+        smpl_body_pose = self.smpl_body_pose_multi_future.reshape(
+            -1, self.smpl_num_future_frames, 23, 3
+        )
+        quat = torch_transform.angle_axis_to_quaternion(smpl_body_pose)
+        mat = matrix_from_quat(quat)
+        smpl_body_pose_6d = mat[..., :2].reshape(smpl_body_pose.shape[0], -1)
+        return smpl_body_pose_6d
+
+    # --- SOMA skeleton properties ---
+
+    @property
+    def soma_joints(self) -> torch.Tensor:
+        """SOMA joints in Z-up body-local frame. Stored Z-up in PKL (same as SMPL)."""
+        return self.motion_lib.get_soma_joints(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+
+    @property
+    def soma_joints_multi_future(self) -> torch.Tensor:
+        """SOMA joints multi-future in Z-up body-local frame."""
+        soma_joints_mf = self.motion_lib.get_soma_joints(
+            self.smpl_future_motion_ids, self.smpl_future_time_steps
+        )
+        return soma_joints_mf.view(
+            self.num_envs, self.smpl_num_future_frames, *soma_joints_mf.shape[1:]
+        )
+
+    @property
+    def soma_root_quat_w(self) -> torch.Tensor:
+        """SOMA root quaternion in Z-up world frame (wxyz)."""
+        root_quat = self.motion_lib.get_soma_root_quat(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+        if self.motion_lib.soma_y_up:
+            root_quat = self.smpl_root_ytoz_up(root_quat)
+        root_quat = rotations.remove_bvh_base_rot(root_quat, w_last=False)
+        return root_quat
+
+    @property
+    def soma_root_quat_w_multi_future(self) -> torch.Tensor:
+        """SOMA root quaternion multi-future in Z-up world frame (wxyz)."""
+        root_quat = self.motion_lib.get_soma_root_quat(
+            self.smpl_future_motion_ids, self.smpl_future_time_steps
+        )
+        if self.motion_lib.soma_y_up:
+            root_quat = self.smpl_root_ytoz_up(root_quat)
+        root_quat = rotations.remove_bvh_base_rot(root_quat, w_last=False)
+        return root_quat.view(self.num_envs, self.smpl_num_future_frames, 4)
+
+    @property
+    def soma_root_quat_w_dif_l_multi_future(self) -> torch.Tensor:
+        """SOMA root orientation relative to robot anchor, as 6D rotation matrix."""
+        root_quat = self.motion_lib.get_soma_root_quat(
+            self.smpl_future_motion_ids, self.smpl_future_time_steps
+        )
+        if self.motion_lib.soma_y_up:
+            root_quat = self.smpl_root_ytoz_up(root_quat)
+        root_quat = rotations.remove_bvh_base_rot(root_quat, w_last=False)
+        root_rot_dif = quat_mul(
+            quat_inv(
+                self.robot_anchor_quat_w.view(self.num_envs, 1, 4).repeat(
+                    1, self.smpl_num_future_frames, 1
+                )
+            ),
+            root_quat.view(self.num_envs, self.smpl_num_future_frames, 4),
+        )
+        mat = matrix_from_quat(root_rot_dif)
+        root_rot_dif_l_mat = mat[..., :2].reshape(mat.shape[0], -1)
+        return root_rot_dif_l_mat
+
+    @property
+    def soma_transl_multi_future(self) -> torch.Tensor:
+        """SOMA hips translation multi-future. Stored Y-up (same as SMPL transl)."""
+        soma_transl_mf = self.motion_lib.get_soma_transl(
+            self.smpl_future_motion_ids, self.smpl_future_time_steps
+        )
+        return soma_transl_mf.view(
+            self.num_envs, self.smpl_num_future_frames, *soma_transl_mf.shape[1:]
+        )
+
+    @property
+    def object_root_pos(self) -> torch.Tensor:
+        """Return object root position in world frame with env origin and z-offset.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_objects, 3)``.
+        """
+        object_root_pos = self.motion_lib.get_object_root_pos(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+        # Apply z-offset if configured (e.g., to lower chair into ground)
+        z_offset = getattr(self.cfg, "object_z_offset", 0.0)
+        if z_offset != 0.0:
+            object_root_pos = object_root_pos.clone()
+            object_root_pos[..., 2] += z_offset
+        return object_root_pos + self._env.scene.env_origins[:, None, :]
+
+    @property
+    def object_root_quat(self) -> torch.Tensor:
+        return self.motion_lib.get_object_root_quat(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+
+    def _get_object_pos_with_offset(self, env_ids: torch.Tensor) -> torch.Tensor:
+        """Get object position with randomization offset applied.
+
+        This method also resamples the offset for the given env_ids.
+        Called at environment reset to place object with random offset.
+
+        Args:
+            env_ids: Environment indices being reset
+
+        Returns:
+            Object positions with offset applied, shape [len(env_ids), 3]
+        """
+        obj_pos = self.object_root_pos[env_ids, 0].clone()
+
+        if self.cfg.object_position_randomize:
+            # Resample offset for these environments
+            rand_cfg = self.cfg.object_position_randomization or {}
+            x_range = rand_cfg.get("x", 0.0)
+            y_range = rand_cfg.get("y", 0.0)
+            z_range = rand_cfg.get("z", 0.0)
+
+            # Generate new random offsets
+            self._object_position_offset[env_ids, 0] = (
+                torch.rand(len(env_ids), device=self.device) * 2 - 1
+            ) * x_range
+            self._object_position_offset[env_ids, 1] = (
+                torch.rand(len(env_ids), device=self.device) * 2 - 1
+            ) * y_range
+            self._object_position_offset[env_ids, 2] = (
+                torch.rand(len(env_ids), device=self.device) * 2 - 1
+            ) * z_range
+
+            # Apply offset
+            obj_pos = obj_pos + self._object_position_offset[env_ids]
+
+        return obj_pos
+
+    @property
+    def object_root_pos_multi_future(self) -> torch.Tensor:
+        """Return object root positions for all future frames in world frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_future_frames, num_objects, 3)``.
+        """
+        object_root_pos = self.motion_lib.get_object_root_pos(
+            self.future_motion_ids, self.future_time_steps
+        )
+        object_root_pos_view = object_root_pos.view(self.num_envs, self.num_future_frames, -1, 3)
+        # Apply z-offset if configured (e.g., to lower chair into ground)
+        z_offset = getattr(self.cfg, "object_z_offset", 0.0)
+        if z_offset != 0.0:
+            object_root_pos_view = object_root_pos_view.clone()
+            object_root_pos_view[..., 2] += z_offset
+        return object_root_pos_view + self._env.scene.env_origins[:, None, None, :]
+
+    @property
+    def object_root_quat_multi_future(self) -> torch.Tensor:
+        object_root_quat = self.motion_lib.get_object_root_quat(
+            self.future_motion_ids, self.future_time_steps
+        )
+        return object_root_quat.view(self.num_envs, self.num_future_frames, -1, 4)
+
+    def _get_contact_center_world(self, hand: str) -> torch.Tensor | None:
+        """Get object contact center in world frame for the given hand.
+
+        Args:
+            hand: "left_hand" or "right_hand"
+
+        Returns:
+            Tensor of shape (num_envs, 3) in world frame, or None if not available.
+        """
+        contact_center = self.motion_lib.get_object_contact_center(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps, hand=hand
+        )
+        if contact_center is None:
+            return None
+
+        # Transform from object-local to world frame
+        obj_pos = self.object_root_pos[:, 0, :]  # (num_envs, 3)
+        obj_quat = self.object_root_quat[:, 0, :]  # (num_envs, 4)
+
+        from gear_sonic.isaac_utils import rotations
+
+        rotated_center = rotations.quat_rotate(obj_quat, contact_center, w_last=False)
+        world_center = rotated_center + obj_pos
+        return world_center
+
+    @property
+    def object_contact_center_left(self) -> torch.Tensor | None:
+        """Get left hand object contact center in world frame. Shape: (num_envs, 3)."""
+        return self._get_contact_center_world("left_hand")
+
+    @property
+    def object_contact_center_right(self) -> torch.Tensor | None:
+        """Get right hand object contact center in world frame. Shape: (num_envs, 3)."""
+        return self._get_contact_center_world("right_hand")
+
+    def get_in_contact(self, hand: str = "right_hand") -> torch.Tensor | None:
+        """Get binary in_contact label for the given hand at current timestep.
+
+        Args:
+            hand: "left_hand" or "right_hand"
+
+        Returns:
+            Tensor of shape (num_envs,) with 1.0 if in contact, 0.0 otherwise,
+            or None if not available.
+        """
+        return self.motion_lib.get_object_in_contact(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps, hand=hand
+        )
+
+    def get_hand_action(self, hand: str = "right_hand") -> torch.Tensor | None:
+        """Get discrete hand action (open/closed) for the given hand at current timestep.
+
+        Args:
+            hand: "left_hand" or "right_hand"
+
+        Returns:
+            Tensor of shape (num_envs,) with -1.0 = open, +1.0 = closed,
+            or None if not available.
+        """
+        return self.motion_lib.get_hand_action(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps, hand=hand
+        )
+
+    @property
+    def joint_pos(self) -> torch.Tensor:
+        """Return reference joint positions for the current frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_dof)``.
+        """
+        return self.motion_lib.get_dof_pos(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+
+    @property
+    def joint_pos_multi_future(self) -> torch.Tensor:
+        """Return reference joint positions for all future frames, flattened.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_future_frames * num_dof)``.
+        """
+        return self.motion_lib.get_dof_pos(self.future_motion_ids, self.future_time_steps).view(
+            self.num_envs, -1
+        )
+
+    @property
+    def joint_pos_multi_future_for_smpl(self) -> torch.Tensor:
+        return self.motion_lib.get_dof_pos(
+            self.smpl_future_motion_ids, self.smpl_future_time_steps
+        ).view(self.num_envs, -1)
+
+    @property
+    def joint_pos_lower_body_multi_future(self) -> torch.Tensor:
+        return self.motion_lib.get_dof_pos(self.future_motion_ids, self.future_time_steps)[
+            ..., self.lower_joint_isaaclab_indices
+        ].view(self.num_envs, -1)
+
+    @property
+    def joint_vel(self) -> torch.Tensor:
+        """Return reference joint velocities for the current frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_dof)``.
+        """
+        return self.motion_lib.get_dof_vel(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+
+    @property
+    def joint_vel_multi_future(self) -> torch.Tensor:
+        """Return reference joint velocities for all future frames, flattened.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_future_frames * num_dof)``.
+        """
+        return self.motion_lib.get_dof_vel(self.future_motion_ids, self.future_time_steps).view(
+            self.num_envs, -1
+        )
+
+    @property
+    def root_pos_multi_future(self) -> torch.Tensor:
+        return self.motion_lib.get_root_pos_w(self.future_motion_ids, self.future_time_steps).view(
+            self.num_envs, -1
+        )
+
+    @property
+    def root_quat_multi_future(self) -> torch.Tensor:
+        return self.motion_lib.get_root_quat_w(self.future_motion_ids, self.future_time_steps).view(
+            self.num_envs, -1
+        )
+
+    @property
+    def root_z(self) -> torch.Tensor:
+        return self.motion_lib.get_root_pos_w(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )[:, 2:3].view(self.num_envs, -1)
+
+    @property
+    def root_z_multi_future(self) -> torch.Tensor:
+        return self.motion_lib.get_root_pos_w(self.future_motion_ids, self.future_time_steps)[
+            :, 2:3
+        ].view(self.num_envs, -1)
+
+    @property
+    def joint_vel_lower_body_multi_future(self) -> torch.Tensor:
+        return self.motion_lib.get_dof_vel(self.future_motion_ids, self.future_time_steps)[
+            ..., self.lower_joint_isaaclab_indices
+        ].view(self.num_envs, -1)
+
+    @property
+    def body_pos_w(self) -> torch.Tensor:
+        """Return reference body positions in world frame for the current frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_bodies, 3)``.
+        """
+        return (
+            self.motion_lib.get_body_pos_w(
+                self.motion_ids, self.motion_start_time_steps + self.time_steps
+            )
+            + self._env.scene.env_origins[:, None, :]
+        )
+
+    @property
+    def body_pos_w_multi_future(self) -> torch.Tensor:
+        """Return reference body positions in world frame for all future frames.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_future_frames * num_bodies * 3)``.
+        """
+        body_pos_w = self.motion_lib.get_body_pos_w(
+            self.future_motion_ids, self.future_time_steps
+        ).view(self.num_envs, self.num_future_frames, -1, 3)
+        body_pos_w_env = body_pos_w + self._env.scene.env_origins[:, None, None, :]
+        return body_pos_w_env.reshape(self.num_envs, -1)
+
+    @property
+    def body_pos_dif_w(self) -> torch.Tensor:
+        """Return position difference (reference - robot) in world frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_bodies, 3)``.
+        """
+        body_pos_w = self.motion_lib.get_body_pos_w(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+        body_pos_w_env = body_pos_w + self._env.scene.env_origins[:, None, :]
+        body_pos_dif = body_pos_w_env - self.robot_body_pos_w
+        return body_pos_dif
+
+    @property
+    def body_pos_dif_w_multi_future(self) -> torch.Tensor:
+        body_pos_w = self.motion_lib.get_body_pos_w(
+            self.future_motion_ids, self.future_time_steps
+        ).view(self.num_envs, self.num_future_frames, -1, 3)
+        body_pos_w_env = body_pos_w + self._env.scene.env_origins[:, None, None, :]
+        body_pos_dif = body_pos_w_env - self.robot_body_pos_w[:, None, :, :]
+        return body_pos_dif.reshape(self.num_envs, -1)
+
+    @property
+    def body_pos_dif_l(self) -> torch.Tensor:
+        """Return body position difference de-headed into robot-local frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_bodies, 3)``.
+        """
+        body_pos_dif_w = self.body_pos_dif_w
+        root_quat = self.robot_anchor_quat_w.view(self.num_envs, 1, 4).repeat(1, self.num_bodies, 1)
+        deheaded_dif_l = quat_apply_yaw(quat_inv(root_quat), body_pos_dif_w)
+        return deheaded_dif_l
+
+    @property
+    def body_pos_dif_l_multi_future(self) -> torch.Tensor:
+        body_pos_dif_w = self.body_pos_dif_w_multi_future
+        root_quat = self.robot_anchor_quat_w.view(self.num_envs, 1, 1, 4).repeat(
+            1, self.num_future_frames, self.num_bodies, 1
+        )
+        deheaded_dif_l = quat_apply_yaw(quat_inv(root_quat), body_pos_dif_w)
+        return deheaded_dif_l
+
+    @property
+    def root_lin_vel_b_2d(self) -> torch.Tensor:
+        """Return reference root linear velocity (XY only) in body frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, 2)``.
+        """
+        root_lin_vel_w = self.motion_lib.get_root_lin_vel_w(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+        root_quat = self.anchor_quat_w.view(self.num_envs, 1, 4)
+        root_lin_vel_l = quat_apply_yaw(quat_inv(root_quat), root_lin_vel_w)[:, :2]
+        return root_lin_vel_l
+
+    @property
+    def root_ang_vel_b_1d(self) -> torch.Tensor:
+        """Return reference root angular velocity (yaw only) in body frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, 1)``.
+        """
+        root_ang_vel_w = self.motion_lib.get_root_ang_vel_w(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+        root_quat = self.anchor_quat_w.view(self.num_envs, 1, 4)
+        root_ang_vel_l = quat_apply_yaw(quat_inv(root_quat), root_ang_vel_w)[:, 2:3]
+        return root_ang_vel_l
+
+    @property
+    def body_quat_w(self) -> torch.Tensor:
+        """Return reference body quaternions in world frame for the current frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_bodies, 4)``.
+        """
+        return self.motion_lib.get_body_quat_w(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+
+    @property
+    def body_quat_w_multi_future(self) -> torch.Tensor:
+        return self.motion_lib.get_body_quat_w(self.future_motion_ids, self.future_time_steps).view(
+            self.num_envs, -1
+        )
+
+    @property
+    def body_quat_dif_w(self) -> torch.Tensor:
+        body_quat_w = self.motion_lib.get_body_quat_w(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+        body_quat_dif = quat_mul(quat_inv(body_quat_w), self.robot_body_quat_w)
+        return body_quat_dif
+
+    @property
+    def body_quat_dif_w_multi_future(self) -> torch.Tensor:
+        ref_body_quat_w = self.motion_lib.get_body_quat_w(
+            self.future_motion_ids, self.future_time_steps
+        ).view(self.num_envs, self.num_future_frames, -1, 4)
+        robot_body_quat_w = self.robot_body_quat_w.view(
+            self.num_envs, 1, self.num_bodies, 4
+        ).repeat(1, self.num_future_frames, 1, 1)
+        body_quat_dif = quat_mul(quat_inv(ref_body_quat_w), robot_body_quat_w)
+        return body_quat_dif
+
+    @property
+    def anchor_heading_quat(self) -> torch.Tensor:
+        """Return robot anchor heading quaternion (yaw-only, pitch/roll removed).
+
+        Returns:
+            Tensor of shape ``(num_envs, 4)``.
+        """
+        return torch_transform.get_heading_q(self.robot_anchor_quat_w)
+
+    # @property
+    # def body_quat_dif_l(self) -> torch.Tensor:
+    #     body_quat_dif_w = self.body_quat_dif_w
+    #     root_heading = self.anchor_heading_quat.view(self.num_envs, 1, 4).repeat(1, self.num_bodies, 1)
+    #     root_heading_inv = quat_inv(self.anchor_heading_quat).view(self.num_envs, 1, 4).repeat(1, self.num_bodies, 1)  # noqa: E501
+    #     deheaded_dif_l = quat_mul(quat_mul(root_heading_inv, body_quat_dif_w), root_heading)
+    #     mat = matrix_from_quat(deheaded_dif_l)
+    #     deheaded_dif_l_mat = mat[..., :2].reshape(mat.shape[0], -1)
+    #     return deheaded_dif_l_mat
+
+    # @property
+    # def body_quat_dif_l_multi_future(self) -> torch.Tensor:
+    #     body_quat_dif_w = self.body_quat_dif_w_multi_future
+    #     root_heading = self.anchor_heading_quat.view(self.num_envs, 1, 1, 4).repeat(1, self.num_future_frames, self.num_bodies, 1)  # noqa: E501
+    #     root_heading_inv = quat_inv(self.anchor_heading_quat).view(self.num_envs, 1, 1, 4).repeat(1, self.num_future_frames, self.num_bodies, 1)  # noqa: E501
+    #     deheaded_dif_l = quat_mul(quat_mul(root_heading_inv, body_quat_dif_w), root_heading)
+    #     mat = matrix_from_quat(deheaded_dif_l)
+    #     deheaded_dif_l_mat = mat[..., :2].reshape(mat.shape[0], -1)
+    #     return deheaded_dif_l_mat
+
+    @property
+    def root_rot_dif_l(self) -> torch.Tensor:
+        """Return reference root orientation relative to robot orientation in 6D repr.
+
+        Returns:
+            Tensor of shape ``(num_envs, 6)``.
+        """
+        ref_root_quat = self.motion_lib.get_root_quat_w(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+        root_rot_dif_w = quat_mul(quat_inv(self.robot_anchor_quat_w), ref_root_quat)
+        # root_heading = self.anchor_heading_quat.view(self.num_envs, 1, 4).repeat(1, 1, 1)
+        # root_heading_inv = quat_inv(self.anchor_heading_quat).view(self.num_envs, 1, 4).repeat(1, 1, 1)
+        # deheaded_rot_dif_l = quat_mul(quat_mul(root_heading_inv, root_rot_dif_w), root_heading)
+        mat = matrix_from_quat(root_rot_dif_w)
+        deheaded_rot_dif_l_mat = mat[..., :2].reshape(mat.shape[0], -1)
+        return deheaded_rot_dif_l_mat
+
+    @property
+    def root_rot_dif_l_multi_future(self) -> torch.Tensor:
+        """Return reference root orientation relative to robot for all future frames.
+
+        Uses the full robot orientation for canonicalization (preserves heading diff).
+
+        Returns:
+            Tensor of shape ``(num_envs, num_future_frames * 6)``.
+        """
+        ref_root_quat = self.motion_lib.get_root_quat_w(
+            self.future_motion_ids, self.future_time_steps
+        )
+        root_rot_dif = quat_mul(
+            quat_inv(
+                self.robot_anchor_quat_w.view(self.num_envs, 1, 4).repeat(
+                    1, self.num_future_frames, 1
+                )
+            ),
+            ref_root_quat.view(self.num_envs, self.num_future_frames, 4),
+        )
+        mat = matrix_from_quat(root_rot_dif)
+        root_rot_dif_l_mat = mat[..., :2].reshape(mat.shape[0], -1)
+        return root_rot_dif_l_mat
+
+    @property
+    def root_rot_dif_heading_multi_future(self) -> torch.Tensor:
+        """Reference root orientation canonicalized by robot heading (yaw) only.
+
+        Unlike root_rot_dif_l_multi_future which uses the full robot orientation,
+        this version only uses the heading (yaw) for canonicalization. This preserves
+        the reference motion's pitch/roll relative to gravity while removing the
+        heading difference.
+
+        Returns:
+            torch.Tensor: 6D rotation matrix representation (first 2 columns),
+                shape (num_envs, num_future_frames * 6)
+        """
+        ref_root_quat = self.motion_lib.get_root_quat_w(
+            self.future_motion_ids, self.future_time_steps
+        )
+        # Use only the heading (yaw) of the robot orientation for canonicalization
+        root_rot_dif = quat_mul(
+            quat_inv(
+                self.anchor_heading_quat.view(self.num_envs, 1, 4).expand(
+                    -1, self.num_future_frames, -1
+                )
+            ),
+            ref_root_quat.view(self.num_envs, self.num_future_frames, 4),
+        )
+        mat = matrix_from_quat(root_rot_dif)
+        root_rot_dif_heading_mat = mat[..., :2].reshape(mat.shape[0], -1)
+        return root_rot_dif_heading_mat
+
+    @property
+    def root_rot_dif_refheading_multi_future(self) -> torch.Tensor:
+        """Reference root orientation canonicalized by the first future frame's heading.
+
+        Instead of using the robot's current heading for canonicalization, this uses
+        the heading of the first (immediate) target frame from the reference motion.
+        This makes the trajectory representation independent of the robot's heading.
+
+        Returns:
+            torch.Tensor: 6D rotation matrix representation (first 2 columns),
+                shape (num_envs, num_future_frames * 6)
+        """
+        ref_root_quat = self.motion_lib.get_root_quat_w(
+            self.future_motion_ids, self.future_time_steps
+        ).view(self.num_envs, self.num_future_frames, 4)
+        # Use the heading of the first future frame as the canonical frame
+        ref_first_heading = torch_transform.get_heading_q(ref_root_quat[:, 0, :])
+        root_rot_dif = quat_mul(
+            quat_inv(
+                ref_first_heading.view(self.num_envs, 1, 4).expand(-1, self.num_future_frames, -1)
+            ),
+            ref_root_quat,
+        )
+        mat = matrix_from_quat(root_rot_dif)
+        return mat[..., :2].reshape(mat.shape[0], -1)
+
+    @property
+    def heading_diff_robot_ref(self) -> torch.Tensor:
+        """Relative heading rotation from robot heading to reference first frame heading.
+
+        Computes quat_mul(quat_inv(robot_heading), ref_first_heading), expressing
+        the reference motion's heading direction in the robot's heading frame.
+
+        Returns:
+            torch.Tensor: 6D rotation matrix representation, shape (num_envs, 6)
+        """
+        ref_root_quat = self.motion_lib.get_root_quat_w(
+            self.future_motion_ids, self.future_time_steps
+        ).view(self.num_envs, self.num_future_frames, 4)
+        ref_first_heading = torch_transform.get_heading_q(ref_root_quat[:, 0, :])
+        heading_diff = quat_mul(quat_inv(self.anchor_heading_quat), ref_first_heading)
+        mat = matrix_from_quat(heading_diff)
+        return mat[..., :2].reshape(self.num_envs, -1)  # (num_envs, 6)
+
+    @property
+    def raw_root_quat_w_multi_future(self) -> torch.Tensor:
+        ref_root_quat = self.motion_lib.get_root_quat_w(
+            self.future_motion_ids, self.future_time_steps
+        )
+        return ref_root_quat.reshape(self.num_envs, self.num_future_frames, 4)
+
+    @property
+    def root_rot_w_multi_future(self) -> torch.Tensor:
+        ref_root_quat = self.motion_lib.get_root_quat_w(
+            self.future_motion_ids, self.future_time_steps
+        )
+        mat = matrix_from_quat(ref_root_quat)
+        root_rot_w_mat = mat[..., :2].reshape(mat.shape[0], -1)
+        return root_rot_w_mat
+
+    @property
+    def body_lin_vel_w(self) -> torch.Tensor:
+        return self.motion_lib.get_body_lin_vel_w(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+
+    @property
+    def body_lin_vel_w_multi_future(self) -> torch.Tensor:
+        return self.motion_lib.get_body_lin_vel_w(
+            self.future_motion_ids, self.future_time_steps
+        ).view(self.num_envs, -1)
+
+    @property
+    def body_lin_vel_l(self) -> torch.Tensor:
+        body_lin_vel_w = self.body_lin_vel_w
+        root_quat = self.robot_anchor_quat_w.view(self.num_envs, 1, 4).repeat(1, self.num_bodies, 1)
+        deheaded_vel_l = quat_apply_yaw(quat_inv(root_quat), body_lin_vel_w)
+        return deheaded_vel_l
+
+    @property
+    def body_lin_vel_l_multi_future(self) -> torch.Tensor:
+        body_lin_vel_w = self.body_lin_vel_w_multi_future
+        root_quat = self.robot_anchor_quat_w.view(self.num_envs, 1, 1, 4).repeat(
+            1, self.num_future_frames, self.num_bodies, 1
+        )
+        deheaded_vel_l = quat_apply_yaw(quat_inv(root_quat), body_lin_vel_w)
+        return deheaded_vel_l
+
+    @property
+    def body_ang_vel_w(self) -> torch.Tensor:
+        return self.motion_lib.get_body_ang_vel_w(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+
+    @property
+    def body_ang_vel_w_multi_future(self) -> torch.Tensor:
+        return self.motion_lib.get_body_ang_vel_w(
+            self.future_motion_ids, self.future_time_steps
+        ).view(self.num_envs, -1)
+
+    @property
+    def body_ang_vel_l(self) -> torch.Tensor:
+        body_ang_vel_w = self.body_ang_vel_w
+        root_quat = self.robot_anchor_quat_w.view(self.num_envs, 1, 4).repeat(1, self.num_bodies, 1)
+        deheaded_vel_l = quat_apply_yaw(quat_inv(root_quat), body_ang_vel_w)
+        return deheaded_vel_l
+
+    @property
+    def body_ang_vel_l_multi_future(self) -> torch.Tensor:
+        body_ang_vel_w = self.body_ang_vel_w_multi_future
+        root_quat = self.robot_anchor_quat_w.view(self.num_envs, 1, 1, 4).repeat(
+            1, self.num_future_frames, self.num_bodies, 1
+        )
+        deheaded_vel_l = quat_apply_yaw(quat_inv(root_quat), body_ang_vel_w)
+        return deheaded_vel_l
+
+    @property
+    def anchor_pos_w(self) -> torch.Tensor:
+        """Return reference anchor body position in world frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, 3)``.
+        """
+        return (
+            self.motion_lib.get_body_pos_w(
+                self.motion_ids, self.motion_start_time_steps + self.time_steps
+            )[:, self.motion_anchor_body_index]
+            + self._env.scene.env_origins
+        )
+
+    @property
+    def anchor_pos_w_multi_future(self) -> torch.Tensor:
+        """Return reference anchor positions for all future frames in world frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_future_frames * 3)``.
+        """
+        anchor_pos_w = self.motion_lib.get_body_pos_w(
+            self.future_motion_ids, self.future_time_steps
+        )[:, self.motion_anchor_body_index].view(self.num_envs, self.num_future_frames, -1)
+        anchor_pos_w_env = anchor_pos_w + self._env.scene.env_origins[:, None, :]
+        return anchor_pos_w_env.reshape(self.num_envs, -1)
+
+    @property
+    def anchor_quat_w(self) -> torch.Tensor:
+        """Return reference anchor body quaternion in world frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, 4)``.
+        """
+        return self.motion_lib.get_body_quat_w(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )[:, self.motion_anchor_body_index]
+
+    @property
+    def anchor_quat_w_multi_future(self) -> torch.Tensor:
+        """Return reference anchor quaternions for all future frames.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_future_frames * 4)``.
+        """
+        return self.motion_lib.get_body_quat_w(self.future_motion_ids, self.future_time_steps)[
+            :, self.motion_anchor_body_index
+        ].reshape(self.num_envs, -1)
+
+    @property
+    def anchor_ori_refheading(self) -> torch.Tensor:
+        """Current anchor orientation canonicalized by its own heading.
+
+        Uses get_heading_q(anchor_quat_w) as the canonical frame, removing the
+        heading component and preserving pitch/roll relative to gravity.
+
+        Returns:
+            torch.Tensor: 6D rotation matrix representation, shape (num_envs, 6)
+        """
+        ref_heading = torch_transform.get_heading_q(self.anchor_quat_w)
+        ori = quat_mul(quat_inv(ref_heading), self.anchor_quat_w)
+        mat = matrix_from_quat(ori)
+        return mat[..., :2].reshape(self.num_envs, -1)
+
+    @property
+    def anchor_ori_heading(self) -> torch.Tensor:
+        """Current anchor orientation canonicalized by robot heading (yaw).
+
+        Uses get_heading_q(robot_anchor_quat_w) as the canonical frame, preserving
+        the reference motion's pitch/roll relative to gravity while removing the
+        robot's heading.
+
+        Returns:
+            torch.Tensor: 6D rotation matrix representation, shape (num_envs, 6)
+        """
+        robot_heading = self.anchor_heading_quat
+        ori = quat_mul(quat_inv(robot_heading), self.anchor_quat_w)
+        mat = matrix_from_quat(ori)
+        return mat[..., :2].reshape(self.num_envs, -1)
+
+    @property
+    def anchor_lin_vel_w(self) -> torch.Tensor:
+        return self.motion_lib.get_body_lin_vel_w(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )[:, self.motion_anchor_body_index]
+
+    @property
+    def anchor_lin_vel_w_multi_future(self) -> torch.Tensor:
+        return self.motion_lib.get_body_lin_vel_w(self.future_motion_ids, self.future_time_steps)[
+            :, self.motion_anchor_body_index
+        ].view(self.num_envs, -1)
+
+    @property
+    def anchor_ang_vel_w(self) -> torch.Tensor:
+        return self.motion_lib.get_body_ang_vel_w(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )[:, self.motion_anchor_body_index]
+
+    @property
+    def anchor_ang_vel_w_multi_future(self) -> torch.Tensor:
+        return self.motion_lib.get_body_ang_vel_w(self.future_motion_ids, self.future_time_steps)[
+            :, self.motion_anchor_body_index
+        ].view(self.num_envs, -1)
+
+    @property
+    def vr_3point_body_quat_w(self) -> torch.Tensor:
+        return self.motion_lib.get_body_quat_w(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )[:, self.vr_3point_body_indices_motion]
+
+    @property
+    def reward_point_body_quat_w(self) -> torch.Tensor:
+        return self.motion_lib.get_body_quat_w(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )[:, self.reward_point_body_indices_motion]
+
+    @property
+    def vr_3point_body_quat_w_multi_future(self) -> torch.Tensor:
+        return self.motion_lib.get_body_quat_w(self.future_motion_ids, self.future_time_steps)[
+            :, self.vr_3point_body_indices_motion
+        ].view(self.num_envs, self.num_future_frames, len(self.cfg.vr_3point_body), -1)
+
+    @property
+    def head_orn_w_multi_future(self) -> torch.Tensor:
+        return self.motion_lib.get_body_quat_w(self.future_motion_ids, self.future_time_steps)[
+            :, self.vr_3point_body_indices_motion[2]
+        ].view(self.num_envs, self.num_future_frames, -1)
+
+    @property
+    def reward_point_body_pos_w(self) -> torch.Tensor:
+        reward_point_original = self.motion_lib.get_body_pos_w(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )[:, self.reward_point_body_indices_motion]
+        return (
+            reward_point_original
+            + quat_apply(self.reward_point_body_quat_w, self.reward_point_body_offsets)
+            + self._env.scene.env_origins[:, None, :]
+        )
+
+    @property
+    def vr_3point_body_pos_w(self) -> torch.Tensor:
+        """Return reference VR 3-point body positions (with offsets) in world frame.
+
+        The 3 points are typically left wrist, right wrist, and head. Offsets
+        allow tracking a point displaced from the body origin (e.g., palm center).
+
+        Returns:
+            Tensor of shape ``(num_envs, 3, 3)``.
+        """
+        vr_3point_original = self.motion_lib.get_body_pos_w(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )[:, self.vr_3point_body_indices_motion]
+        return (
+            vr_3point_original
+            + quat_apply(self.vr_3point_body_quat_w, self.vr_3point_body_offsets)
+            + self._env.scene.env_origins[:, None, :]
+        )
+
+    @property
+    def vr_3point_body_pos_w_multi_future(self) -> torch.Tensor:
+        """Return reference VR 3-point positions for all future frames in world frame.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_future_frames, 3, 3)``.
+        """
+        vr_3point_original = self.motion_lib.get_body_pos_w(
+            self.future_motion_ids, self.future_time_steps
+        )[:, self.vr_3point_body_indices_motion].view(
+            self.num_envs, self.num_future_frames, len(self.cfg.vr_3point_body), -1
+        )
+        vr_3point_offset_extend = self.vr_3point_body_offsets.unsqueeze(1).repeat(
+            1, self.num_future_frames, 1, 1
+        )
+        vr_3point_pos_w = (
+            vr_3point_original
+            + quat_apply(self.vr_3point_body_quat_w_multi_future, vr_3point_offset_extend)
+            + self._env.scene.env_origins[:, None, None, :]
+        )
+        return vr_3point_pos_w
+
+    @property
+    def robot_joint_pos(self) -> torch.Tensor:
+        return self.robot.data.joint_pos
+
+    @property
+    def robot_joint_vel(self) -> torch.Tensor:
+        return self.robot.data.joint_vel
+
+    @property
+    def robot_body_pos_w(self) -> torch.Tensor:
+        return self.robot.data.body_pos_w[:, self.body_indexes]
+
+    @property
+    def robot_body_quat_w(self) -> torch.Tensor:
+        return self.robot.data.body_quat_w[:, self.body_indexes]
+
+    @property
+    def robot_body_lin_vel_w(self) -> torch.Tensor:
+        return self.robot.data.body_lin_vel_w[:, self.body_indexes]
+
+    @property
+    def robot_body_ang_vel_w(self) -> torch.Tensor:
+        return self.robot.data.body_ang_vel_w[:, self.body_indexes]
+
+    @property
+    def robot_anchor_pos_w(self) -> torch.Tensor:
+        """Return the robot's current anchor body position in world frame.
+
+        In offline mode, falls back to the reference motion root position.
+
+        Returns:
+            Tensor of shape ``(num_envs, 3)``.
+        """
+        if getattr(self, "_offline", False):
+            return self.motion_lib.get_root_pos_w(
+                self.motion_ids, self.motion_start_time_steps + self.time_steps
+            )
+        return self.robot.data.body_pos_w[:, self.robot_anchor_body_index]
+
+    @property
+    def robot_anchor_quat_w(self) -> torch.Tensor:
+        """Return the robot's current anchor body quaternion in world frame.
+
+        When ``use_ref_motion_root_quat_w_as_anchor`` is True or in offline mode,
+        returns the reference motion root orientation (optionally with added noise)
+        instead of the simulated robot state.
+
+        Returns:
+            Tensor of shape ``(num_envs, 4)``.
+        """
+        if self.use_ref_motion_root_quat_w_as_anchor or getattr(self, "_offline", False):
+            ref_root_quat = self.motion_lib.get_root_quat_w(
+                self.motion_ids, self.motion_start_time_steps + self.time_steps
+            )
+            if self.ref_motion_root_rotation_noise is not None:
+                ref_root_quat = quat_mul(ref_root_quat, self.ref_motion_root_rotation_noise)
+            return ref_root_quat
+        return self.robot.data.body_quat_w[:, self.robot_anchor_body_index]
+
+    @property
+    def robot_anchor_lin_vel_w(self) -> torch.Tensor:
+        assert not getattr(
+            self, "_offline", False
+        ), "robot_anchor_lin_vel_w is not available in offline mode"
+        return self.robot.data.body_lin_vel_w[:, self.robot_anchor_body_index]
+
+    @property
+    def robot_anchor_ang_vel_w(self) -> torch.Tensor:
+        assert not getattr(
+            self, "_offline", False
+        ), "robot_anchor_ang_vel_w is not available in offline mode"
+        return self.robot.data.body_ang_vel_w[:, self.robot_anchor_body_index]
+
+    @property
+    def robot_vr_3point_quat_w(self) -> torch.Tensor:
+        return self.robot.data.body_quat_w[:, self.vr_3point_body_indices]
+
+    @property
+    def robot_reward_point_body_pos_w(self) -> torch.Tensor:
+        return self.robot.data.body_pos_w[:, self.reward_point_body_indices] + quat_apply(
+            self.robot.data.body_quat_w[:, self.reward_point_body_indices],
+            self.reward_point_body_offsets,
+        )
+
+    @property
+    def robot_vr_3point_pos_w(self) -> torch.Tensor:
+        return self.robot.data.body_pos_w[:, self.vr_3point_body_indices] + quat_apply(
+            self.robot.data.body_quat_w[:, self.vr_3point_body_indices], self.vr_3point_body_offsets
+        )
+
+    @property
+    def feet_l(self) -> torch.Tensor:
+        return self.motion_lib.get_feet_l(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+
+    @property
+    def feet_r(self) -> torch.Tensor:
+        return self.motion_lib.get_feet_r(
+            self.motion_ids, self.motion_start_time_steps + self.time_steps
+        )
+
+    @property
+    def episode_encoder_index(self) -> torch.Tensor:
+        return self.encoder_index
+
+    def _update_metrics(self):
+        """Compute tracking error metrics between reference motion and robot state.
+
+        Populates ``self.metrics`` with per-env errors for anchor position/rotation,
+        body position/rotation, and joint position/velocity. When there is a DOF
+        mismatch (extra finger joints), only body joints are compared.
+        """
+        self.metrics["error_anchor_pos"] = torch.norm(
+            self.anchor_pos_w - self.robot_anchor_pos_w, dim=-1
+        )
+        self.metrics["error_anchor_rot"] = quat_error_magnitude(
+            self.anchor_quat_w, self.robot_anchor_quat_w
+        )
+        self.metrics["error_anchor_lin_vel"] = torch.norm(
+            self.anchor_lin_vel_w - self.robot_anchor_lin_vel_w, dim=-1
+        )
+        self.metrics["error_anchor_ang_vel"] = torch.norm(
+            self.anchor_ang_vel_w - self.robot_anchor_ang_vel_w, dim=-1
+        )
+
+        self.metrics["error_body_pos"] = torch.norm(
+            self.body_pos_relative_w - self.robot_body_pos_w, dim=-1
+        ).mean(dim=-1)
+        self.metrics["error_body_rot"] = quat_error_magnitude(
+            self.body_quat_relative_w, self.robot_body_quat_w
+        ).mean(dim=-1)
+
+        self.metrics["error_body_lin_vel"] = torch.norm(
+            self.body_lin_vel_w - self.robot_body_lin_vel_w, dim=-1
+        ).mean(dim=-1)
+        self.metrics["error_body_ang_vel"] = torch.norm(
+            self.body_ang_vel_w - self.robot_body_ang_vel_w, dim=-1
+        ).mean(dim=-1)
+
+        # Compare only body joints when there's a DOF mismatch (motion lib has fewer DOFs than robot)
+        if self.has_dof_mismatch:
+            robot_body_joint_pos = self.robot_joint_pos[:, self.body_joint_indices]
+            robot_body_joint_vel = self.robot_joint_vel[:, self.body_joint_indices]
+            self.metrics["error_joint_pos"] = torch.abs(self.joint_pos - robot_body_joint_pos).mean(
+                dim=-1
+            )
+            self.metrics["error_joint_vel"] = torch.abs(self.joint_vel - robot_body_joint_vel).mean(
+                dim=-1
+            )
+        else:
+            self.metrics["error_joint_pos"] = torch.abs(self.joint_pos - self.robot_joint_pos).mean(
+                dim=-1
+            )
+            self.metrics["error_joint_vel"] = torch.abs(self.joint_vel - self.robot_joint_vel).mean(
+                dim=-1
+            )
+
+    def resample_all_commands(self):
+        """Resample motion clips and reset state for all environments at once."""
+        self._resample_command(torch.arange(self.num_envs))
+
+    def _load_contact_data(self):
+        """Load contact data from file or directory and validate frame counts against motion library.
+
+        Supports both:
+        - Single file: loads one pkl file
+        - Directory: loads all pkl files and merges them
+
+        Sets:
+            self._contact_data: Raw contact data dict
+            self._first_contact_frame: Dict mapping motion_name -> first contact frame index
+        """
+        import joblib
+
+        self._first_contact_frame = None
+        self._contact_data = None
+        self._first_contact_lookup = None
+        self._per_env_first_contact = torch.zeros(
+            self.num_envs, device=self.device, dtype=torch.long
+        )
+        self._motion_contact_flags = None
+
+        if self.cfg.contact_file is None or not os.path.exists(self.cfg.contact_file):
+            # Fallback: derive first contact from motion lib's in_contact labels
+            self._derive_first_contact_from_in_contact_labels()
+            return
+
+        # Support both file and directory modes
+        if os.path.isfile(self.cfg.contact_file):
+            contact_data = joblib.load(self.cfg.contact_file)
+        elif os.path.isdir(self.cfg.contact_file):
+            # Directory mode: load all pkl files
+            # Only load data if internal key EXACTLY matches filename (without .pkl)
+            contact_data = {}
+            pkl_files = glob.glob(os.path.join(self.cfg.contact_file, "*.pkl"))
+            for pkl_file in sorted(pkl_files):
+                expected_key = os.path.splitext(os.path.basename(pkl_file))[0]
+                try:
+                    data = joblib.load(pkl_file)
+                    if expected_key in data:
+                        contact_data[expected_key] = data[expected_key]
+                except Exception as e:  # noqa: BLE001
+                    print(f"  Warning: Failed to load {pkl_file}: {e}")  # noqa: T201
+            print(  # noqa: T201
+                f"[TrackingCommand] Loaded {len(contact_data)} contact sequences from {self.cfg.contact_file}"
+            )
+        else:
+            print(  # noqa: T201
+                f"[TrackingCommand] Warning: contact_file path invalid: {self.cfg.contact_file}"
+            )
+            return
+
+        # Filter contact data to only include motions that exist in motion_lib
+        # This respects filter_motion_keys used by motion_lib
+        motion_keys_set = set(self.motion_lib.curr_motion_keys)
+        filtered_contact_data = {k: v for k, v in contact_data.items() if k in motion_keys_set}
+        if len(filtered_contact_data) < len(contact_data):
+            print(  # noqa: T201
+                f"[TrackingCommand] Filtered contact data: {len(filtered_contact_data)}/{len(contact_data)} "
+                f"motions match loaded motion keys"
+            )
+        self._contact_data = filtered_contact_data
+
+        # Find first frame with contact for each motion
+        # Contact file structure: {motion_name: {body: (N, 10475), object: (N, obj_verts), ...}}
+        first_contact_frames = {}
+        for motion_name, motion_data in filtered_contact_data.items():
+            object_contact = motion_data.get("object", None)
+            if object_contact is not None:
+                # Find first frame where any body vertex is in contact
+                contact_per_frame = (object_contact != 0).sum(axis=1)
+                contact_frames = np.where(contact_per_frame > 0)[0]
+                if len(contact_frames) > 0:
+                    first_contact_frames[motion_name] = int(contact_frames[0])
+                else:
+                    # No contact, use last frame
+                    first_contact_frames[motion_name] = object_contact.shape[0]
+            else:
+                first_contact_frames[motion_name] = 0
+
+        self._first_contact_frame = first_contact_frames
+        print(f"[TrackingCommand] Loaded contact data from {self.cfg.contact_file}")  # noqa: T201
+        for motion_name, first_frame in first_contact_frames.items():
+            print(f"  {motion_name}: first contact at frame {first_frame}")  # noqa: T201
+
+        # Validate and align contact data frame counts to match motion data
+        self._validate_and_align_contact_frame_counts(filtered_contact_data)
+
+        # Build motion_key -> first_contact_frame lookup tensor for efficient per-env updates
+        self._build_first_contact_lookup()
+
+        # Preprocess contact flags for each motion: motion_id -> (num_frames,) bool tensor
+        self._build_motion_contact_flags()
+
+    def _derive_first_contact_from_in_contact_labels(self):
+        """Derive first contact frames from motion lib's in_contact labels.
+
+        When no contact_file is provided, but the motion lib has per-frame
+        in_contact labels (from contact_points_left_hand/right_hand in the
+        object motion data), derive the first contact frame for each motion.
+        """
+        hand = getattr(self.cfg, "sample_before_contact_hand", "right_hand")
+        side = "left" if hand == "left_hand" else "right"
+        attr = f"_motion_object_in_contact_{side}"
+
+        if not hasattr(self.motion_lib, attr):
+            return
+
+        in_contact_tensor = getattr(self.motion_lib, attr)  # (total_frames,)
+        length_starts = self.motion_lib.length_starts
+        num_motions = len(self.motion_lib.curr_motion_keys)
+
+        first_contact_frames = {}
+        for motion_idx in range(num_motions):
+            motion_key = self.motion_lib.curr_motion_keys[motion_idx]
+            start = length_starts[motion_idx].item()
+            num_frames = self.motion_lib._motion_num_frames[motion_idx].item()  # noqa: SLF001
+            end = start + num_frames
+
+            motion_in_contact = in_contact_tensor[start:end]
+            contact_indices = torch.nonzero(motion_in_contact > 0.5, as_tuple=False)
+            if len(contact_indices) > 0:
+                first_contact_frames[motion_key] = int(contact_indices[0].item())
+            else:
+                first_contact_frames[motion_key] = num_frames
+
+        self._first_contact_frame = first_contact_frames
+        print(  # noqa: T201
+            f"[TrackingCommand] Derived first contact from in_contact labels ({hand}):"
+        )
+        for motion_name, first_frame in first_contact_frames.items():
+            print(f"  {motion_name}: first contact at frame {first_frame}")  # noqa: T201
+
+        self._build_first_contact_lookup()
+
+    def _build_first_contact_lookup(self):
+        """Build a tensor for efficient per-env first contact frame lookup."""
+        if self._first_contact_frame is None:
+            self._first_contact_lookup = None
+            return
+
+        # Create a tensor indexed by motion_id for O(1) lookup
+        num_motions = len(self.motion_lib.curr_motion_keys)
+        self._first_contact_lookup = torch.zeros(num_motions, device=self.device, dtype=torch.long)
+
+        for motion_idx, motion_key in enumerate(self.motion_lib.curr_motion_keys):
+            if motion_key in self._first_contact_frame:
+                self._first_contact_lookup[motion_idx] = self._first_contact_frame[motion_key]
+            else:
+                # Fallback to 0 if not found
+                self._first_contact_lookup[motion_idx] = 0
+
+    def _build_motion_contact_flags(self):
+        """Preprocess contact flags for each motion: motion_id -> (num_frames,) bool tensor."""
+        if self._contact_data is None or len(self._contact_data) == 0:
+            self._motion_contact_flags = None
+            return
+
+        self._motion_contact_flags = {}
+
+        for motion_idx, motion_key in enumerate(self.motion_lib.curr_motion_keys):
+            if motion_key not in self._contact_data:
+                continue
+
+            motion_contact_data = self._contact_data[motion_key]
+            object_contact = motion_contact_data.get("object", None)
+
+            if object_contact is not None:
+                # Convert to torch tensor if needed
+                if not isinstance(object_contact, torch.Tensor):
+                    object_contact_tensor = torch.from_numpy(object_contact).to(self.device)
+                else:
+                    object_contact_tensor = object_contact.to(self.device)
+
+                # Compute contact flags: (num_frames,) - True if any vertex has contact
+                contact_per_frame = (object_contact_tensor != 0).sum(dim=1)  # (num_frames,)
+                contact_flags = contact_per_frame > 0  # (num_frames,) bool
+
+                self._motion_contact_flags[motion_idx] = contact_flags
+
+    def _update_per_env_first_contact(self, env_ids):
+        """Update _per_env_first_contact for given env_ids based on their assigned motion."""
+        if self._first_contact_lookup is None:
+            return
+
+        # Vectorized lookup using motion_ids as indices
+        motion_ids = self.motion_ids[env_ids]
+        self._per_env_first_contact[env_ids] = self._first_contact_lookup[motion_ids]
+
+    def _validate_and_align_contact_frame_counts(self, contact_data: dict):
+        """Validate contact data frame counts match the loaded motion frame counts.
+
+        Allows a tolerance of ±3 frames due to slight duration differences between
+        GRAB (120Hz) and robot motion (30Hz) sources.
+
+        Args:
+            contact_data: Dict mapping motion_name -> contact arrays
+
+        Raises:
+            AssertionError: If frame count difference exceeds 3 frames
+        """
+        FRAME_TOLERANCE = 3
+
+        for motion_name, motion_contact in contact_data.items():
+            # Get contact frame count
+            contact_frames = None
+            if "object" in motion_contact and motion_contact["object"] is not None:
+                contact_frames = motion_contact["object"].shape[0]
+            elif "body" in motion_contact and motion_contact["body"] is not None:
+                contact_frames = motion_contact["body"].shape[0]
+
+            if contact_frames is None:
+                continue
+
+            # Find matching motion in motion_lib and get actual frame count
+            motion_idx = None
+            for idx, key in enumerate(self.motion_lib.curr_motion_keys):
+                if key == motion_name:
+                    motion_idx = idx
+                    break
+
+            if motion_idx is None:
+                print(  # noqa: T201
+                    f"[TrackingCommand] Warning: Contact motion '{motion_name}' "
+                    f"not found in loaded motions"
+                )
+                continue
+
+            # Get actual frame count from loaded motion
+            motion_frames = int(
+                self.motion_lib._motion_num_frames[motion_idx].item()  # noqa: SLF001
+            )
+            frame_diff = abs(contact_frames - motion_frames)
+
+            if frame_diff == 0:
+                print(f"  {motion_name}: {contact_frames} frames (exact match)")  # noqa: T201
+            elif frame_diff <= FRAME_TOLERANCE:
+                print(  # noqa: T201
+                    f"  {motion_name}: {contact_frames} frames "
+                    f"(motion lib: {motion_frames}, diff: {contact_frames - motion_frames})"
+                )
+            else:
+                raise AssertionError(
+                    f"[TrackingCommand] Frame count mismatch too large for motion '{motion_name}':\n"
+                    f"  Contact data: {contact_frames} frames\n"
+                    f"  Motion lib: {motion_frames} frames\n"
+                    f"  Difference: {frame_diff} frames (max allowed: ±{FRAME_TOLERANCE})"
+                )
+
+    def _load_table_meta(self, motion_key: str):
+        """Load and cache table meta info (table_pos, table_quat, table_size) for a motion.
+        Derives meta path from motion_file path. Falls back to motion file if meta not found.
+
+        For GeniHOI data:
+            - table_pos: [x, y, z] center position of the table
+            - table_quat: [w, x, y, z] quaternion (identity for cuboid tables)
+            - table_size: [width, depth, thickness] dimensions of the table
+
+        For GRAB data:
+            - table_pos, table_quat: from meta file
+            - table_size: not provided (uses USD with scene_scale)
+        """  # noqa: D205
+        import joblib
+
+        try:
+            # Derive meta path from motion_file path
+            motion_file = (
+                self.cfg.motion_lib_cfg.get("motion_file", "") if self.cfg.motion_lib_cfg else ""
+            )
+            if motion_file:
+                if os.path.isdir(motion_file):
+                    meta_dir = motion_file.replace("/robot", "/meta")
+                elif "/robot" in motion_file:
+                    meta_dir = os.path.dirname(motion_file).replace("/robot", "/meta")
+                else:
+                    meta_dir = "data/motion_lib_grab/meta"
+                meta_file = os.path.join(meta_dir, f"{motion_key}.pkl")
+            else:
+                meta_file = f"data/motion_lib_grab/meta/{motion_key}.pkl"
+
+            if os.path.exists(meta_file):
+                meta = joblib.load(meta_file)
+                self._table_meta_cache[motion_key] = {
+                    "table_pos": torch.tensor(
+                        meta.get("table_pos", [0.0, 0.0, 0.8]), device=self.device
+                    ).float(),
+                    "table_quat": torch.tensor(
+                        meta.get("table_quat", [1.0, 0.0, 0.0, 0.0]), device=self.device
+                    ).float(),
+                    "table_size": torch.tensor(
+                        meta.get("table_size", [1.0, 0.6, 0.04]), device=self.device
+                    ).float(),
+                }
+                return
+
+            # Fallback: try to get table data from motion file directly
+            # Cache ALL motions at once to avoid reloading the file
+            if motion_file and os.path.isfile(motion_file):
+                motion_file_data = joblib.load(motion_file)
+                # Cache table data for ALL motions in the file
+                for mk, motion_data in motion_file_data.items():
+                    if mk in self._table_meta_cache:
+                        continue  # Already cached
+                    if (
+                        isinstance(motion_data, dict)
+                        and "table_pos" in motion_data
+                        and "table_quat" in motion_data
+                    ):
+                        self._table_meta_cache[mk] = {
+                            "table_pos": torch.tensor(
+                                motion_data["table_pos"], device=self.device
+                            ).float(),
+                            "table_quat": torch.tensor(
+                                motion_data["table_quat"], device=self.device
+                            ).float(),
+                        }
+                    else:
+                        self._table_meta_cache[mk] = None
+                # After caching all, check if we got the one we needed
+                if motion_key in self._table_meta_cache:
+                    return
+        except Exception:  # noqa: BLE001, S110
+            pass
+
+        self._table_meta_cache[motion_key] = None
+
+    def _sample_before_contact(
+        self, env_ids: Sequence[int], sampled_times: torch.Tensor
+    ) -> torch.Tensor:
+        """Sample timestamps before the first contact frame for contact-based initialization.
+
+        Args:
+            env_ids: Environment indices to resample
+            sampled_times: Originally sampled time steps
+
+        Returns:
+            Modified sampled_times with timestamps clamped to before first contact
+        """
+        if self._first_contact_frame is None or len(self._first_contact_frame) == 0:
+            return sampled_times
+        # Get motion keys from motion library
+        curr_motion_keys = getattr(self.motion_lib, "curr_motion_keys", None)
+
+        for i, env_idx in enumerate(env_ids):
+            motion_id = self.motion_ids[env_idx].item()
+
+            # Get motion key for this motion_id
+            if curr_motion_keys is not None and motion_id < len(curr_motion_keys):
+                motion_key = curr_motion_keys[motion_id]
+            else:
+                # Fallback: use first contact key
+                motion_key = list(self._first_contact_frame.keys())[0]  # noqa: RUF015
+            first_contact = None
+            if motion_key in self._first_contact_frame:
+                first_contact = self._first_contact_frame[motion_key]
+            else:
+                for contact_key in self._first_contact_frame.keys():  # noqa: SIM118
+                    if contact_key in motion_key or motion_key in contact_key:
+                        first_contact = self._first_contact_frame[contact_key]
+                        break
+                if first_contact is None:
+                    first_contact = list(self._first_contact_frame.values())[0]  # noqa: RUF015
+
+            # Sample uniformly from [0, first_contact - margin)
+            margin = getattr(self.cfg, "sample_before_contact_margin", 10)
+            if first_contact > margin:
+                sampled_times[i] = torch.randint(
+                    0, first_contact - margin, (1,), device=self.device, dtype=sampled_times.dtype
+                )
+            else:
+                sampled_times[i] = 0
+
+        return sampled_times
+
+    def _resample_command(self, env_ids: Sequence[int]):
+        """Resample motion clips, reset robot state, and position objects for given envs.
+
+        This is the main episode-reset handler. It performs the following in order:
+
+        1. Sample new motion IDs and start times (respecting evaluation mode,
+           paired motions, multi-object mode, and adaptive sampling).
+        2. Resample encoder mode (G1/SMPL/teleop) per env.
+        3. Apply pose and velocity randomization (skipped during evaluation).
+        4. Handle DOF mismatch by mapping motion lib joints to robot joints.
+        5. Write joint state and root state to the simulator.
+        6. Position scene objects (single or multi-object mode) and tables.
+        7. Cache body-relative poses for reward computation.
+
+        Args:
+            env_ids: Environment indices being reset.
+        """
+        self.time_steps[env_ids] = 0
+        # Variable frames: resample per-env num_frames at episode reset
+        if self.variable_frames_enabled and len(env_ids) > 0:
+            idx = torch.randint(0, len(self._frame_choices), (len(env_ids),), device=self.device)
+            self.per_env_num_frames[env_ids] = self._frame_choices[idx]
+        if len(env_ids) > 0:
+            if self.is_evaluating:
+                self.motion_ids[env_ids] = (
+                    torch.arange(self.num_envs).to(self.device)
+                    % self.motion_lib._num_motions  # noqa: SLF001
+                )[env_ids]
+                self.motion_start_time_steps[env_ids] = 0
+            elif self.cfg.use_paired_motions:
+                self.motion_ids[env_ids] = (
+                    torch.arange(self.num_envs).to(self.device)
+                    % self.motion_lib._num_motions  # noqa: SLF001
+                )[env_ids]
+
+            elif self._multi_object_mode:
+                # MULTI-OBJECT MODE: Resetting envs sample a new motion (and corresponding object)
+                # Over time, staggered resets lead to different envs using different objects,
+                # which provides training diversity. Object positioning (below) handles per-env instances.
+                new_motion_id = self.motion_lib.sample_motions(1)[0]
+                self.motion_ids[env_ids] = new_motion_id
+                self.motion_start_time_steps[env_ids] = self.motion_lib.sample_time_steps(
+                    self.motion_ids[env_ids], truncate_time=None
+                )
+            else:
+                if self.use_adaptive_sampling:
+                    sampled_ids, sampled_times = self.motion_lib.sample_motion_ids_and_time_steps(
+                        len(env_ids)
+                    )
+                    self.motion_ids[env_ids] = sampled_ids.to(self.motion_ids.dtype)
+                    sampled_times = sampled_times.to(self.motion_start_time_steps.dtype)
+                else:
+                    self.motion_ids[env_ids] = self.motion_lib.sample_motions(len(env_ids))
+                    sampled_times = self.motion_lib.sample_time_steps(
+                        self.motion_ids[env_ids], truncate_time=None
+                    )
+
+                # Override to sample from initial frames if configured
+                if self.cfg.sample_from_n_initial_frames is not None:
+                    # Sample uniformly from first N frames
+                    n_frames = self.cfg.sample_from_n_initial_frames
+                    sampled_times = torch.randint(
+                        0, n_frames, (len(env_ids),), dtype=sampled_times.dtype, device=self.device
+                    )
+                elif self.cfg.start_from_first_frame:
+                    sampled_times.zero_()
+
+                # Contact-based initialization: sample timestamps before first contact frame
+                if self.cfg.sample_before_contact and self._first_contact_frame is not None:
+                    sampled_times = self._sample_before_contact(env_ids, sampled_times)
+
+                self.motion_start_time_steps[env_ids] = sampled_times
+
+            if self.encoder_sample_probs is not None:
+                has_smpl = self.motion_lib.motion_has_smpl[self.motion_ids[env_ids]]
+                if self.soma_encoder_index is not None and hasattr(
+                    self.motion_lib, "motion_has_soma"
+                ):
+                    has_soma = self.motion_lib.motion_has_soma[self.motion_ids[env_ids]]
+                    sampling_cases = [
+                        (env_ids[has_smpl & has_soma], self.encoder_sample_probs),
+                        (env_ids[has_smpl & ~has_soma], self.encoder_sample_probs_no_soma),
+                        (env_ids[~has_smpl & has_soma], self.encoder_sample_probs_no_smpl),
+                        (
+                            env_ids[~has_smpl & ~has_soma],
+                            self.encoder_sample_probs_no_smpl_no_soma,
+                        ),
+                    ]
+                else:
+                    sampling_cases = [
+                        (env_ids[has_smpl], self.encoder_sample_probs),
+                        (env_ids[~has_smpl], self.encoder_sample_probs_no_smpl),
+                    ]
+                for subset_ids, probs in sampling_cases:
+                    if len(subset_ids) > 0:
+                        encoder_index = torch.multinomial(
+                            probs, len(subset_ids), replacement=True
+                        ).to(self.device)
+                        self.encoder_index[subset_ids] = 0
+                        self.encoder_index[subset_ids, encoder_index] = 1
+
+                # =============================================================
+                # Legacy behavior: SMPL-native envs also activate G1 encoder
+                # This causes G1 tokens to be computed (then overwritten by SMPL)
+                # in the main encoding loop, enabling G1-SMPL latent alignment.
+                #
+                # When optimize_encoders_ratio_for_CHIP=True:
+                #   - Skip this OR logic (cleaner native encoder selection)
+                #   - G1 encoder only runs for G1-native envs in main loop
+                #   - G1 latents for SMPL-native are computed separately in
+                #     aux losses, ONLY when compliance=0 (stiff mode)
+                # =============================================================
+                if (
+                    self.smpl_encoder_index is not None
+                    and not self.optimize_encoders_ratio_for_CHIP
+                ):
+                    use_smpl = self.encoder_index[env_ids, self.smpl_encoder_index]
+                    self.encoder_index[env_ids, self.g1_encoder_index] = (
+                        self.encoder_index[env_ids, self.g1_encoder_index] | use_smpl
+                    )
+                    # Also sample teleop mode when smpl mode is active (for latent alignment)
+                    if (
+                        self.teleop_encoder_index is not None
+                        and self.teleop_sample_prob_when_smpl > 0.0
+                    ):
+                        smpl_env_ids = env_ids[use_smpl.bool()]
+                        if len(smpl_env_ids) > 0:
+                            sample_teleop = (
+                                torch.rand(len(smpl_env_ids), device=self.device)
+                                < self.teleop_sample_prob_when_smpl
+                            )
+                            self.encoder_index[smpl_env_ids, self.teleop_encoder_index] = (
+                                self.encoder_index[smpl_env_ids, self.teleop_encoder_index]
+                                | sample_teleop.long()
+                            )
+
+                # When soma is sampled, also activate g1 (for g1-soma latent alignment)
+                if (
+                    self.soma_encoder_index is not None
+                    and not self.optimize_encoders_ratio_for_CHIP
+                ):
+                    use_soma = self.encoder_index[env_ids, self.soma_encoder_index]
+                    self.encoder_index[env_ids, self.g1_encoder_index] = (
+                        self.encoder_index[env_ids, self.g1_encoder_index] | use_soma
+                    )
+
+            self.motion_num_steps[env_ids] = self.motion_lib.get_motion_num_steps(
+                self.motion_ids[env_ids]
+            )
+            if self.num_future_frames > 1:
+                self.future_motion_ids = self.motion_ids.repeat_interleave(self.num_future_frames)
+                self.smpl_future_motion_ids = self.motion_ids.repeat_interleave(
+                    self.smpl_num_future_frames
+                )
+                self.motion_num_steps[env_ids] = self.motion_lib.get_motion_num_steps(
+                    self.motion_ids[env_ids]
+                )
+
+            # Update per-env first contact frame based on assigned motion
+            if self._first_contact_frame is not None:
+                self._update_per_env_first_contact(env_ids)
+
+        root_pos = self.body_pos_w[:, 0].clone()
+        root_ori = self.body_quat_w[:, 0].clone()
+        root_lin_vel = self.body_lin_vel_w[:, 0].clone()
+        root_ang_vel = self.body_ang_vel_w[:, 0].clone()
+
+        self.running_ref_root_height[env_ids] = self.anchor_pos_w[env_ids, 2]
+
+        # Skip reset randomizations during evaluation — they cause visible stumbling
+        # at the start of rendered episodes and are only needed for training robustness.
+        if not self.is_evaluating:
+            range_list = [
+                self.cfg.pose_range.get(key, (0.0, 0.0))
+                for key in ["x", "y", "z", "roll", "pitch", "yaw"]
+            ]
+            ranges = torch.tensor(range_list, device=self.device)
+            rand_samples = sample_uniform(
+                ranges[:, 0], ranges[:, 1], (len(env_ids), 6), device=self.device
+            )
+            root_pos[env_ids] += rand_samples[:, 0:3]
+            orientations_delta = quat_from_euler_xyz(
+                rand_samples[:, 3], rand_samples[:, 4], rand_samples[:, 5]
+            )
+            root_ori[env_ids] = quat_mul(orientations_delta, root_ori[env_ids])
+            range_list = [
+                self.cfg.velocity_range.get(key, (0.0, 0.0))
+                for key in ["x", "y", "z", "roll", "pitch", "yaw"]
+            ]
+            ranges = torch.tensor(range_list, device=self.device)
+            rand_samples = sample_uniform(
+                ranges[:, 0], ranges[:, 1], (len(env_ids), 6), device=self.device
+            )
+            root_lin_vel[env_ids] += rand_samples[:, :3]
+            root_ang_vel[env_ids] += rand_samples[:, 3:]
+
+        # Handle DOF mismatch between motion library and robot
+        motion_lib_joint_pos = self.joint_pos.clone()  # Shape: [num_envs, motion_lib_num_dof]
+        motion_lib_joint_vel = self.joint_vel.clone()  # Shape: [num_envs, motion_lib_num_dof]
+
+        if self.has_dof_mismatch:
+            # Create full robot joint tensors and map using name-based indices
+            joint_pos = torch.zeros(
+                self.num_envs, self.robot_num_dof, dtype=torch.float32, device=self.device
+            )
+            joint_vel = torch.zeros(
+                self.num_envs, self.robot_num_dof, dtype=torch.float32, device=self.device
+            )
+            joint_pos[:, self.body_joint_indices] = motion_lib_joint_pos
+            joint_vel[:, self.body_joint_indices] = motion_lib_joint_vel
+            joint_pos[:, self.extra_joint_indices] = self.extra_default_positions
+            joint_vel[:, self.extra_joint_indices] = self.extra_default_velocities
+        else:
+            joint_pos = motion_lib_joint_pos
+            joint_vel = motion_lib_joint_vel
+
+        if not self.is_evaluating:
+            joint_pos += sample_uniform(
+                *self.cfg.joint_position_range, joint_pos.shape, joint_pos.device
+            )
+            joint_vel += sample_uniform(
+                *self.cfg.joint_velocity_range, joint_vel.shape, joint_vel.device
+            )
+
+        soft_joint_pos_limits = self.robot.data.soft_joint_pos_limits[env_ids]
+        joint_pos[env_ids] = torch.clip(
+            joint_pos[env_ids], soft_joint_pos_limits[:, :, 0], soft_joint_pos_limits[:, :, 1]
+        )
+
+        ####### Resetting Humaonid States #######
+        self.robot.write_joint_state_to_sim(joint_pos[env_ids], joint_vel[env_ids], env_ids=env_ids)
+        self.robot.write_root_state_to_sim(
+            torch.cat(
+                [
+                    root_pos[env_ids],
+                    root_ori[env_ids],
+                    root_lin_vel[env_ids],
+                    root_ang_vel[env_ids],
+                ],
+                dim=-1,
+            ),
+            env_ids=env_ids,
+        )
+        # Handle object positioning
+        if self._multi_object_mode and len(self._object_names) > 0:
+            # MULTI-OBJECT MODE: Position active object, move others far away
+            # Get active object name from motion key (original name with hyphens)
+            active_motion_id = self.motion_ids[env_ids[0]].item() if len(env_ids) > 0 else 0
+            active_motion_key = self.motion_lib.curr_motion_keys[active_motion_id]
+            # Convert motion key to safe name (hyphens → underscores) for scene lookup
+            active_obj_safe_name = active_motion_key.replace("-", "_")
+            self._active_object_name = active_motion_key
+
+            # Position the active object from motion trajectory
+            active_key = f"object_{active_obj_safe_name}"
+            if active_key in self._env.scene.rigid_objects:
+                obj = self._env.scene[active_key]
+                obj_pos = self._get_object_pos_with_offset(env_ids)
+                # Reset with zero velocity to prevent velocity carryover between episodes
+                zero_vel = torch.zeros(len(env_ids), 6, device=self.device)
+                obj.write_root_state_to_sim(
+                    torch.cat(
+                        [obj_pos, self.object_root_quat[env_ids, 0], zero_vel],
+                        dim=-1,
+                    ),
+                    env_ids=env_ids,
+                )
+            else:
+                print(  # noqa: T201
+                    f"[Warning] Active object '{active_key}' not found in scene. Available: {list(self._env.scene.rigid_objects.keys())[:5]}..."  # noqa: E501
+                )
+
+            # Move all other objects to inactive positions (far away from robot)
+            # Spread objects vertically (Z-axis) - envs naturally separate in X,Y via env_origins
+            inactive_base = INACTIVE_OBJECT_BASE_OFFSET.to(self.device)
+            inactive_quat = torch.tensor([[1.0, 0.0, 0.0, 0.0]], device=self.device).expand(
+                len(env_ids), -1
+            )
+
+            inactive_idx = 0
+            for safe_name in self._object_names:
+                obj_key = f"object_{safe_name}"
+                if obj_key != active_key and obj_key in self._env.scene.rigid_objects:
+                    # Spread in Z: each object type at different depth underground
+                    # X,Y comes from env_origins → naturally separates different envs
+                    z_offset = torch.tensor(
+                        [0.0, 0.0, -inactive_idx * INACTIVE_OBJECT_Z_SPACING], device=self.device
+                    )
+                    inactive_pos = self._env.scene.env_origins[env_ids] + inactive_base + z_offset
+                    # Reset with zero velocity to prevent velocity carryover between episodes
+                    inactive_zero_vel = torch.zeros(len(env_ids), 6, device=self.device)
+                    inactive_state = torch.cat(
+                        [inactive_pos, inactive_quat, inactive_zero_vel], dim=-1
+                    )
+                    self._env.scene[obj_key].write_root_state_to_sim(
+                        inactive_state, env_ids=env_ids
+                    )
+                    inactive_idx += 1
+        elif "object" in self._env.scene.rigid_objects:
+            # SINGLE OBJECT MODE: Existing behavior
+            obj = self._env.scene["object"]
+            obj_pos = self._get_object_pos_with_offset(env_ids)
+            # Reset with zero velocity to prevent velocity carryover between episodes
+            zero_vel = torch.zeros(len(env_ids), 6, device=self.device)
+            obj.write_root_state_to_sim(
+                torch.cat([obj_pos, self.object_root_quat[env_ids, 0], zero_vel], dim=-1),
+                env_ids=env_ids,
+            )
+        if "table" in self._env.scene.rigid_objects:
+            table = self._env.scene["table"]
+
+            # Per-env table positions based on each env's assigned motion
+            table_pos_list = []
+            table_quat_list = []
+
+            # Prefetch motion_ids to CPU to avoid per-iteration GPU sync
+            motion_ids_cpu = self.motion_ids.cpu()
+
+            for i, env_idx in enumerate(env_ids):  # noqa: B007
+                motion_idx = motion_ids_cpu[env_idx].item()
+                motion_key = self.motion_lib.curr_motion_keys[motion_idx]
+
+                # Load and cache meta if not already cached
+                if motion_key not in self._table_meta_cache:
+                    self._load_table_meta(motion_key)
+
+                # Use cached meta if available
+                cached = self._table_meta_cache.get(motion_key)
+                if cached is not None:
+                    pos = cached["table_pos"].clone()
+                    quat = cached["table_quat"].clone()
+                    # Add env_origin offset
+                    pos = pos + self._env.scene.env_origins[env_idx]
+                    table_pos_list.append(pos)
+                    table_quat_list.append(quat)
+                else:
+                    # Fallback: derive from object position with hardcoded offset
+                    pos = self.object_root_pos[env_idx, 0].clone()
+                    pos[2] = 0.76  # Table height
+                    pos[1] -= 0.15
+                    quat = torch.tensor([1.0, 0.0, 0.0, 0.0], device=self._env.device)
+                    table_pos_list.append(pos)
+                    table_quat_list.append(quat)
+
+            if len(table_pos_list) > 0:
+                table_pos = torch.stack(table_pos_list, dim=0)
+                table_quat = torch.stack(table_quat_list, dim=0)
+
+                # Apply table_offset if configured
+                if self.cfg.table_offset is not None:
+                    table_offset = torch.tensor(
+                        self.cfg.table_offset, device=self._env.device, dtype=table_pos.dtype
+                    )
+                    table_pos = table_pos + table_offset
+
+                table_root_pose = torch.cat([table_pos, table_quat], dim=-1)
+                table.write_root_pose_to_sim(table_root_pose, env_ids=env_ids)
+
+        anchor_pos_w_repeat = self.anchor_pos_w[:, None, :].repeat(1, len(self.cfg.body_names), 1)
+        anchor_quat_w_repeat = self.anchor_quat_w[:, None, :].repeat(1, len(self.cfg.body_names), 1)
+        robot_anchor_pos_w_repeat = self.robot_anchor_pos_w[:, None, :].repeat(
+            1, len(self.cfg.body_names), 1
+        )
+        robot_anchor_quat_w_repeat = self.robot_anchor_quat_w[:, None, :].repeat(
+            1, len(self.cfg.body_names), 1
+        )
+
+        delta_pos_w = robot_anchor_pos_w_repeat  # Root position of the robot
+        delta_pos_w[..., 2] = anchor_pos_w_repeat[..., 2]
+        delta_ori_w = torch_transform.get_heading_q(
+            quat_mul(robot_anchor_quat_w_repeat, quat_inv(anchor_quat_w_repeat))
+        )
+
+        self.body_quat_relative_w = quat_mul(delta_ori_w, self.body_quat_w)
+        self.body_pos_relative_w = delta_pos_w + quat_apply(
+            delta_ori_w, self.body_pos_w - anchor_pos_w_repeat
+        )
+
+    def _update_command(self):
+        """Advance the motion time cursor by one step and handle episode wrap-around.
+
+        Called once per simulation step. Updates adaptive sampling statistics,
+        increments the time cursor, resamples any environments that have reached
+        the end of their motion clip, updates the running reference root height
+        EMA, and optionally performs height-map raycasting for object-aware
+        observations.
+        """
+        if self.use_adaptive_sampling:
+            with common.Timer("update_adaptive_sampling"):
+                cur_time_steps = self.motion_start_time_steps + self.time_steps
+                self.motion_lib.update_adaptive_sampling(
+                    self._env.reset_terminated, self.motion_ids, cur_time_steps
+                )
+        self.time_steps += 1
+        env_ids = torch.where(
+            self.time_steps + self.motion_start_time_steps
+            >= self.motion_lib.get_time_step_total(self.motion_ids)
+        )[0]
+        self._resample_command(env_ids)
+
+        # Exponential moving average update for running_ref_root_height.
+        # ZL this should be moved to the recorders???
+        ema_alpha = 0.1  # Smoothing factor, adjust as needed
+        self.running_ref_root_height = (
+            ema_alpha * self.anchor_pos_w[:, 2] + (1 - ema_alpha) * self.running_ref_root_height
+        )
+
+        if self.cfg.use_height_map:
+            root_pos_w = self.robot.data.root_pos_w
+            root_quat_w = self.robot.data.root_quat_w
+
+            ray_starts_w = root_pos_w.unsqueeze(1).expand(-1, self.num_rays, -1)
+            root_quat_expanded = (
+                torch_transform.get_heading_q(root_quat_w)
+                .unsqueeze(1)
+                .expand(-1, self.num_rays, -1)
+            )
+            ray_dirs_w = quat_apply(
+                root_quat_expanded,
+                self.ray_dirs_local,
+            )
+
+            scan_dot_pos_w, _ = self.height_map.raycast_fused(
+                self.object_root_pos,
+                self.object_root_quat,
+                ray_starts_w,
+                ray_dirs_w,
+                n_mesh_per_cam=self.num_mesh_per_cam,
+                mesh_ids_flattened=self.mesh_ids,
+                cam_ids_flattened=self.cam_ids,
+                min_dist=0.0,
+                max_dist=self.cfg.height_map_max_dist,
+            )
+            # Adjust hits that fall below the ground plane so all zs are non-negative.
+            denom = ray_starts_w[:, :, 2] - scan_dot_pos_w[:, :, 2]
+            denom = torch.clamp_min(denom, 1e-8)
+            scale = torch.clamp(ray_starts_w[:, :, 2] / denom, max=1.0)
+            self.scan_dot_pos_w[:] = (
+                ray_starts_w + (scan_dot_pos_w - ray_starts_w) * scale.unsqueeze(-1)
+            ).view(self.num_envs, self.num_rays_x, self.num_rays_y, 3)
+
+    @property
+    def future_time_steps(self) -> torch.Tensor:
+        """Compute absolute time-step indices for all future reference frames.
+
+        Clamps to the last valid frame of each motion to avoid out-of-bounds access.
+
+        Returns:
+            Flattened tensor of shape ``(num_envs * num_future_frames,)``.
+        """
+        return (
+            torch.clip(
+                self.future_time_steps_init
+                + self.time_steps[:, None]
+                + self.motion_start_time_steps[:, None],
+                max=self.motion_num_steps[:, None] - 1,
+            )
+            .flatten()
+            .long()
+        )
+
+    @property
+    def smpl_future_time_steps(self) -> torch.Tensor:
+        """Compute absolute time-step indices for SMPL future reference frames.
+
+        SMPL future frames may use different count and spacing than robot frames.
+
+        Returns:
+            Flattened tensor of shape ``(num_envs * smpl_num_future_frames,)``.
+        """
+        return (
+            torch.clip(
+                self.smpl_future_time_steps_init
+                + self.time_steps[:, None]
+                + self.motion_start_time_steps[:, None],
+                max=self.motion_num_steps[:, None] - 1,
+            )
+            .flatten()
+            .long()
+        )
+
+    def _set_debug_vis_impl(self, debug_vis: bool):
+        """Create or toggle visibility of debug visualization markers.
+
+        Lazily initializes feet contact markers, height-map dot markers, and
+        contact center sphere markers on first enable. Subsequent calls toggle
+        visibility without re-creating prims.
+
+        Args:
+            debug_vis: Whether to enable or disable debug visualization.
+        """
+        if debug_vis:
+            if not hasattr(self, "goal_pos_visualizer"):
+                self.goal_pos_visualizer = VisualizationMarkers(
+                    self.cfg.body_pos_visualizer_cfg.replace(
+                        prim_path="/Visuals/goal_marker_sphere"
+                    )
+                )
+
+                self.feet_contact_goal_visualizers = []
+
+                for name in self.cfg.feet_body_names:
+                    self.feet_contact_goal_visualizers.append(
+                        VisualizationMarkers(
+                            self.cfg.feet_contact_visualizer_cfg.replace(
+                                prim_path="/Visuals/Command/goal/" + name
+                            )
+                        )
+                    )
+
+            self.goal_pos_visualizer.set_visibility(True)
+            for i in range(len(self.cfg.feet_body_names)):
+                self.feet_contact_goal_visualizers[i].set_visibility(True)
+
+            if self.cfg.use_height_map:
+                if not hasattr(self, "height_map_visualizer"):
+                    height_map_cfg = VisualizationMarkersCfg(
+                        prim_path="/Visuals/height_map",
+                        markers={
+                            "scan_dots": sim_utils.SphereCfg(
+                                radius=0.05,
+                                visual_material=sim_utils.PreviewSurfaceCfg(
+                                    diffuse_color=(1.0, 1.0, 0.0)
+                                ),
+                            ),
+                        },
+                    )
+                    self.height_map_visualizer = VisualizationMarkers(height_map_cfg)
+                self.height_map_visualizer.set_visibility(True)
+
+            # Contact center visualizers: deferred to _debug_vis_callback (lazy init)
+            # because motion_lib is not yet available during super().__init__()
+            if not hasattr(self, "contact_center_visualizers"):
+                self.contact_center_visualizers = None
+
+        else:
+            if hasattr(self, "goal_pos_visualizer"):
+                self.goal_pos_visualizer.set_visibility(False)
+            if hasattr(self, "feet_contact_goal_visualizers"):
+                for vis in self.feet_contact_goal_visualizers:
+                    vis.set_visibility(False)
+            if hasattr(self, "height_map_visualizer"):
+                self.height_map_visualizer.set_visibility(False)
+            if hasattr(self, "contact_center_visualizers") and self.contact_center_visualizers:
+                for vis in self.contact_center_visualizers.values():
+                    vis.set_visibility(False)
+
+    def _debug_vis_callback(self, event):  # noqa: ARG002
+        """Update debug visualization marker positions each render frame.
+
+        Draws current robot body frames, reference (goal) body frames, feet
+        contact indicators, height-map hit points, and object contact centers.
+
+        Args:
+            event: Render event from the simulation (unused).
+        """
+        if not self.robot.is_initialized:
+            return
+
+        if not hasattr(self, "goal_pos_visualizer"):
+            return
+
+        self.goal_pos_visualizer.visualize(self.body_pos_w.view(-1, 3))
+
+        if hasattr(self, "feet_contact_goal_visualizers"):
+            for i in range(len(self.cfg.body_names)):
+                if self.cfg.body_names[i] == "left_ankle_roll_link":
+                    self.feet_contact_goal_visualizers[0].visualize(
+                        translations=self.body_pos_relative_w[:, i],
+                        marker_indices=self.feet_l.int().reshape(-1),
+                    )
+                if self.cfg.body_names[i] == "right_ankle_roll_link":
+                    self.feet_contact_goal_visualizers[1].visualize(
+                        translations=self.body_pos_relative_w[:, i],
+                        marker_indices=self.feet_r.int().reshape(-1),
+                    )
+
+        if self.cfg.use_height_map:
+            self.height_map_visualizer.visualize(
+                translations=self.scan_dot_pos_w.view(-1, 3),
+            )
+
+        # Contact center visualization (lazy init on first callback)
+        if hasattr(self, "contact_center_visualizers") and self.contact_center_visualizers is None:  # noqa: SIM102
+            if hasattr(self, "motion_lib"):
+                self.contact_center_visualizers = {}
+                if hasattr(self.motion_lib, "_motion_object_contact_center_left"):
+                    left_cfg = VisualizationMarkersCfg(
+                        prim_path="/Visuals/Command/contact_center_left",
+                        markers={
+                            "contact": sim_utils.SphereCfg(
+                                radius=0.05,
+                                visual_material=sim_utils.PreviewSurfaceCfg(
+                                    diffuse_color=(0.0, 0.0, 1.0),
+                                ),
+                            ),
+                        },
+                    )
+                    self.contact_center_visualizers["left_hand"] = VisualizationMarkers(left_cfg)
+                if hasattr(self.motion_lib, "_motion_object_contact_center_right"):
+                    right_cfg = VisualizationMarkersCfg(
+                        prim_path="/Visuals/Command/contact_center_right",
+                        markers={
+                            "contact": sim_utils.SphereCfg(
+                                radius=0.05,
+                                visual_material=sim_utils.PreviewSurfaceCfg(
+                                    diffuse_color=(0.0, 1.0, 1.0),
+                                ),
+                            ),
+                        },
+                    )
+                    self.contact_center_visualizers["right_hand"] = VisualizationMarkers(right_cfg)
+                for vis in self.contact_center_visualizers.values():
+                    vis.set_visibility(True)
+
+        if self.contact_center_visualizers:
+            hidden_pos = torch.tensor([[0.0, 0.0, -1000.0]], device=self._env.device)
+            for hand, visualizer in self.contact_center_visualizers.items():
+                contact_center = self.motion_lib.get_object_contact_center(
+                    self.motion_ids, self.motion_start_time_steps + self.time_steps, hand=hand
+                )
+                if contact_center is None:
+                    visualizer.visualize(translations=hidden_pos)
+                    continue
+                valid_mask = torch.norm(contact_center, dim=-1) > 1e-6
+                world_center = self._get_contact_center_world(hand)
+                world_center[~valid_mask] = hidden_pos
+                visualizer.visualize(translations=world_center)
+
+    def set_motion_state(self, motion_ids, time_steps, motion_start_time_steps=None):
+        """Update which motion clip/frame this command serves (for offline use)."""
+        self.motion_ids = motion_ids
+        self.time_steps = time_steps
+        self.motion_start_time_steps = (
+            motion_start_time_steps
+            if motion_start_time_steps is not None
+            else torch.zeros_like(motion_ids)
+        )
+        self.future_motion_ids = motion_ids.repeat_interleave(self.num_future_frames)
+        self.smpl_future_motion_ids = motion_ids.repeat_interleave(self.smpl_num_future_frames)
+        self.motion_num_steps = self.motion_lib.get_time_step_total(motion_ids)
+
+
+class ForceTrackingCommand(CommandTerm):
+    """Apply external perturbation forces and manage compliance state for training robustness.
+
+    This command term works alongside ``TrackingCommand`` to add domain-randomized
+    external forces on specified robot bodies (typically wrists and torso). It also
+    tracks per-env compliance levels that control how stiffly the end-effectors
+    track their targets, enabling the policy to learn compliant manipulation
+    behaviors.
+
+    Key responsibilities:
+        - Maintain per-body force direction and magnitude buffers that event terms
+          sample into periodically (every ``force_update_frequency`` steps).
+        - Track per-env compliance levels for left wrist, right wrist, and head
+          via ``eef_stiffness_buf``.
+        - Compute Jacobian-based end-effector analysis for compliance calculations.
+        - Report detailed force and compliance metrics per encoder type
+          (G1/teleop/SMPL) for W&B logging.
+
+    NOTE: Force application itself is handled by event terms that read and write
+    this command's buffers. This command only manages the state and metrics.
+    """
+
+    cfg: ForceTrackingCommandCfg
+
+    def __init__(self, cfg: ForceTrackingCommandCfg, env: ManagerBasedRLEnv):
+        """Initialize force tracking state, Jacobian indices, and compliance buffers.
+
+        Args:
+            cfg: Configuration specifying force bodies, max force, update frequency,
+                joint dependencies for Jacobian computation, and debug settings.
+            env: The manager-based RL environment that owns this command.
+        """
+        super().__init__(cfg, env)
+
+        self.is_evaluating = False
+        self.robot: Articulation = env.scene[cfg.asset_name]
+        self.robot_anchor_body_index = self.robot.body_names.index(self.cfg.anchor_body)
+        self.motion_anchor_body_index = self.cfg.body_names.index(self.cfg.anchor_body)
+
+        self.left_eef_deps_ids = np.array(
+            [self.robot.joint_names.index(joint) for joint in self.cfg.left_eef_deps]
+        )
+        self.right_eef_deps_ids = np.array(
+            [self.robot.joint_names.index(joint) for joint in self.cfg.right_eef_deps]
+        )
+        self.waist_deps_ids = np.array(
+            [self.robot.joint_names.index(joint) for joint in self.cfg.waist_joints]
+        )
+        self.upper_body_deps_ids = np.concatenate(
+            [self.waist_deps_ids, self.left_eef_deps_ids, self.right_eef_deps_ids]
+        )
+
+        self.kp = self.robot.data.joint_stiffness.clone()
+        self.kd = self.robot.data.joint_damping.clone()
+
+        self.force_update_frequency = self.cfg.force_update_frequency
+        self.max_force = self.cfg.max_force
+        self.vr_3point_body_indices = [
+            self.robot.body_names.index(name) for name in self.cfg.vr_3point_body
+        ]
+        self.vr_3point_body_indices_motion = [
+            self.cfg.body_names.index(name) for name in self.cfg.vr_3point_body
+        ]
+        self.vr_3point_body_offsets = (
+            torch.tensor(self.cfg.vr_3point_body_offset, dtype=torch.float32, device=self.device)
+            .view(1, -1, 3)
+            .repeat(self.num_envs, 1, 1)
+        )
+        self.body_indexes = torch.tensor(
+            self.robot.find_bodies(self.cfg.body_names, preserve_order=True)[0],
+            dtype=torch.long,
+            device=self.device,
+        )
+
+        ### Force related
+        self.num_bodies = len(self.cfg.body_names)
+        self.body_force_dir_buf = torch.randn(
+            self.num_envs,
+            self.num_bodies,
+            3,
+            dtype=torch.float,
+            device=self.device,
+            requires_grad=False,
+        )
+        self.body_force_dir_buf /= torch.norm(
+            self.body_force_dir_buf, dim=-1, keepdim=True
+        )  # normalize
+
+        # NOTE: We initialize force_push_ids first so we can use len(force_push_ids) for buffer shape
+        self.force_push_ids = self.robot.find_bodies(self.cfg.force_push_body, preserve_order=True)[
+            0
+        ]
+        self.num_force_push_bodies = len(self.force_push_ids)
+
+        # Per-body force magnitude buffer: [num_envs, num_force_push_bodies]
+        # Each body can have different force magnitudes, enabling differentiated
+        # force application (e.g., stronger forces on wrists than torso)
+        self.body_force_magnitude_buf = torch.rand(
+            self.num_envs,
+            self.num_force_push_bodies,
+            dtype=torch.float,
+            device=self.device,
+            requires_grad=False,
+        )  # [0, 1] per body
+
+        self.force_push_counter = torch.zeros(self.num_envs, dtype=torch.int, device=self.device)
+        self.force_duration_per_env = torch.zeros(
+            self.num_envs, dtype=torch.int, device=self.device
+        )
+        self.force_config_init = False
+        self.non_force_push_ids_rel = []
+        self.force_push_ids_rel = []
+        for i, idx in enumerate(self.body_indexes.tolist()):
+            if idx not in self.force_push_ids:
+                self.non_force_push_ids_rel.append(i)
+            else:
+                self.force_push_ids_rel.append(i)
+        # self.non_force_push_ids = [i for i in self.body_indexes.tolist() if i not in self.force_push_ids]
+        self.force_push_body_offsets = (
+            torch.tensor(self.cfg.force_push_body_offset, dtype=torch.float32, device=self.device)
+            .view(1, -1, 3)
+            .repeat(self.num_envs, 1, 1)
+        )
+        self.last_force_applied = torch.zeros(
+            self.num_envs,
+            len(self.force_push_ids),
+            3,
+            dtype=torch.float,
+            device=self.device,
+            requires_grad=False,
+        )
+
+        # compliance related counters
+        self.compliance_counter = torch.zeros(self.num_envs, dtype=torch.int, device=self.device)
+        self.compliance_duration_per_env = torch.zeros(
+            self.num_envs, dtype=torch.int, device=self.device
+        )
+        self.eef_stiffness_buf = torch.zeros(
+            self.num_envs, 3, dtype=torch.float32, device=self.device
+        )
+        self.compliance_config_init = False
+
+        # Compliance monitoring - track cumulative stats for sanity checks
+        self._force_update_count = 0  # Tracks how many times force was applied (non-zero)
+        self._compliance_update_count = 0  # Tracks how many times compliance was updated
+        self._total_steps = 0
+        self._warned_no_force = False  # Prevent spamming warnings
+        # Debug print frequency (0 = disabled, nonzero = print every N steps)
+        # Can be set via config: manager_env.commands.force.debug_print_every_n_steps=10
+        self._debug_print_every_n_steps = self.cfg.debug_print_every_n_steps
+        # Note: Metrics are now created dynamically in _update_metrics()
+        # The old "force applied" metric (with space) has been removed to avoid
+        # confusion with "force_applied" (with underscore)
+
+    def set_is_evaluating(self, is_evaluating: bool):
+        """Toggle evaluation mode."""
+        self.is_evaluating = is_evaluating
+
+    @property
+    def jacobian(self) -> torch.Tensor:
+        """Return the full articulation Jacobian from PhysX.
+
+        Returns:
+            Tensor of shape ``(num_envs, num_bodies, 6, num_dof + 6)``.
+        """
+        return self.robot.root_physx_view.get_jacobians()
+
+    @property
+    def eef_jacobian(self) -> torch.Tensor:
+        """Return the translational Jacobian for left and right end-effectors.
+
+        Returns:
+            Tensor of shape ``(num_envs, 2, 3, num_dof + 6)``.
+        """
+        return self.jacobian[:, self.vr_3point_body_indices[:2], :3, :]
+
+    @property
+    def matrix_M(self) -> torch.Tensor:
+        """Compute the combined upper-body Jacobian mapping for compliance control.
+
+        Assembles a ``(6, 17)`` matrix per env that maps the 17 upper-body joint
+        velocities (3 waist + 7 left arm + 7 right arm) to 6D end-effector
+        velocities (3 left + 3 right translational).
+
+        Returns:
+            Tensor of shape ``(num_envs, 6, 17)``.
+        """
+        eef_jacobian = self.eef_jacobian
+        left_eef_jac = eef_jacobian[:, 0, :, self.left_eef_deps_ids + 6]
+        right_eef_jac = eef_jacobian[:, 1, :, self.right_eef_deps_ids + 6]
+        waist_jac_left = eef_jacobian[:, 0, :, self.waist_deps_ids + 6]
+        waist_jac_right = eef_jacobian[:, 1, :, self.waist_deps_ids + 6]
+
+        M = torch.zeros(self.num_envs, 6, 17).to(self.device)
+        M[:, :3, :3] = waist_jac_left
+        M[:, 3:, :3] = waist_jac_right
+        M[:, :3, 3:10] = left_eef_jac
+        M[:, 3:, 10:] = right_eef_jac
+        return M
+
+    def _resample_command(self, env_ids: Sequence[int]):
+        """Reset force-related buffers when environments are reset.
+
+        This prevents stale force values from affecting compliance calculations
+        immediately after environment reset.
+
+        NOTE: We intentionally do NOT reset force_push_counter here.
+        The counter needs to reach force_update_frequency (default: 100) before
+        forces are applied. If we reset it on every env reset, and episodes are
+        shorter than 100 steps, forces would never be applied!
+        """
+        if len(env_ids) > 0:
+            self.last_force_applied[env_ids] = 0.0
+
+    def _update_command(self):
+        """No-op; force state is updated by event terms, not per-step."""
+        pass
+
+    def _update_compliance_force_push_related_metrics(self):
+        """Compute and log per-body force magnitudes, compliance levels, and encoder ratios.
+
+        Populates ``self.metrics`` with force norms per body, compliance levels
+        per encoder type (G1/teleop/SMPL), stiff vs. compliant environment ratios,
+        and cross-referenced force-by-compliance-status metrics. Optionally prints
+        debug summaries every ``_debug_print_every_n_steps`` steps.
+        """
+        self._total_steps += 1
+
+        # =====================================================================
+        # Get encoder masks from TrackingCommand (for per-encoder metrics)
+        # encoder_index: [num_envs, num_encoders] one-hot encoding
+        # =====================================================================
+        motion_command = None
+        encoder_g1_mask = None
+        encoder_teleop_mask = None
+        encoder_smpl_mask = None
+        has_encoder_info = False
+
+        try:
+            motion_command = self._env.command_manager.get_term("motion")
+            if hasattr(motion_command, "encoder_index") and hasattr(
+                motion_command, "encoder_sample_probs_dict"
+            ):
+                encoder_names = list(motion_command.encoder_sample_probs_dict.keys())
+                encoder_index = motion_command.encoder_index  # [num_envs, num_encoders]
+
+                # Get masks for each encoder type
+                if "g1" in encoder_names:
+                    g1_idx = encoder_names.index("g1")
+                    encoder_g1_mask = encoder_index[:, g1_idx].bool()
+                if "teleop" in encoder_names:
+                    teleop_idx = encoder_names.index("teleop")
+                    encoder_teleop_mask = encoder_index[:, teleop_idx].bool()
+                if "smpl" in encoder_names:
+                    smpl_idx = encoder_names.index("smpl")
+                    encoder_smpl_mask = encoder_index[:, smpl_idx].bool()
+
+                has_encoder_info = True
+        except Exception:  # noqa: BLE001, S110
+            pass  # Motion command not available, skip encoder-specific metrics
+
+        # =====================================================================
+        # Per-body force metrics
+        # last_force_applied: [num_envs, num_force_push_bodies, 3]
+        # force_push_body order: ["left_wrist_yaw_link", "right_wrist_yaw_link", "torso_link"]
+        # =====================================================================
+        force_norm = torch.norm(self.last_force_applied, dim=-1)  # [num_envs, num_bodies]
+        num_force_bodies = self.last_force_applied.shape[1]
+
+        # Per-body force magnitudes (for wrists and torso)
+        if num_force_bodies >= 1:
+            self.metrics["force_left_wrist"] = force_norm[:, 0]
+        if num_force_bodies >= 2:
+            self.metrics["force_right_wrist"] = force_norm[:, 1]
+        if num_force_bodies >= 3:
+            self.metrics["force_torso"] = force_norm[:, 2]
+
+        # Combined wrist force
+        if num_force_bodies >= 2:
+            self.metrics["force_wrists_sum"] = force_norm[:, 0] + force_norm[:, 1]
+            self.metrics["force_wrists_max"] = torch.max(force_norm[:, 0], force_norm[:, 1])
+
+        # Overall force stats
+        self.metrics["force_applied"] = force_norm.mean(dim=-1)
+        self.metrics["force_applied_max"] = force_norm.max(dim=-1).values
+
+        # Track if any force was actually applied
+        force_nonzero = (force_norm.sum() > 0.01).float()
+        if force_nonzero > 0:
+            self._force_update_count += 1
+        self.metrics["force_nonzero_ratio"] = torch.tensor(
+            self._force_update_count / max(1, self._total_steps), device=self.device
+        ).expand(self.num_envs)
+
+        # =====================================================================
+        # Compliance distribution masks
+        # eef_stiffness_buf: [num_envs, 3] -> [left_wrist, right_wrist, head]
+        # =====================================================================
+        compliance_threshold = 0.001  # Threshold to consider compliance as "active"
+        is_compliant = self.eef_stiffness_buf[:, :2].abs().sum(dim=-1) > compliance_threshold
+        is_stiff = ~is_compliant
+
+        # =====================================================================
+        # ENCODER RATIO METRICS
+        # ratio_of_total_envs/encoder_*: Fraction of envs using each encoder
+        # =====================================================================
+        if has_encoder_info:
+            if encoder_g1_mask is not None:
+                n_g1 = encoder_g1_mask.sum().float()
+                self.metrics["ratio_of_total_envs/encoder_g1"] = (n_g1 / self.num_envs).expand(
+                    self.num_envs
+                )
+            if encoder_teleop_mask is not None:
+                n_teleop = encoder_teleop_mask.sum().float()
+                self.metrics["ratio_of_total_envs/encoder_teleop"] = (
+                    n_teleop / self.num_envs
+                ).expand(self.num_envs)
+            if encoder_smpl_mask is not None:
+                n_smpl = encoder_smpl_mask.sum().float()
+                self.metrics["ratio_of_total_envs/encoder_smpl"] = (n_smpl / self.num_envs).expand(
+                    self.num_envs
+                )
+
+        # =====================================================================
+        # COMPLIANCE LEVEL METRICS PER ENCODER
+        # compliance_level_LH/encoder_*: Mean compliance for each encoder type
+        # =====================================================================
+        if has_encoder_info:
+            # G1 encoder (should always be stiff = 0)
+            if encoder_g1_mask is not None and encoder_g1_mask.sum() > 0:
+                g1_compliance = self.eef_stiffness_buf[encoder_g1_mask]
+                self.metrics["compliance_level_LH/encoder_g1"] = (
+                    g1_compliance[:, 0].mean().expand(self.num_envs)
+                )
+                self.metrics["compliance_level_RH/encoder_g1"] = (
+                    g1_compliance[:, 1].mean().expand(self.num_envs)
+                )
+                self.metrics["compliance_level_Head/encoder_g1"] = (
+                    g1_compliance[:, 2].mean().expand(self.num_envs)
+                )
+            else:
+                self.metrics["compliance_level_LH/encoder_g1"] = torch.zeros(
+                    self.num_envs, device=self.device
+                )
+                self.metrics["compliance_level_RH/encoder_g1"] = torch.zeros(
+                    self.num_envs, device=self.device
+                )
+                self.metrics["compliance_level_Head/encoder_g1"] = torch.zeros(
+                    self.num_envs, device=self.device
+                )
+
+            # Teleop encoder
+            if encoder_teleop_mask is not None and encoder_teleop_mask.sum() > 0:
+                teleop_compliance = self.eef_stiffness_buf[encoder_teleop_mask]
+                self.metrics["compliance_level_LH/encoder_teleop"] = (
+                    teleop_compliance[:, 0].mean().expand(self.num_envs)
+                )
+                self.metrics["compliance_level_RH/encoder_teleop"] = (
+                    teleop_compliance[:, 1].mean().expand(self.num_envs)
+                )
+                self.metrics["compliance_level_Head/encoder_teleop"] = (
+                    teleop_compliance[:, 2].mean().expand(self.num_envs)
+                )
+            else:
+                self.metrics["compliance_level_LH/encoder_teleop"] = torch.zeros(
+                    self.num_envs, device=self.device
+                )
+                self.metrics["compliance_level_RH/encoder_teleop"] = torch.zeros(
+                    self.num_envs, device=self.device
+                )
+                self.metrics["compliance_level_Head/encoder_teleop"] = torch.zeros(
+                    self.num_envs, device=self.device
+                )
+
+            # SMPL encoder
+            if encoder_smpl_mask is not None and encoder_smpl_mask.sum() > 0:
+                smpl_compliance = self.eef_stiffness_buf[encoder_smpl_mask]
+                self.metrics["compliance_level_LH/encoder_smpl"] = (
+                    smpl_compliance[:, 0].mean().expand(self.num_envs)
+                )
+                self.metrics["compliance_level_RH/encoder_smpl"] = (
+                    smpl_compliance[:, 1].mean().expand(self.num_envs)
+                )
+                self.metrics["compliance_level_Head/encoder_smpl"] = (
+                    smpl_compliance[:, 2].mean().expand(self.num_envs)
+                )
+            else:
+                self.metrics["compliance_level_LH/encoder_smpl"] = torch.zeros(
+                    self.num_envs, device=self.device
+                )
+                self.metrics["compliance_level_RH/encoder_smpl"] = torch.zeros(
+                    self.num_envs, device=self.device
+                )
+                self.metrics["compliance_level_Head/encoder_smpl"] = torch.zeros(
+                    self.num_envs, device=self.device
+                )
+
+        # =====================================================================
+        # FORCE METRICS BY COMPLIANCE STATUS
+        # compliance_related_force_on_LH/nonzero_compliance_envs: Force on envs with compliance > threshold
+        # compliance_related_force_on_LH/stiff_envs: Force on stiff envs (G1 + stiff teleop/smpl)
+        # =====================================================================
+        if num_force_bodies >= 1:
+            # Forces on envs with nonzero compliance (teleop/SMPL with active compliance)
+            if is_compliant.sum() > 0:
+                force_compliant = force_norm[is_compliant]
+                self.metrics["compliance_related_force_on_LH/nonzero_compliance_envs"] = (
+                    force_compliant[:, 0].mean().expand(self.num_envs)
+                )
+                if num_force_bodies >= 2:
+                    self.metrics["compliance_related_force_on_RH/nonzero_compliance_envs"] = (
+                        force_compliant[:, 1].mean().expand(self.num_envs)
+                    )
+                if num_force_bodies >= 3:
+                    self.metrics["compliance_related_force_on_Head/nonzero_compliance_envs"] = (
+                        force_compliant[:, 2].mean().expand(self.num_envs)
+                    )
+            else:
+                self.metrics["compliance_related_force_on_LH/nonzero_compliance_envs"] = (
+                    torch.zeros(self.num_envs, device=self.device)
+                )
+                if num_force_bodies >= 2:
+                    self.metrics["compliance_related_force_on_RH/nonzero_compliance_envs"] = (
+                        torch.zeros(self.num_envs, device=self.device)
+                    )
+                if num_force_bodies >= 3:
+                    self.metrics["compliance_related_force_on_Head/nonzero_compliance_envs"] = (
+                        torch.zeros(self.num_envs, device=self.device)
+                    )
+
+            # Forces on stiff envs (G1 envs + teleop/smpl envs with stiff mode)
+            if is_stiff.sum() > 0:
+                force_stiff = force_norm[is_stiff]
+                self.metrics["compliance_related_force_on_LH/stiff_envs"] = (
+                    force_stiff[:, 0].mean().expand(self.num_envs)
+                )
+                if num_force_bodies >= 2:
+                    self.metrics["compliance_related_force_on_RH/stiff_envs"] = (
+                        force_stiff[:, 1].mean().expand(self.num_envs)
+                    )
+                if num_force_bodies >= 3:
+                    self.metrics["compliance_related_force_on_Head/stiff_envs"] = (
+                        force_stiff[:, 2].mean().expand(self.num_envs)
+                    )
+            else:
+                self.metrics["compliance_related_force_on_LH/stiff_envs"] = torch.zeros(
+                    self.num_envs, device=self.device
+                )
+                if num_force_bodies >= 2:
+                    self.metrics["compliance_related_force_on_RH/stiff_envs"] = torch.zeros(
+                        self.num_envs, device=self.device
+                    )
+                if num_force_bodies >= 3:
+                    self.metrics["compliance_related_force_on_Head/stiff_envs"] = torch.zeros(
+                        self.num_envs, device=self.device
+                    )
+
+        # =====================================================================
+        # Debug print every N steps
+        # =====================================================================
+        if (
+            self._debug_print_every_n_steps > 0
+            and self._total_steps % self._debug_print_every_n_steps == 0
+        ):
+            force_mean = force_norm.mean(dim=0)
+            force_vals = [f"{force_mean[i].item():.2f}" for i in range(min(3, num_force_bodies))]
+            force_str = ", ".join(force_vals) if force_vals else "N/A"
+            body_labels = ["L_wrist", "R_wrist", "Torso"][:num_force_bodies]
+
+            n_compliant = is_compliant.sum().item()
+            n_stiff = is_stiff.sum().item()
+
+            # Per-encoder counts
+            encoder_str = ""
+            if has_encoder_info:
+                n_g1 = encoder_g1_mask.sum().item() if encoder_g1_mask is not None else 0
+                n_teleop = (
+                    encoder_teleop_mask.sum().item() if encoder_teleop_mask is not None else 0
+                )
+                n_smpl = encoder_smpl_mask.sum().item() if encoder_smpl_mask is not None else 0
+                encoder_str = f" | Encoders(G1/Teleop/SMPL): {n_g1}/{n_teleop}/{n_smpl}"
+
+            # Compliance stats for compliant envs only
+            if n_compliant > 0:
+                compliance_compliant_mean = self.eef_stiffness_buf[is_compliant].mean(dim=0)
+                comp_str = f"[{compliance_compliant_mean[0].item():.3f}, {compliance_compliant_mean[1].item():.3f}, {compliance_compliant_mean[2].item():.3f}]"  # noqa: E501
+            else:
+                comp_str = "[N/A - all stiff]"
+
+            # Debug info about event calls
+            event_calls = getattr(self, "_event_call_count", 0)
+            counter_stats = f"counter_min={self.force_push_counter.min().item()} max={self.force_push_counter.max().item()}"  # noqa: E501
+
+            print(  # noqa: T201
+                f"[ForceDebug @ step {self._total_steps}] "
+                f"Force({'/'.join(body_labels)}): [{force_str}] | "
+                f"Compliant: {n_compliant} Stiff: {n_stiff}{encoder_str} | "
+                f"Compliance(L/R/H): {comp_str} | "
+                f"Force ratio: {self._force_update_count}/{self._total_steps} | "
+                f"Event calls: {event_calls} | {counter_stats}"
+            )
+
+        # =====================================================================
+        # Legacy compliance metrics (kept for backward compatibility)
+        # =====================================================================
+        self.metrics["compliance_left_wrist"] = self.eef_stiffness_buf[:, 0]
+        self.metrics["compliance_right_wrist"] = self.eef_stiffness_buf[:, 1]
+        self.metrics["compliance_head"] = self.eef_stiffness_buf[:, 2]
+        self.metrics["compliance_mean"] = self.eef_stiffness_buf.mean(dim=-1)
+
+        # =====================================================================
+        # Compliance ratio metrics FOR NON-G1 ENCODER ENVS ONLY
+        # G1 envs are always forced to stiff, so including them is misleading.
+        # These metrics show the stiff/compliant distribution among TELEOP/SMPL envs.
+        # =====================================================================
+        if has_encoder_info and encoder_g1_mask is not None:
+            non_g1_mask = ~encoder_g1_mask
+            n_non_g1 = non_g1_mask.sum()
+            if n_non_g1 > 0:
+                stiff_ratio_non_g1 = is_stiff[non_g1_mask].float().mean()
+                self.metrics["stiff_ratio_for_non_G1_encoder_envs"] = stiff_ratio_non_g1.expand(
+                    self.num_envs
+                )
+                self.metrics["compliant_ratio_for_non_G1_encoder_envs"] = (
+                    1.0 - stiff_ratio_non_g1
+                ).expand(self.num_envs)
+            else:
+                # All envs are G1, ratio is undefined (report 0)
+                self.metrics["stiff_ratio_for_non_G1_encoder_envs"] = torch.zeros(
+                    self.num_envs, device=self.device
+                )
+                self.metrics["compliant_ratio_for_non_G1_encoder_envs"] = torch.zeros(
+                    self.num_envs, device=self.device
+                )
+        else:
+            # No encoder info available, fall back to global ratio
+            stiff_ratio = is_stiff.float().mean()
+            self.metrics["stiff_ratio_for_non_G1_encoder_envs"] = stiff_ratio.expand(self.num_envs)
+            self.metrics["compliant_ratio_for_non_G1_encoder_envs"] = (1.0 - stiff_ratio).expand(
+                self.num_envs
+            )
+
+        # Sanity check warning: If compliance is being varied but no forces are applied
+        compliance_is_active = self.compliance_config_init and (
+            self.eef_stiffness_buf.abs().sum() > 0
+        )
+        if compliance_is_active and self._total_steps > 1000:  # noqa: SIM102
+            if self._force_update_count == 0 and not self._warned_no_force:
+                import warnings
+
+                warnings.warn(  # noqa: B028
+                    "\n" + "=" * 80 + "\n"
+                    "[COMPLIANCE SANITY CHECK FAILED]\n"
+                    "Compliance values are being changed, but NO external forces have been applied!\n"
+                    "This means vr_3point_local_target_compliant == vr_3point_local_target always.\n"
+                    "The policy will NOT learn to use the compliance signal.\n\n"
+                    "FIX: Add 'compliance_force_push' event to your events config:\n"
+                    "  defaults:\n"
+                    "    - terms/compliance_force_push@_here_\n"
+                    "=" * 80
+                )
+                self._warned_no_force = True
+
+    def _update_metrics(self):
+        """Refresh max_force from config and compute all force/compliance metrics."""
+        self.max_force = self.cfg.max_force
+        self._update_compliance_force_push_related_metrics()
+
+    @property
+    def command(self):
+        """Return None; ForceTrackingCommand has no direct observation tensor."""
+        return None
+
+
+@configclass
+class TrackingCommandCfg(CommandTermCfg):
+    """Configuration for TrackingCommand motion tracking.
+
+    Controls motion library loading, future reference frame layout, encoder
+    sampling probabilities (G1/SMPL/teleop), episode initialization strategy,
+    object handling, and debug visualization marker styles.
+    """
+
+    class_type: type = TrackingCommand
+
+    asset_name: str = dataclasses.MISSING
+
+    motion_lib_cfg: dict = None
+    motion_file: str = None
+    smpl_motion_file: str = None
+    filter_motion_keys: list[str] = None
+    use_paired_motions: bool = False
+    # Contact-based initialization: sample timestamps before first contact frame
+    # Path to contact pickle file (e.g., data/motion_lib_grab/contact/s2_apple_lift.pkl)
+    contact_file: str = None
+    # If True, sample only timestamps before the first contact frame
+    sample_before_contact: bool = False
+    # Margin (in frames) before first contact to sample from
+    sample_before_contact_margin: int = 10
+    # Which hand's in_contact label to use for deriving first contact frame
+    sample_before_contact_hand: str = "right_hand"
+    anchor_body: str = dataclasses.MISSING
+    body_names: list[str] = dataclasses.MISSING
+
+    vr_3point_body: list[str] = []  # noqa: RUF012
+    vr_3point_body_offset: list[list[float]] = []  # noqa: RUF012
+
+    reward_point_body: list[str] = []  # noqa: RUF012
+    reward_point_body_offset: list[list[float]] = []  # noqa: RUF012
+
+    # For backward compatibility (to remove)
+    force_push_body: list[str] = []  # noqa: RUF012
+    force_push_body_offset: list[list[float]] = []  # noqa: RUF012
+
+    num_future_frames: int = 1
+    dt_future_ref_frames: float = 0.1
+    randomize_heading: bool = False
+
+    # Variable frame support: when enabled, num_future_frames serves as max_frames
+    # and each environment/sample gets a random num_frames from
+    # [variable_frames_min, num_future_frames] with step variable_frames_step.
+    # Step must be a power of 2 compatible with down_t (step=4 works for down_t=2).
+    variable_frames_enabled: bool = False
+    variable_frames_min: int = 16
+    variable_frames_step: int = 4
+
+    freeze_frame_aug: bool = False
+    freeze_frame_aug_prob: float = 0.1
+
+    smpl_num_future_frames: int = None
+    smpl_dt_future_ref_frames: float = None
+
+    smpl_num_future_frames: int = None  # noqa: PIE794
+    smpl_dt_future_ref_frames: float = None  # noqa: PIE794
+
+    encoder_sample_probs: dict[str, float] = None
+
+    # ==========================================================================
+    # CHIP Compliance Training Optimization Flag
+    # ==========================================================================
+    # When True, enables cleaner encoder logic for compliance-aware training:
+    #   1. SMPL-native envs do NOT automatically activate G1 encoder in main loop
+    #   2. G1 encoder only runs for G1-native envs (policy tokens)
+    #   3. G1 latents for SMPL-native envs are computed ONLY in aux losses
+    #      and ONLY when compliance=0 (stiff mode)
+    #
+    # This eliminates wasted computation where G1 tokens were computed for
+    # SMPL-native envs but immediately overwritten by SMPL tokens.
+    #
+    # Set to False (default) for backward compatibility with non-compliance runs.
+    # ==========================================================================
+    optimize_encoders_ratio_for_CHIP: bool = False
+
+    # Probability to also sample teleop mode when smpl mode is active
+    # This enables teleop-smpl and g1-teleop latent alignment losses
+    teleop_sample_prob_when_smpl: float = 0.0
+
+    # Always start from the first frame of the motion file during resampling
+    # Useful for debugging and replaying specific motions from the beginning
+    start_from_first_frame: bool = False
+
+    # Sample each motion at most once (no duplicates across environments)
+    # Requires num_envs <= num_available_motions, otherwise will error
+    # Useful for replay/evaluation to ensure coverage of all unique motions
+    sample_unique_motions: bool = False
+
+    # Sample from the first N frames of the motion (random uniform in [0, N-1])
+    # If set, this takes precedence over start_from_first_frame
+    # Useful for adding slight variation while still starting near the beginning
+    sample_from_n_initial_frames: int = None
+
+    # Object position randomization (adds random offset to object position at reset)
+    object_position_randomize: bool = False
+    # Randomization range for each axis: {"x": 0.05, "y": 0.05, "z": 0.0}
+    # Values are half-range, so 0.05 means uniform random in [-0.05, 0.05]
+    object_position_randomization: dict[str, float] = None
+
+    # Table position offset: [x, y, z] added to table position from meta file
+    # Useful for adjusting table position relative to the robot
+    table_offset: list[float] = None
+
+    pose_range: dict[str, tuple[float, float]] = {}  # noqa: RUF012
+    velocity_range: dict[str, tuple[float, float]] = {}  # noqa: RUF012
+
+    joint_position_range: tuple[float, float] = (-0.52, 0.52)
+    joint_velocity_range: tuple[float, float] = (-0, 0)
+
+    body_pos_visualizer_cfg: VisualizationMarkersCfg = DEFORMABLE_TARGET_MARKER_CFG.replace(
+        prim_path="/Visuals/goal_marker_sphere"
+    )
+    body_pos_visualizer_cfg.markers["target"].radius = 0.05
+    body_pos_visualizer_cfg.markers["target"].visual_material = sim_utils.PreviewSurfaceCfg(
+        diffuse_color=(1.0, 1.0, 0.0)
+    )
+
+    feet_contact_visualizer_cfg: VisualizationMarkersCfg = DEFORMABLE_TARGET_MARKER_CFG.replace(
+        prim_path="/Visuals/goal_marker_sphere"
+    )
+    feet_contact_visualizer_cfg.markers = {  # noqa: RUF012
+        "target_small": sim_utils.SphereCfg(
+            radius=0.0001,
+            visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0)),
+        ),
+        "target_big": sim_utils.SphereCfg(
+            radius=0.04,
+            visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0)),
+        ),
+    }
+
+    waist_joints = ["waist_yaw_joint", "waist_roll_joint", "waist_pitch_joint"]  # noqa: RUF012
+
+    left_eef_deps = [  # noqa: RUF012
+        "left_shoulder_pitch_joint",
+        "left_shoulder_roll_joint",
+        "left_shoulder_yaw_joint",
+        "left_elbow_joint",
+        "left_wrist_pitch_joint",
+        "left_wrist_roll_joint",
+        "left_wrist_yaw_joint",
+    ]
+
+    right_eef_deps = [  # noqa: RUF012
+        "right_shoulder_pitch_joint",
+        "right_shoulder_roll_joint",
+        "right_shoulder_yaw_joint",
+        "right_elbow_joint",
+        "right_wrist_pitch_joint",
+        "right_wrist_roll_joint",
+        "right_wrist_yaw_joint",
+    ]
+
+    force_push_body: list[str] = []  # default only two wrists  # noqa: PIE794, RUF012
+    force_push_body_offset: list[list[float]] = []  # noqa: PIE794, RUF012
+
+    feet_body_names = ["left_ankle_roll_link", "right_ankle_roll_link"]  # noqa: RUF012
+    cat_upper_body_poses: bool = False
+    cat_upper_body_poses_prob: float = 0.5
+    randomize_wrist_poses: bool = False
+    randomize_wrist_prob: float = 0.3
+    randomize_wrist_std: float = 0.1  # radians (~5.7 degrees)
+
+    use_height_map: bool = False
+    height_map_resolution: float = 0.15
+    height_map_size: float = 1.5
+    height_map_max_dist: float = 5.0
+    motion_lib_num_dof: int | None = None  # If None, assumes motion lib DOF matches robot DOF
+    hand_default_positions: list[float] | None = (
+        None  # Default positions for extra joints (e.g., hands)
+    )
+    hand_default_velocities: list[float] | None = (
+        None  # Default velocities for extra joints (defaults to 0)
+    )
+    # Object z-offset (e.g., -0.05 to lower chair 5cm into ground)
+    object_z_offset: float = 0.0
+
+
+@configclass
+class ForceTrackingCommandCfg(CommandTermCfg):
+    """Configuration for ForceTrackingCommand external perturbation and compliance."""
+
+    class_type: type = ForceTrackingCommand
+
+    asset_name: str = dataclasses.MISSING
+
+    anchor_body: str = dataclasses.MISSING
+    body_names: list[str] = dataclasses.MISSING
+
+    force_update_frequency: int = 100
+    max_force: float = 20.0
+
+    vr_3point_body: list[str] = []  # noqa: RUF012
+    vr_3point_body_offset: list[list[float]] = []  # noqa: RUF012
+
+    force_push_body: list[str] = []  # default only two wrists  # noqa: RUF012
+    force_push_body_offset: list[list[float]] = []  # noqa: RUF012
+
+    # Debug print frequency (0 = disabled, nonzero = print every N steps)
+    # Usage: manager_env.commands.force.debug_print_every_n_steps=10
+    debug_print_every_n_steps: int = 0
+
+    waist_joints = ["waist_yaw_joint", "waist_roll_joint", "waist_pitch_joint"]  # noqa: RUF012
+
+    left_eef_deps = [  # noqa: RUF012
+        "left_shoulder_pitch_joint",
+        "left_shoulder_roll_joint",
+        "left_shoulder_yaw_joint",
+        "left_elbow_joint",
+        "left_wrist_pitch_joint",
+        "left_wrist_roll_joint",
+        "left_wrist_yaw_joint",
+    ]
+
+    right_eef_deps = [  # noqa: RUF012
+        "right_shoulder_pitch_joint",
+        "right_shoulder_roll_joint",
+        "right_shoulder_yaw_joint",
+        "right_elbow_joint",
+        "right_wrist_pitch_joint",
+        "right_wrist_roll_joint",
+        "right_wrist_yaw_joint",
+    ]
+
+
+def _get_body_indexes(command: TrackingCommand, body_names: list[str] | None) -> list[int]:
+    """Return indices into command.cfg.body_names for the requested subset.
+
+    Args:
+        command: TrackingCommand instance whose cfg.body_names defines the full body list.
+        body_names: Subset of body names to select. If None, return all indices.
+
+    Returns:
+        List of integer indices into ``command.cfg.body_names``.
+    """
+    return [
+        idx
+        for idx, body_name in enumerate(command.cfg.body_names)
+        if body_names is None or body_name in body_names
+    ]
+
+
+# Backward compat — remove after all checkpoints migrated
+MotionCommand = TrackingCommand
+MotionCommandCfg = TrackingCommandCfg
+ForceCommand = ForceTrackingCommand
+ForceCommandCfg = ForceTrackingCommandCfg
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/curriculum.py b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/curriculum.py
new file mode 100644
index 0000000000000000000000000000000000000000..8992c4c84e34ca6ef51b785ca060c20c994bb04c
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/curriculum.py
@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Common functions that can be used to create curriculum for the learning environment.
+
+The functions can be passed to the :class:`isaaclab.managers.CurriculumTermCfg` object to enable
+the curriculum introduced by the function.
+"""
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+from isaaclab.assets import Articulation
+from isaaclab.managers import SceneEntityCfg
+from isaaclab.terrains import TerrainImporter
+from isaaclab.utils import configclass
+import torch
+
+if TYPE_CHECKING:
+    from isaaclab.envs import ManagerBasedRLEnv
+
+
+@configclass
+class CurriculumCfg:
+    """Curriculum terms for the MDP."""
+
+    force_push_curriculum = None
+    force_push_linear_curriculum = None
+
+
+def terrain_levels_vel(
+    env: ManagerBasedRLEnv,
+    env_ids: Sequence[int],
+    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
+) -> torch.Tensor:
+    """Curriculum based on the distance the robot walked when commanded to move at a desired velocity.
+
+    This term is used to increase the difficulty of the terrain when the robot walks far enough and decrease the
+    difficulty when the robot walks less than half of the distance required by the commanded velocity.
+
+    .. note::
+        It is only possible to use this term with the terrain type ``generator``. For further information
+        on different terrain types, check the :class:`isaaclab.terrains.TerrainImporter` class.
+
+    Returns:
+        The mean terrain level for the given environment ids.
+    """
+    # extract the used quantities (to enable type-hinting)
+    asset: Articulation = env.scene[asset_cfg.name]
+    terrain: TerrainImporter = env.scene.terrain
+    command = env.command_manager.get_command("base_velocity")
+    # compute the distance the robot walked
+    distance = torch.norm(
+        asset.data.root_pos_w[env_ids, :2] - env.scene.env_origins[env_ids, :2], dim=1
+    )
+    # robots that walked far enough progress to harder terrains
+    move_up = distance > terrain.cfg.terrain_generator.size[0] / 2
+    # robots that walked less than half of their required distance go to simpler terrains
+    move_down = distance < torch.norm(command[env_ids, :2], dim=1) * env.max_episode_length_s * 0.5
+    move_down *= ~move_up
+    # update terrain levels
+    terrain.update_env_origins(env_ids, move_up, move_down)
+    # return the mean terrain level
+    return torch.mean(terrain.terrain_levels.float())
+
+
+def step_curriculum(env, env_ids, original_value, values, num_steps):
+    # Override after num_steps
+    assert len(values) == len(num_steps)
+    for i in range(len(values)):
+        if env.common_step_counter > num_steps[len(num_steps) - i - 1]:
+            return values[len(num_steps) - i - 1]
+    return original_value
+
+
+def linear_curriculum(env, env_ids, original_value, values, num_steps):
+    """
+    Linearly interpolates training curriculum values based on step counter.
+
+    Args:
+        env: IsaacLab environment (must have `common_step_counter`).
+        env_ids: Unused here, but kept for API consistency.
+        original_value (float): The base value before curriculum starts.
+        values (list of float): Target values at milestones.
+        num_steps (list of int): Step milestones corresponding to values.
+                                 Must be same length as `values`.
+
+    Returns:
+        float: interpolated value at current step.
+    """
+    assert len(values) == len(num_steps), "values and num_steps must match"
+
+    step = env.common_step_counter
+
+    # Between milestones → interpolate
+    for i in range(1, len(num_steps)):
+        if step <= num_steps[i]:
+            t0, t1 = num_steps[i - 1], num_steps[i]
+            v0, v1 = values[i - 1], values[i]
+            alpha = (step - t0) / (t1 - t0)
+            return v0 + alpha * (v1 - v0)
+
+    # After last milestone → final value
+    return values[-1]
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/events.py b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/events.py
new file mode 100644
index 0000000000000000000000000000000000000000..ee6b131581e8943ba46eed20cc2c42701d1ef92e
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/events.py
@@ -0,0 +1,142 @@
+"""Event functions for domain randomization and environment resets in RL training."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+from isaaclab.assets import Articulation
+from isaaclab.envs.mdp.events import _randomize_prop_by_op
+from isaaclab.managers import SceneEntityCfg
+from isaaclab.utils import configclass
+import isaaclab.utils.math as math_utils
+import torch
+
+if TYPE_CHECKING:
+    from isaaclab.envs import ManagerBasedEnv
+
+
+@configclass
+class EventCfg:
+    """Configuration for events."""
+
+    # startup
+    physics_material = None
+    add_joint_default_pos = None
+    add_hand_joint_default_pos = None
+    base_com = None
+
+    # interval - balance training
+    push_robot = None
+
+    randomize_rigid_body_mass = None
+
+
+def randomize_joint_default_pos(
+    env: ManagerBasedEnv,
+    env_ids: torch.Tensor | None,
+    asset_cfg: SceneEntityCfg,
+    pos_distribution_params: tuple[float, float] | None = None,
+    operation: Literal["add", "scale", "abs"] = "abs",
+    distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform",
+):
+    """Randomize joint default positions to simulate calibration errors.
+
+    Applies random offsets to the default joint positions of the robot, modeling
+    real-world joint encoder calibration inaccuracies. Also updates the action
+    manager offset to keep action space aligned with the new defaults.
+
+    Args:
+        env: The environment instance.
+        env_ids: Environment indices to randomize. If None, randomizes all.
+        asset_cfg: Scene entity config with joint IDs to randomize.
+        pos_distribution_params: Min/max range for the position offset distribution.
+        operation: How to combine the random value with the original ("add", "scale", "abs").
+        distribution: Sampling distribution type.
+    """
+    # extract the used quantities (to enable type-hinting)
+    asset: Articulation = env.scene[asset_cfg.name]
+
+    # save nominal value for export
+    asset.data.default_joint_pos_nominal = torch.clone(asset.data.default_joint_pos[0])
+
+    # resolve environment ids
+    if env_ids is None:
+        env_ids = torch.arange(env.scene.num_envs, device=asset.device)
+
+    # resolve joint indices
+    if asset_cfg.joint_ids == slice(None):
+        joint_ids = slice(None)  # for optimization purposes
+    else:
+        joint_ids = torch.tensor(asset_cfg.joint_ids, dtype=torch.int, device=asset.device)
+
+    if pos_distribution_params is not None:
+        pos = asset.data.default_joint_pos.to(asset.device).clone()
+        pos = _randomize_prop_by_op(
+            pos,
+            pos_distribution_params,
+            env_ids,
+            joint_ids,
+            operation=operation,
+            distribution=distribution,
+        )[env_ids][:, joint_ids]
+
+        if env_ids != slice(None) and joint_ids != slice(None):
+            env_ids = env_ids[:, None]
+        asset.data.default_joint_pos[env_ids, joint_ids] = pos
+        # update the offset in action since it is not updated automatically
+
+        action_joint_names = env.action_manager.get_term("joint_pos")._joint_names
+        asset_joint_names = asset.joint_names
+        shared_joint_names = list(set(action_joint_names).intersection(set(asset_joint_names)))
+        shared_joint_indices_action = [
+            action_joint_names.index(name) for name in shared_joint_names
+        ]
+        shared_joint_indices_asset = [asset_joint_names.index(name) for name in shared_joint_names]
+
+        shared_offset = asset.data.default_joint_pos[env_ids, shared_joint_indices_asset]
+        env.action_manager.get_term("joint_pos")._offset[
+            env_ids, shared_joint_indices_action
+        ] = shared_offset
+
+
+def randomize_rigid_body_com(
+    env: ManagerBasedEnv,
+    env_ids: torch.Tensor | None,
+    com_range: dict[str, tuple[float, float]],
+    asset_cfg: SceneEntityCfg,
+):
+    """Randomize the center of mass (CoM) of rigid bodies by adding a random value sampled from the given ranges.
+
+    .. note::
+        This function uses CPU tensors to assign the CoM. It is recommended to use this function
+        only during the initialization of the environment.
+    """
+    # extract the used quantities (to enable type-hinting)
+    asset: Articulation = env.scene[asset_cfg.name]
+    # resolve environment ids
+    if env_ids is None:
+        env_ids = torch.arange(env.scene.num_envs, device="cpu")
+    else:
+        env_ids = env_ids.cpu()
+
+    # resolve body indices
+    if asset_cfg.body_ids == slice(None):
+        body_ids = torch.arange(asset.num_bodies, dtype=torch.int, device="cpu")
+    else:
+        body_ids = torch.tensor(asset_cfg.body_ids, dtype=torch.int, device="cpu")
+
+    # sample random CoM values
+    range_list = [com_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]]
+    ranges = torch.tensor(range_list, device="cpu")
+    rand_samples = math_utils.sample_uniform(
+        ranges[:, 0], ranges[:, 1], (len(env_ids), 3), device="cpu"
+    ).unsqueeze(1)
+
+    # get the current com of the bodies (num_assets, num_bodies)
+    coms = asset.root_physx_view.get_coms().clone()
+
+    # Randomize the com in range
+    coms[:, body_ids, :3] += rand_samples
+
+    # Set the new coms
+    asset.root_physx_view.set_coms(coms, env_ids)
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/observations.py b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/observations.py
new file mode 100644
index 0000000000000000000000000000000000000000..a9c9c5b62887dd3ea778860e14de41e6da73b381
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/observations.py
@@ -0,0 +1,2290 @@
+"""Observation functions for the manager-based RL environment MDP."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from isaaclab.utils.math import (
+    matrix_from_quat,
+    quat_apply,
+    quat_apply_inverse,
+    quat_apply_yaw,
+    quat_conjugate,
+    quat_inv,
+    quat_mul,
+    subtract_frame_transforms,
+)
+import torch
+
+from gear_sonic.envs.env_utils import joint_utils
+from gear_sonic.envs.manager_env.mdp import commands, utils
+from gear_sonic.trl.utils import torch_transform
+
+if TYPE_CHECKING:
+    from isaaclab.envs import ManagerBasedEnv
+
+
+from isaaclab.managers import ObservationGroupCfg as ObsGroup
+from isaaclab.managers import SceneEntityCfg
+from isaaclab.utils import configclass
+
+# Joint ordering constants (Mujoco order for compatibility)
+G1_MUJOCO_ORDER = [
+    "left_hip_pitch_joint",
+    "left_hip_roll_joint",
+    "left_hip_yaw_joint",
+    "left_knee_joint",
+    "left_ankle_pitch_joint",
+    "left_ankle_roll_joint",
+    "right_hip_pitch_joint",
+    "right_hip_roll_joint",
+    "right_hip_yaw_joint",
+    "right_knee_joint",
+    "right_ankle_pitch_joint",
+    "right_ankle_roll_joint",
+    "waist_yaw_joint",
+    "waist_roll_joint",
+    "waist_pitch_joint",
+    "left_shoulder_pitch_joint",
+    "left_shoulder_roll_joint",
+    "left_shoulder_yaw_joint",
+    "left_elbow_joint",
+    "left_wrist_roll_joint",
+    "left_wrist_pitch_joint",
+    "left_wrist_yaw_joint",
+    "right_shoulder_pitch_joint",
+    "right_shoulder_roll_joint",
+    "right_shoulder_yaw_joint",
+    "right_elbow_joint",
+    "right_wrist_roll_joint",
+    "right_wrist_pitch_joint",
+    "right_wrist_yaw_joint",
+]
+
+# Index mappings for 29 DOF
+isaaclab_to_mujoco_dof = [joint_utils.G1_ISAACLab_ORDER.index(i) for i in G1_MUJOCO_ORDER]
+mujoco_to_isaaclab = [G1_MUJOCO_ORDER.index(i) for i in joint_utils.G1_ISAACLab_ORDER]
+
+
+@configclass
+class PolicyCfg(ObsGroup):
+    """Observations for policy group."""
+
+    # observation terms (order preserved)
+    command = None
+    command_vel = None
+    command_max = None
+    command_multi_future = None
+    command_multi_future_joint_pos = None
+    command_multi_future_joint_body_pos = None
+    command_multi_future_joint_body_diff_pos = None
+    command_multi_future_joint_body_abs_pos = None
+    command_multi_future_lower_body = None
+    command_max_multi_future = None
+    command_max_diff_w = None
+    command_max_diff_w_multi_future = None
+    command_max_diff_l = None
+    command_max_diff_l_multi_future = None
+    command_z_multi_future = None
+    root_pos_multi_future = None
+    root_quat_multi_future = None
+    joint_pos_multi_future = None
+    smpl_pose = None
+    smpl_body_pose = None
+    smpl_joints_multi_future = None
+    smpl_joints_lower_multi_future = None
+    smpl_root_ori_b = None
+    motion_anchor_pos_b = None
+    motion_anchor_ori_b_mf = None
+    motion_anchor_pos_b_xy = None
+    motion_anchor_ori_w_mf = None
+    motion_anchor_ori_heading_mf = None
+    motion_anchor_ori_b = None
+    motion_anchor_ori_w = None
+    motion_anchor_yaw_b = None
+    robot_anchor_ori_w = None
+    base_lin_vel = None
+    base_ang_vel = None
+    joint_pos = None
+    joint_vel = None
+    actions = None
+    # Body-only observations (for pre-trained action_transform_module)
+    joint_pos_wo_hand = None
+    joint_vel_wo_hand = None
+    actions_wo_hand = None
+    vr_3point_target = None
+    vr_3point_target_compliant = None
+    vr_3point_target_compliant_multi_future = None
+    vr_3point_target_multi_future = None
+    vr_3point_orn_target_multi_future = None
+    vr_3point_local_target = None
+    vr_3point_local_target_compliant = None
+    vr_3point_local_target_multi_future = None
+    vr_3point_local_orn_target = None
+    head_orn_target_multi_future = None
+    vr_wrists_local_pos_target = None
+    vr_wrists_local_orn_target = None
+    vr_head_local_orn_target = None
+    gravity_dir = None
+    compliance = None
+    ext_forces = None
+    motion_anchor_gravity_dir = None
+    body_pos = None
+    body_pos_diff_l = None
+    # HOI manipulation task specific
+    target_object_pos = None
+    hand_object_transform = None
+    finger_tips_force = None
+    # Object future motion observations
+    object_pos_b_multi_future = None
+    object_ori_b_multi_future = None
+    object_pos_delta_multi_future = None
+    object_ori_delta_multi_future_6d = None
+    grab_contact_flag = None
+    # Hand-object transform
+    hand_object_transform_6d = None
+    # Object-to-root observations (current state in body frame)
+    object_pos_b = None
+    object_ori_b_6d = None
+    # Table observations
+    table_pos_b = None
+    table_ori_b = None
+
+    # Policy output from last step (latent residual + primitives, e.g., 64+2=66 dims)
+    last_meta_action = None
+
+    # Terrain observations
+    height_map_flat = None
+
+
+@configclass
+class PolicyAtmCfg(ObsGroup):
+    """Observations for action_transform_module (ATM).
+
+    This observation group provides body-only observations (29 DOF) for use with
+    pre-trained action_transform_module. When using a 43 DOF robot (29 body + 14 hand),
+    this group extracts only the body joint observations matching ATM's expected input.
+
+    NOTE: Order must match the observations the pretrained ATM was trained with.
+    """
+
+    # Order matches PolicyCfg: base_ang_vel, joint_pos, joint_vel, actions, gravity_dir
+    base_ang_vel = None
+    joint_pos_wo_hand = None
+    joint_vel_wo_hand = None
+    actions_wo_hand = None
+    gravity_dir = None
+
+
+@configclass
+class TeacherCfg(ObsGroup):
+    """Teacher observations for distillation.
+
+    This provides privileged state observations that the teacher policy uses.
+    The student policy learns to imitate the teacher's actions using vision instead.
+    """
+
+    # # Basic proprioception
+    # command = None
+    # motion_anchor_pos_b = None
+    # motion_anchor_ori_b = None
+    # base_lin_vel = None
+    # base_ang_vel = None
+    # joint_pos = None
+    # joint_vel = None
+    # actions = None
+    # # Object and manipulation observations
+    # target_object_pos = None
+    # hand_object_transform = None
+    # finger_tips_force = None
+    # grab_contact_flag = None
+    # # Object future motion observations
+    # object_pos_b_multi_future = None
+    # object_ori_b_multi_future = None
+    # # Body multi-future observations
+    # command_multi_future = None
+    # motion_anchor_ori_b_mf = None
+
+    # observation terms (order preserved)
+    command = None
+    command_vel = None
+    command_max = None
+    command_multi_future = None
+    command_multi_future_joint_pos = None
+    command_multi_future_joint_body_pos = None
+    command_multi_future_joint_body_diff_pos = None
+    command_multi_future_joint_body_abs_pos = None
+    command_multi_future_lower_body = None
+    command_max_multi_future = None
+    command_max_diff_w = None
+    command_max_diff_w_multi_future = None
+    command_max_diff_l = None
+    command_max_diff_l_multi_future = None
+    command_z_multi_future = None
+    root_pos_multi_future = None
+    root_quat_multi_future = None
+    joint_pos_multi_future = None
+    smpl_pose = None
+    smpl_body_pose = None
+    smpl_joints_multi_future = None
+    smpl_joints_lower_multi_future = None
+    smpl_root_ori_b = None
+    motion_anchor_pos_b = None
+    motion_anchor_ori_b_mf = None
+    motion_anchor_pos_b_xy = None
+    motion_anchor_ori_w_mf = None
+    motion_anchor_ori_b = None
+    motion_anchor_ori_w = None
+    motion_anchor_yaw_b = None
+    robot_anchor_ori_w = None
+    base_lin_vel = None
+    base_ang_vel = None
+    joint_pos = None
+    joint_vel = None
+    actions = None
+    # Body-only observations (for pre-trained action_transform_module)
+    joint_pos_wo_hand = None
+    joint_vel_wo_hand = None
+    actions_wo_hand = None
+    vr_3point_target = None
+    vr_3point_target_compliant = None
+    vr_3point_target_compliant_multi_future = None
+    vr_3point_target_multi_future = None
+    vr_3point_orn_target_multi_future = None
+    vr_3point_local_target = None
+    vr_3point_local_target_compliant = None
+    vr_3point_local_target_multi_future = None
+    vr_3point_local_orn_target = None
+    head_orn_target_multi_future = None
+    vr_wrists_local_pos_target = None
+    vr_wrists_local_orn_target = None
+    vr_head_local_orn_target = None
+    gravity_dir = None
+    compliance = None
+    ext_forces = None
+    motion_anchor_gravity_dir = None
+    body_pos = None
+    body_pos_diff_l = None
+    # HOI manipulation task specific
+    target_object_pos = None
+    hand_object_transform = None
+    finger_tips_force = None
+    # Object future motion observations
+    object_pos_b_multi_future = None
+    object_ori_b_multi_future = None
+    grab_contact_flag = None
+    # Policy output from last step (latent residual + primitives, e.g., 64+2=66 dims)
+    last_meta_action = None
+
+
+@configclass
+class PrivilegedCfg(ObsGroup):
+    """Privileged observations for the critic network (asymmetric actor-critic)."""
+
+    command = None
+    command_max = None
+    command_multi_future = None
+    command_multi_future_lower_body = None
+    command_multi_future_lower_body_joint_pos = None
+    command_max_multi_future = None
+    command_max_diff_w = None
+    command_max_diff_w_multi_future = None
+    command_max_diff_l = None
+    command_max_diff_l_multi_future = None
+    command_z_multi_future = None
+    motion_anchor_pos_b = None
+    motion_anchor_ori_b = None
+    motion_anchor_ori_b_mf = None
+    motion_anchor_ori_heading_mf = None
+    body_pos = None
+    body_ori = None
+    base_lin_vel = None
+    base_ang_vel = None
+    joint_pos = None
+    joint_vel = None
+    actions = None
+    vr_3point_target = None
+    vr_3point_target_compliant_multi_future = None
+    vr_3point_target_multi_future = None
+    vr_3point_orn_target_multi_future = None
+    head_orn_target_multi_future = None
+    vr_3point_local_target = None
+    vr_3point_local_target_compliant = None
+    vr_3point_local_target_multi_future = None
+    vr_3point_local_orn_target = None
+    vr_wrists_local_pos_target = None
+    vr_wrists_local_orn_target = None
+    vr_head_local_orn_target = None
+    gravity_dir = None
+    compliance = None
+    ext_forces = None
+    motion_anchor_gravity_dir = None
+    # HOI manipulation task specific
+    target_object_pos = None
+    hand_object_transform = None
+    finger_tips_force = None
+    # Object future motion observations
+    object_pos_b_multi_future = None
+    object_ori_b_multi_future = None
+    object_pos_delta_multi_future = None
+    object_ori_delta_multi_future_6d = None
+    # Hand-object transform
+    hand_object_transform_6d = None
+    # Object-to-root observations (current state in body frame)
+    object_pos_b = None
+    object_ori_b_6d = None
+    # Table observations
+    table_pos_b = None
+    table_ori_b = None
+    # Staged training
+    task_stage = None
+
+    # g1 token obs
+    ref_root_pos_future_b = None
+    ref_root_ori_future_b = None
+    diff_body_pos_future_local = None
+    diff_body_ori_future_local = None
+    diff_body_lin_vel_future_local = None
+    diff_body_ang_vel_future_local = None
+    grab_contact_flag = None
+
+    # Terrain observations
+    height_map_flat = None
+
+
+@configclass
+class DiscriminatorCfg:
+    """Observation specifications for the discriminator."""
+
+    disc_obs = None
+    ref_disc_obs = None
+
+
+@configclass
+class TokenizerCfg(ObsGroup):
+    """Observations for the tokenizer (SONIC/ATM encoder input)."""
+
+    encoder_index = None
+    command_multi_future_nonflat = None
+    motion_anchor_ori_w = None
+    command_z_multi_future_nonflat = None
+    command_z = None
+    motion_anchor_ori_b = None
+    motion_anchor_ori_heading_b = None
+    motion_anchor_ori_b_nonflat = None
+    motion_anchor_ori_b_mf_nonflat = None
+    motion_anchor_ori_w_mf_nonflat = None
+    command_multi_future_egocentric_joint_transforms = None
+    command_multi_future_egocentric_joint_transforms_nonflat = None
+    command_multi_future_egocentric_joint_positions = None
+    command_multi_future_egocentric_joint_positions_nonflat = None
+    command_multi_future_egocentric_joint_rotations = None
+    command_multi_future_egocentric_joint_rotations_nonflat = None
+    command_multi_future_root_transforms = None
+    command_multi_future_root_transforms_nonflat = None
+    motion_anchor_ori_heading_mf_nonflat = None
+    motion_anchor_ori_refheading_mf_nonflat = None
+    heading_diff_robot_ref = None
+    motion_anchor_ori_refheading = None
+    motion_anchor_ori_heading = None
+    command_multi_future_lower_body = None
+    vr_3point_local_target = None
+    vr_3point_local_orn_target = None
+    vr_3point_local_target_compliant = None
+    smpl_joints_multi_future_nonflat = None
+    smpl_joints_multi_future_local_nonflat = None
+    smpl_joints_multi_future_local_flatten = None
+    smpl_lower_body_joints_multi_future_local_nonflat = None
+    smpl_lower_body_joints_multi_future_local_flatten = None
+    smpl_joints_lower_multi_future_local_flatten = None
+    smpl_transl_z_multi_future_nonflat = None
+    smpl_root_ori_b_multi_future = None
+    smpl_root_ori_b_multi_future_flatten = None
+    smpl_root_ori_refheading_multi_future = None
+    smpl_root_ori_heading_multi_future = None
+    smpl_elbow_wrist_pose_multi_future = None
+    smpl_wrist_pose_multi_future = None
+    joint_pos_multi_future_wrist = None
+    joint_pos_multi_future_wrist_flatten = None
+    joint_pos_multi_future_wrist_for_smpl = None
+    # SOMA skeleton observations
+    soma_joints_multi_future_local_nonflat = None
+    soma_root_ori_b_multi_future = None
+    joint_pos_multi_future_wrist_for_soma = None
+    # Object goal observations (position and orientation in robot body frame)
+    object_pos_b = None
+    object_ori_b = None
+    object_ori_b_6d = None
+    ref_root_pos_future_b = None
+    ref_root_ori_future_b = None
+    diff_body_pos_future_local = None
+    diff_body_ori_future_local = None
+    diff_body_lin_vel_future_local = None
+    diff_body_ang_vel_future_local = None
+    compliance = None
+    # HOI encoder observations (for end-to-end SONIC-HOI training)
+    command = None
+    motion_anchor_pos_b = None
+    motion_anchor_ori_b_mf = None
+    base_lin_vel = None
+    base_ang_vel = None
+    joint_pos = None
+    joint_vel = None
+    actions = None
+    target_object_pos = None
+    hand_object_transform_6d = None
+    finger_tips_force = None
+    table_pos_b = None
+    table_ori_b = None
+    object_pos_delta_multi_future = None
+    object_ori_delta_multi_future_6d = None
+    command_multi_future = None
+
+
+@configclass
+class HeightMapCfg(ObsGroup):
+    """Height map observation group for terrain-aware locomotion."""
+
+    height_map = None
+
+
+@configclass
+class CameraRGBCfg(ObsGroup):
+    """Camera RGB image observation group.
+
+    This is a separate observation group for vision observations.
+    It will be passed through the wrapper as 'camera_rgb' key in obs_dict,
+    not concatenated with other policy observations.
+    """
+
+    camera_rgb = None
+
+
+@configclass
+class ResidualAction(ObsGroup):
+    """Observation group for residual action feedback."""
+
+    residual_action = None
+
+
+@configclass
+class ObservationsCfg:
+    """Observation specifications for the MDP."""
+
+    # observation groups
+    policy: PolicyCfg = None
+    critic: PrivilegedCfg = None
+    disc: DiscriminatorCfg = None
+    tokenizer: TokenizerCfg = None
+    policy_atm: PolicyAtmCfg = None  # Body-only obs for pre-trained action_transform_module
+    height_map: HeightMapCfg = None
+    teacher: TeacherCfg = None  # Teacher observations for distillation
+    camera_rgb: CameraRGBCfg = None  # Separate vision observation group
+    residual_action: ResidualAction = None
+
+
+def command_max(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get maximum reference body positions in heading-local frame.
+
+    Returns:
+        torch.Tensor: Flattened body positions, shape (num_envs, num_bodies * 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+
+    return command.command_max
+
+
+def command_vel(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get reference body velocities.
+
+    Returns:
+        torch.Tensor: Flattened body velocities, shape (num_envs, vel_dim).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    return command.command_vel
+
+
+def command_max_diff(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get difference between reference and robot body positions.
+
+    Returns:
+        torch.Tensor: Flattened position differences, shape (num_envs, num_bodies * 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    return command.command_max_diff
+
+
+def command_max_diff_multi_future(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get reference-minus-robot body position differences for multiple future frames.
+
+    Returns:
+        torch.Tensor: Flattened multi-future differences,
+            shape (num_envs, num_future_frames * num_bodies * 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    return command.command_max_diff_multi_future
+
+
+def command_z(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get reference root z-height.
+
+    Returns:
+        torch.Tensor: Root z-height, shape (num_envs, 1).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    return command.command_z
+
+
+def command_z_multi_future(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Get reference root z-heights for multiple future frames.
+
+    Args:
+        command_name: Name of the tracking command term.
+        non_flatten: If True, return shape (num_envs, num_future_frames, 1).
+
+    Returns:
+        torch.Tensor: Z-heights, shape (num_envs, num_future_frames) when flattened.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    if non_flatten:
+        return command.command_z_multi_future.reshape(
+            command.num_envs, command.num_future_frames, -1
+        )
+    else:
+        return command.command_z_multi_future
+
+
+def command_max_multi_future(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get maximum reference body positions for multiple future frames in heading-local frame.
+
+    Returns:
+        torch.Tensor: Flattened multi-future body positions,
+            shape (num_envs, num_future_frames * num_bodies * 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    return command.command_max_multi_future
+
+
+def command_multi_future(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Get reference body positions in body-local frame for multiple future frames.
+
+    Args:
+        command_name: Name of the tracking command term.
+        non_flatten: If True, return shape (num_envs, num_future_frames, num_bodies * 3).
+
+    Returns:
+        torch.Tensor: Body positions, shape (num_envs, num_future_frames * num_bodies * 3)
+            when flattened.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    if non_flatten:
+        return command.command_multi_future.reshape(command.num_envs, command.num_future_frames, -1)
+    else:
+        return command.command_multi_future
+
+
+def command_multi_future_joint_pos(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Get reference joint positions for multiple future frames.
+
+    Args:
+        command_name: Name of the tracking command term.
+        non_flatten: If True, return shape (num_envs, num_future_frames, num_joints).
+
+    Returns:
+        torch.Tensor: Joint positions, shape (num_envs, num_future_frames * num_joints)
+            when flattened.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    if non_flatten:
+        return command.command_multi_future_joint_pos.reshape(
+            command.num_envs, command.num_future_frames, -1
+        )
+    else:
+        return command.command_multi_future_joint_pos
+
+
+def command_multi_future_joint_body_pos(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Get reference joint body positions for multiple future frames.
+
+    Args:
+        command_name: Name of the tracking command term.
+        non_flatten: If True, return shape (num_envs, num_future_frames, num_joints * 3).
+
+    Returns:
+        torch.Tensor: Joint body positions, shape (num_envs, num_future_frames * num_joints * 3)
+            when flattened.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    if non_flatten:
+        return command.command_multi_future_joint_body_pos.reshape(
+            command.num_envs, command.num_future_frames, -1
+        )
+    else:
+        return command.command_multi_future_joint_body_pos
+
+
+def command_multi_future_joint_body_diff_pos(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Get reference-minus-robot joint body position differences for multiple future frames.
+
+    Args:
+        command_name: Name of the tracking command term.
+        non_flatten: If True, return shape (num_envs, num_future_frames, num_joints * 3).
+
+    Returns:
+        torch.Tensor: Joint body position differences,
+            shape (num_envs, num_future_frames * num_joints * 3) when flattened.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    if non_flatten:
+        return command.command_multi_future_joint_body_diff_pos.reshape(
+            command.num_envs, command.num_future_frames, -1
+        )
+    else:
+        return command.command_multi_future_joint_body_diff_pos
+
+
+def command_multi_future_joint_body_abs_pos(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Get absolute reference joint body positions for multiple future frames.
+
+    Args:
+        command_name: Name of the tracking command term.
+        non_flatten: If True, return shape (num_envs, num_future_frames, num_joints * 3).
+
+    Returns:
+        torch.Tensor: Absolute joint body positions,
+            shape (num_envs, num_future_frames * num_joints * 3) when flattened.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    if non_flatten:
+        return command.command_multi_future_joint_body_abs_pos.reshape(
+            command.num_envs, command.num_future_frames, -1
+        )
+    else:
+        return command.command_multi_future_joint_body_abs_pos
+
+
+def command_multi_future_lower_body(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get reference lower-body positions for multiple future frames in body-local frame.
+
+    Returns:
+        torch.Tensor: Flattened lower-body positions,
+            shape (num_envs, num_future_frames * lower_body_dim).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    return command.command_multi_future_lower_body
+
+
+def command_multi_future_lower_body_joint_pos(
+    env: ManagerBasedEnv, command_name: str
+) -> torch.Tensor:
+    """Get reference lower-body joint positions for multiple future frames.
+
+    Returns:
+        torch.Tensor: Flattened lower-body joint positions,
+            shape (num_envs, num_future_frames * lower_body_joints).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    return command.command_multi_future_lower_body_joint_pos
+
+
+def command_max_diff_w(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get reference-minus-robot body position differences in world frame.
+
+    Returns:
+        torch.Tensor: Flattened position differences, shape (num_envs, num_bodies * 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    return command.command_max_diff_w
+
+
+def command_max_diff_w_multi_future(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get reference-minus-robot body position differences in world frame for multiple future frames.
+
+    Returns:
+        torch.Tensor: Flattened multi-future differences,
+            shape (num_envs, num_future_frames * num_bodies * 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    return command.command_max_diff_w_multi_future
+
+
+def command_max_diff_l(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get reference-minus-robot body position differences in body-local frame.
+
+    Returns:
+        torch.Tensor: Flattened position differences, shape (num_envs, num_bodies * 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    return command.command_max_diff_l
+
+
+def command_max_diff_l_multi_future(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get reference-minus-robot body position differences in body-local frame for multiple future frames.
+
+    Returns:
+        torch.Tensor: Flattened multi-future differences,
+            shape (num_envs, num_future_frames * num_bodies * 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    return command.command_max_diff_l_multi_future
+
+
+def command_num_frames(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Per-env number of valid future frames. Shape [num_envs, 1]."""
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    if getattr(command, "per_env_num_frames", None) is not None:
+        return command.per_env_num_frames.float().reshape(-1, 1)
+    return torch.full(
+        (command.num_envs, 1), float(command.num_future_frames), device=command.device
+    )
+
+
+def robot_anchor_ori_w(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get robot anchor (pelvis) orientation in world frame as 6D rotation.
+
+    Returns:
+        torch.Tensor: First two columns of rotation matrix, shape (num_envs, 6).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    mat = matrix_from_quat(command.robot_anchor_quat_w)
+    return mat[..., :2].reshape(mat.shape[0], -1)
+
+
+def motion_anchor_gravity_dir(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Compute gravity direction in the reference motion anchor frame.
+
+    Transforms the world-frame down vector into the reference anchor's local frame.
+
+    Returns:
+        torch.Tensor: Gravity direction vector, shape (num_envs, 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    gravity_dir = quat_apply(quat_inv(command.anchor_quat_w), command.down_dir)
+    return gravity_dir
+
+
+def gravity_dir(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Compute gravity direction in the robot anchor (pelvis) frame.
+
+    Transforms the world-frame down vector into the robot's local frame.
+    Provides the policy with tilt/orientation information.
+
+    Returns:
+        torch.Tensor: Gravity direction vector, shape (num_envs, 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    gravity_dir = quat_apply(quat_inv(command.robot_anchor_quat_w), command.down_dir)
+    return gravity_dir
+
+
+def robot_anchor_lin_vel_w(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get robot anchor (pelvis) linear velocity in world frame.
+
+    Returns:
+        torch.Tensor: Linear velocity xyz, shape (num_envs, 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+
+    return command.robot_anchor_vel_w[:, :3].view(env.num_envs, -1)
+
+
+def robot_anchor_ang_vel_w(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get robot anchor (pelvis) angular velocity in world frame.
+
+    Returns:
+        torch.Tensor: Angular velocity xyz, shape (num_envs, 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+
+    return command.robot_anchor_vel_w[:, 3:6].view(env.num_envs, -1)
+
+
+def robot_body_pos_b(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get robot body positions in robot anchor (pelvis) local frame.
+
+    Transforms all tracked body positions from world frame into the robot's
+    anchor frame using subtract_frame_transforms.
+
+    Returns:
+        torch.Tensor: Flattened body positions, shape (num_envs, num_bodies * 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+
+    num_bodies = len(command.cfg.body_names)
+    pos_b, _ = subtract_frame_transforms(
+        command.robot_anchor_pos_w[:, None, :].repeat(1, num_bodies, 1),
+        command.robot_anchor_quat_w[:, None, :].repeat(1, num_bodies, 1),
+        command.robot_body_pos_w,
+        command.robot_body_quat_w,
+    )
+
+    return pos_b.view(env.num_envs, -1)
+
+
+def robot_body_pos_diff_l(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get reference-minus-robot body position differences in world frame.
+
+    Computes (reference_body_pos - robot_body_pos) for each tracked body.
+
+    NOTE: Despite the ``_l`` suffix, the subtraction is done in world frame
+    (body_pos_relative_w - robot_body_pos_w).
+
+    Returns:
+        torch.Tensor: Flattened position differences, shape (num_envs, num_bodies * 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    body_indexes = commands._get_body_indexes(command, command.cfg.body_names)  # noqa: SLF001
+    body_pos_diff_l = (
+        command.body_pos_relative_w[:, body_indexes] - command.robot_body_pos_w[:, body_indexes]
+    )
+    return body_pos_diff_l.view(env.num_envs, -1)
+
+
+def robot_body_ori_b(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get robot body orientations in robot anchor (pelvis) local frame as 6D rotation.
+
+    Returns:
+        torch.Tensor: First two columns of each body's rotation matrix,
+            shape (num_envs, num_bodies * 6).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+
+    num_bodies = len(command.cfg.body_names)
+    _, ori_b = subtract_frame_transforms(
+        command.robot_anchor_pos_w[:, None, :].repeat(1, num_bodies, 1),
+        command.robot_anchor_quat_w[:, None, :].repeat(1, num_bodies, 1),
+        command.robot_body_pos_w,
+        command.robot_body_quat_w,
+    )
+    mat = matrix_from_quat(ori_b)
+    return mat[..., :2].reshape(mat.shape[0], -1)
+
+
+def motion_anchor_pos_b(env: ManagerBasedEnv, command_name: str, mask_out_z=False) -> torch.Tensor:
+    """Get reference motion anchor position relative to robot anchor in robot-local frame.
+
+    Computes the position of the reference motion's root relative to the robot's
+    root, expressed in the robot's local coordinate frame.
+
+    Args:
+        command_name: Name of the tracking command term.
+        mask_out_z: If True, return only xy components (discard z).
+
+    Returns:
+        torch.Tensor: Position offset, shape (num_envs, 3) or (num_envs, 2) if mask_out_z.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    pos, _ = subtract_frame_transforms(
+        command.robot_anchor_pos_w,
+        command.robot_anchor_quat_w,
+        command.anchor_pos_w,
+        command.anchor_quat_w,
+    )
+    if mask_out_z:
+        pos = pos[:, :2]
+
+    return pos.view(env.num_envs, -1)
+
+
+def motion_anchor_yaw_b(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get yaw (heading) difference between reference anchor and robot anchor.
+
+    Extracts only the heading component of the relative orientation between the
+    reference motion root and the robot root.
+
+    Returns:
+        torch.Tensor: Heading quaternion, shape (num_envs, 4).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+
+    _, ori = subtract_frame_transforms(
+        command.robot_anchor_pos_w,
+        command.robot_anchor_quat_w,
+        command.anchor_pos_w,
+        command.anchor_quat_w,
+    )
+    yaw = torch_transform.get_heading_q(ori)
+    return yaw.view(env.num_envs, -1)
+
+
+def motion_anchor_ori_heading_b(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get reference root orientation with heading removed as 6D rotation.
+
+    Removes the reference anchor's own heading (yaw) from the reference root
+    quaternion, preserving pitch/roll relative to gravity.
+
+    Returns:
+        torch.Tensor: 6D rotation (first two columns of rotation matrix),
+            shape (num_envs, 6).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    ref_root_quat = command.motion_lib.get_root_quat_w(
+        command.motion_ids, command.motion_start_time_steps + command.time_steps
+    )
+    root_heading_inv = quat_inv(command.anchor_heading_quat).view(env.num_envs, 4)
+    deheaded_ref_rot = quat_mul(root_heading_inv, ref_root_quat)
+    mat = matrix_from_quat(deheaded_ref_rot)
+    deheaded_ref_root_mat = mat[..., :2].reshape(mat.shape[0], -1)
+    return deheaded_ref_root_mat
+
+
+def motion_anchor_ori_b(env: ManagerBasedEnv, command_name: str, non_flatten=False) -> torch.Tensor:
+    """Get reference anchor orientation relative to robot anchor as 6D rotation.
+
+    Computes the orientation difference between the reference motion root and the
+    robot root, expressed in the robot's local frame.
+
+    Args:
+        command_name: Name of the tracking command term.
+        non_flatten: If True, repeat across future frames and return
+            shape (num_envs, num_future_frames, 6).
+
+    Returns:
+        torch.Tensor: 6D rotation representation, shape (num_envs, 6) when flattened.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+
+    _, ori = subtract_frame_transforms(
+        command.robot_anchor_pos_w,  # robot
+        command.robot_anchor_quat_w,
+        command.anchor_pos_w,  # reference
+        command.anchor_quat_w,
+    )
+    mat = matrix_from_quat(ori)
+    ori = mat[..., :2].reshape(mat.shape[0], -1)
+    if non_flatten:
+        return ori.unsqueeze(1).repeat(1, command.num_future_frames, 1)
+    else:
+        return ori
+
+
+def motion_anchor_ori_refheading(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Motion anchor orientation canonicalized by its own heading.
+
+    Removes the heading component from the anchor orientation, preserving
+    pitch/roll relative to gravity.
+
+    Returns:
+        torch.Tensor: 6D rotation matrix representation,
+            shape (num_envs, 6) or (num_envs, num_future_frames, 6) if non_flatten
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    ori = command.anchor_ori_refheading  # (num_envs, 6)
+    if non_flatten:
+        return ori.unsqueeze(1).repeat(1, command.num_future_frames, 1)
+    return ori
+
+
+def motion_anchor_ori_heading(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Motion anchor orientation canonicalized by robot heading (yaw).
+
+    Uses the robot's heading for canonicalization, preserving the reference
+    motion's pitch/roll relative to gravity while removing heading difference.
+
+    Returns:
+        torch.Tensor: 6D rotation matrix representation,
+            shape (num_envs, 6) or (num_envs, num_future_frames, 6) if non_flatten
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    ori = command.anchor_ori_heading  # (num_envs, 6)
+    if non_flatten:
+        return ori.unsqueeze(1).repeat(1, command.num_future_frames, 1)
+    return ori
+
+
+def motion_anchor_ori_w(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get reference motion anchor orientation in world frame as 6D rotation.
+
+    Returns:
+        torch.Tensor: First two columns of rotation matrix, shape (num_envs, 6).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    mat = matrix_from_quat(command.anchor_quat_w)
+    return mat[..., :2].reshape(mat.shape[0], -1)
+
+
+def motion_anchor_ori_b_mf(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Get reference-vs-robot root orientation difference for multiple future frames in body-local frame.
+
+    Uses the robot's full orientation (including pitch/roll) for normalization.
+
+    Args:
+        command_name: Name of the tracking command term.
+        non_flatten: If True, return shape (num_envs, num_future_frames, 6).
+
+    Returns:
+        torch.Tensor: 6D rotation differences,
+            shape (num_envs, num_future_frames * 6) when flattened.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    if non_flatten:
+        return command.root_rot_dif_l_multi_future.reshape(
+            env.num_envs, command.num_future_frames, -1
+        )
+    else:
+        return command.root_rot_dif_l_multi_future.reshape(env.num_envs, -1)
+
+
+def motion_anchor_ori_w_mf(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Get reference root orientations in world frame for multiple future frames as 6D rotation.
+
+    Args:
+        command_name: Name of the tracking command term.
+        non_flatten: If True, return shape (num_envs, num_future_frames, 6).
+
+    Returns:
+        torch.Tensor: 6D rotation representations,
+            shape (num_envs, num_future_frames * 6) when flattened.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    if non_flatten:
+        return command.root_rot_w_multi_future.reshape(env.num_envs, command.num_future_frames, -1)
+    else:
+        return command.root_rot_w_multi_future.reshape(env.num_envs, -1)
+
+
+# =============================================================================
+# Egocentric joint transforms and root transforms relative to first frame
+# =============================================================================
+
+
+def command_multi_future_egocentric_joint_transforms(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Egocentric joint transforms (positions + rotations) for reference frames.
+
+    For each future frame, joint positions and rotations are relative to that frame's
+    projected root (heading/yaw only rotation, z=0).
+
+    Returns:
+        if non_flatten:
+            [num_envs, num_future_frames, num_bodies_full * 9] (3 pos + 6 rot per body, flattened)
+        else:
+            [num_envs, num_future_frames * num_bodies_full * 9]
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    transforms = command.egocentric_joint_transforms_multi_future
+    if non_flatten:
+        return transforms.reshape(env.num_envs, command.num_future_frames, -1)
+    else:
+        return transforms.reshape(env.num_envs, -1)
+
+
+def command_multi_future_egocentric_joint_positions(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Egocentric joint positions for reference frames.
+
+    Returns:
+        if non_flatten:
+            [num_envs, num_future_frames, num_bodies_full * 3]
+        else:
+            [num_envs, num_future_frames * num_bodies_full * 3]
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    positions = command.egocentric_joint_positions_multi_future
+    if non_flatten:
+        return positions.reshape(env.num_envs, command.num_future_frames, -1)
+    else:
+        return positions.reshape(env.num_envs, -1)
+
+
+def command_multi_future_egocentric_joint_rotations(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Egocentric joint rotations (6D) for reference frames.
+
+    Returns:
+        if non_flatten:
+            [num_envs, num_future_frames, num_bodies_full * 6]
+        else:
+            [num_envs, num_future_frames * num_bodies_full * 6]
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    rotations = command.egocentric_joint_rotations_multi_future
+    if non_flatten:
+        return rotations.reshape(env.num_envs, command.num_future_frames, -1)
+    else:
+        return rotations.reshape(env.num_envs, -1)
+
+
+def command_multi_future_root_transforms(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Root transforms (position + rotation) relative to the first reference frame.
+
+    Position: Delta from first frame's projected root position, in first frame's heading frame.
+    Rotation: Relative rotation from first frame's heading quaternion (6D representation).
+
+    Returns:
+        if non_flatten:
+            [num_envs, num_future_frames, 1 * 9] (3 pos + 6 rot)
+        else:
+            [num_envs, num_future_frames * 1 * 9]
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    transforms = command.root_transforms_relative_to_first_frame
+    if non_flatten:
+        return transforms.reshape(env.num_envs, command.num_future_frames, -1)
+    else:
+        return transforms.reshape(env.num_envs, -1)
+
+
+def motion_anchor_ori_heading_mf(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Motion anchor orientation normalized by robot heading (yaw) only, for multi-future frames.
+
+    Unlike motion_anchor_ori_b_mf which normalizes by the robot's full orientation (including
+    pitch/roll), this version only normalizes by the robot's heading (yaw). This preserves
+    the reference motion's pitch/roll relative to gravity while removing the heading difference.
+
+    Returns:
+        torch.Tensor: Orientation as 6D rotation matrix representation (first 2 columns of
+            rotation matrix), shape (num_envs, num_future_frames, 6) if non_flatten else
+            (num_envs, num_future_frames * 6)
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    if non_flatten:
+        return command.root_rot_dif_heading_multi_future.reshape(
+            env.num_envs, command.num_future_frames, -1
+        )
+    else:
+        return command.root_rot_dif_heading_multi_future.reshape(env.num_envs, -1)
+
+
+def motion_anchor_ori_refheading_mf(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Motion anchor orientation canonicalized by the first target frame's heading.
+
+    Instead of using the robot's heading, this uses the heading of the first (immediate)
+    future frame from the reference motion. The trajectory is expressed in the reference
+    motion's own heading frame.
+
+    Returns:
+        torch.Tensor: Orientation as 6D rotation matrix representation,
+            shape (num_envs, num_future_frames, 6) if non_flatten else
+            (num_envs, num_future_frames * 6)
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    if non_flatten:
+        return command.root_rot_dif_refheading_multi_future.reshape(
+            env.num_envs, command.num_future_frames, -1
+        )
+    else:
+        return command.root_rot_dif_refheading_multi_future.reshape(env.num_envs, -1)
+
+
+def heading_diff_robot_ref(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Relative heading from robot to reference motion's first target frame.
+
+    Computes the rotation from the robot's heading to the reference motion's first
+    future frame heading as a single-frame 6D rotation.
+
+    Returns:
+        torch.Tensor: 6D rotation matrix representation, shape (num_envs, 6)
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    return command.heading_diff_robot_ref
+
+
+### 3 point force based tracking
+def vr_3point_target(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get VR 3-point tracking target positions in heading-local frame.
+
+    Transforms the reference 3-point body positions (left wrist, right wrist, head)
+    relative to the robot anchor, then de-heads (removes yaw) to get a
+    heading-invariant representation.
+
+    Returns:
+        torch.Tensor: Flattened 3-point positions, shape (num_envs, 9).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    pos_b = command.vr_3point_body_pos_w - command.robot_anchor_pos_w[:, None, :]
+
+    # transform pos_b in to deheaded root frame
+    root_quat = command.robot_anchor_quat_w.view(env.num_envs, 1, 4).repeat(
+        1, len(command.cfg.vr_3point_body), 1
+    )
+    deheaded_pos_b = quat_apply_yaw(quat_inv(root_quat), pos_b)
+    return deheaded_pos_b.view(env.num_envs, -1)
+
+
+def _phase_to_weight_pyramid(phase, start=0.25, end=0.75):
+    """Compute pyramid-shaped weight from a [0,1] phase value.
+
+    Ramps linearly from 0 to 1 over [0, start], holds at 1 over [start, end],
+    then ramps linearly from 1 to 0 over [end, 1].
+    """
+    weight = torch.zeros_like(phase)
+    mask1 = phase < start
+    mask2 = phase > end
+    weight[mask1] = 1 / start * phase[mask1]
+    weight[mask2] = 1 / (1 - end) * (1 - phase[mask2])
+    weight[~(mask1 | mask2)] = 1.0
+    return weight
+
+
+def vr_3point_target_compliant_multi_future(
+    env: ManagerBasedEnv, motion_command_name: str, force_command_name: str
+) -> torch.Tensor:
+    """Get force-compliant VR 3-point targets for multiple future frames in heading-local frame.
+
+    Modifies the reference 3-point positions by subtracting external force displacements
+    (scaled by compliance/stiffness) to create compliant tracking targets. The force
+    profile uses a pyramid-shaped phase weighting across future frames.
+
+    Returns:
+        torch.Tensor: Flattened compliant targets,
+            shape (num_envs, num_future_frames * num_points * 3).
+    """
+    motion_command: commands.TrackingCommand = env.command_manager.get_term(motion_command_name)
+    force_command: commands.ForceTrackingCommand = env.command_manager.get_term(force_command_name)
+
+    future_phases = (
+        (
+            force_command.force_push_counter[:, None]
+            + motion_command.future_time_steps_init
+            - force_command.force_update_frequency
+        )
+        / force_command.force_duration_per_env[:, None]
+    ).clamp(min=0.0, max=1.0)
+    # future_phases shape: (num_envs, num_future_frames)
+    future_forces = (
+        force_command.body_force_magnitude_buf[:, None, None, None]
+        * _phase_to_weight_pyramid(future_phases)[:, :, None, None]
+        * (force_command.body_force_dir_buf[:, force_command.force_push_ids_rel])[:, None, :, :]
+        * force_command.max_force
+    )
+
+    # future_forces [E, 5, 2, 3]
+    ext_force_disp = future_forces * force_command.eef_stiffness_buf[:, None, :, None]
+
+    mod_3point = motion_command.vr_3point_body_pos_w_multi_future
+    mod_3point -= ext_force_disp
+    pos_b = mod_3point - motion_command.robot_anchor_pos_w[:, None, None, :]
+    # transform pos_b in to deheaded root frame
+    root_quat = motion_command.robot_anchor_quat_w.view(env.num_envs, 1, 1, 4).repeat(
+        1, motion_command.num_future_frames, len(motion_command.cfg.vr_3point_body), 1
+    )
+    deheaded_pos_b = quat_apply_yaw(quat_inv(root_quat), pos_b)
+    return deheaded_pos_b.view(env.num_envs, -1)
+
+
+def vr_3point_target_multi_future(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get VR 3-point tracking target positions for multiple future frames in heading-local frame.
+
+    Returns:
+        torch.Tensor: Flattened multi-future 3-point positions,
+            shape (num_envs, num_future_frames * num_points * 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    pos_b = command.vr_3point_body_pos_w_multi_future - command.robot_anchor_pos_w[:, None, None, :]
+    # transform pos_b in to deheaded root frame
+    root_quat = command.robot_anchor_quat_w.view(env.num_envs, 1, 1, 4).repeat(
+        1, command.num_future_frames, len(command.cfg.vr_3point_body), 1
+    )
+    deheaded_pos_b = quat_apply_yaw(quat_inv(root_quat), pos_b)
+    return deheaded_pos_b.view(env.num_envs, -1)
+
+
+def head_orn_target_multi_future(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get reference head orientation for multiple future frames, de-headed by robot yaw.
+
+    Returns:
+        torch.Tensor: Flattened head quaternions,
+            shape (num_envs, num_future_frames * 4).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    head_quat = command.head_orn_w_multi_future
+    root_quat = command.robot_anchor_quat_w.view(env.num_envs, 1, 4).repeat(
+        1, command.num_future_frames, 1
+    )
+    # deheaded_head_quat = quat_mul(head_quat, get_heading_q(quat_inv(root_quat)))
+    deheaded_head_quat = quat_mul(torch_transform.get_heading_q(quat_inv(root_quat)), head_quat)
+    return deheaded_head_quat.view(env.num_envs, -1)
+
+
+def vr_3point_orn_target_multi_future(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get VR 3-point body orientations for multiple future frames, de-headed by robot yaw.
+
+    Returns:
+        torch.Tensor: Flattened body quaternions,
+            shape (num_envs, num_future_frames * num_points * 4).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    vr_3point_quat = command.vr_3point_body_quat_w_multi_future
+    root_quat = command.robot_anchor_quat_w.view(env.num_envs, 1, 1, 4).repeat(
+        1, command.num_future_frames, len(command.cfg.vr_3point_body), 1
+    )
+    # deheaded_vr_3point_quat = quat_mul(vr_3point_quat, get_heading_q(quat_inv(root_quat)))
+    deheaded_vr_3point_quat = quat_mul(
+        torch_transform.get_heading_q(quat_inv(root_quat)), vr_3point_quat
+    )
+    return deheaded_vr_3point_quat.view(env.num_envs, -1)
+
+
+def vr_3point_local_target(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get VR 3-point positions in reference motion anchor local frame.
+
+    NOTE: "local" here means relative to the reference motion root, not the robot root.
+    This gives the policy the reference pose structure independent of global position.
+
+    Returns:
+        torch.Tensor: Flattened 3-point positions, shape (num_envs, num_points * 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    ref_root_quat = command.anchor_quat_w.view(env.num_envs, 1, 4).repeat(
+        1, len(command.cfg.vr_3point_body), 1
+    )
+    ref_3point_diff = command.vr_3point_body_pos_w - command.anchor_pos_w[:, None, :]
+    ref_3point_root = quat_apply(quat_inv(ref_root_quat), ref_3point_diff)
+    return ref_3point_root.view(env.num_envs, -1)
+
+
+def vr_3point_local_target_multi_future(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get VR 3-point positions in reference anchor local frame for multiple future frames.
+
+    All future frames use the current reference anchor orientation for canonicalization.
+
+    Returns:
+        torch.Tensor: Flattened multi-future local positions,
+            shape (num_envs, num_future_frames * num_points * 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    ref_root_quat = command.anchor_quat_w.view(env.num_envs, 1, 1, 4).repeat(
+        1, command.num_future_frames, len(command.cfg.vr_3point_body), 1
+    )
+    ref_3point_diff = (
+        command.vr_3point_body_pos_w_multi_future - command.anchor_pos_w[:, None, None, :]
+    )  # All w.r.t current anchor
+    ref_3point_root = quat_apply(quat_inv(ref_root_quat), ref_3point_diff)
+    return ref_3point_root.view(env.num_envs, -1)
+
+
+def vr_3point_local_target_compliant(
+    env: ManagerBasedEnv,
+    motion_command_name: str,
+    force_command_name: str,
+    zero_out_head_position: bool = False,
+) -> torch.Tensor:
+    """Get force-compliant VR 3-point targets in reference anchor local frame.
+
+    Subtracts external force displacements (in robot-local frame) from the reference
+    3-point positions to create compliant tracking targets.
+
+    Args:
+        motion_command_name: Name of the motion tracking command.
+        force_command_name: Name of the force tracking command.
+        zero_out_head_position: If True, zero out the head position component.
+
+    Returns:
+        torch.Tensor: Flattened compliant targets, shape (num_envs, num_points * 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(motion_command_name)
+    force_command: commands.ForceTrackingCommand = env.command_manager.get_term(force_command_name)
+    ext_force_disp_w = (
+        force_command.last_force_applied * force_command.eef_stiffness_buf[:, :, None]
+    )  # delta x external force world frame
+    root_quat = command.robot_anchor_quat_w[:, None, :].repeat(
+        1, len(command.cfg.vr_3point_body), 1
+    )  # robot root pos not ref motion root pose as force is applied on robot, ref motion root frame is meaningless
+    ext_force_disp_l = quat_apply(
+        quat_inv(root_quat), ext_force_disp_w
+    )  # delta x external force local frame
+    ref_root_quat = command.anchor_quat_w.view(env.num_envs, 1, 4).repeat(
+        1, len(command.cfg.vr_3point_body), 1
+    )
+    ref_3point_diff = command.vr_3point_body_pos_w - command.anchor_pos_w[:, None, :]
+    ref_3point_root = quat_apply(quat_inv(ref_root_quat), ref_3point_diff)
+    ref_3point_root -= ext_force_disp_l  # [E, 3, 3] -- only a single frame VR 3 point is returned
+
+    # Conditionally zero out the head position (index -1)
+    if zero_out_head_position:
+        ref_3point_root[:, -1, :] = 0.0  # Head position
+
+    return ref_3point_root.view(env.num_envs, -1)
+
+
+def vr_3point_local_orn_target(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get VR 3-point body orientations in reference anchor local frame.
+
+    Returns:
+        torch.Tensor: Flattened 3-point quaternions, shape (num_envs, num_points * 4).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    ref_root_quat = command.anchor_quat_w.view(env.num_envs, 1, 4).repeat(
+        1, len(command.cfg.vr_3point_body), 1
+    )
+    ref_3point_quat = command.vr_3point_body_quat_w
+    ref_3point_root = quat_mul(quat_inv(ref_root_quat), ref_3point_quat)
+    return ref_3point_root.view(env.num_envs, -1)
+
+
+def vr_wrists_local_pos_target(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get wrist positions in reference anchor local frame (wrists only, no head).
+
+    Returns:
+        torch.Tensor: Flattened 2-wrist positions, shape (num_envs, 6).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    ref_root_quat = command.anchor_quat_w.view(env.num_envs, 1, 4).repeat(1, 2, 1)
+    ref_2point_diff = command.vr_3point_body_pos_w[:, :2] - command.anchor_pos_w[:, None, :]
+    ref_2point_root = quat_apply(quat_inv(ref_root_quat), ref_2point_diff)
+    return ref_2point_root.view(env.num_envs, -1)
+
+
+def vr_wrists_local_orn_target(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get wrist orientations in reference anchor local frame (wrists only, no head).
+
+    Returns:
+        torch.Tensor: Flattened 2-wrist quaternions, shape (num_envs, 8).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    ref_root_quat = command.anchor_quat_w.view(env.num_envs, 1, 4).repeat(1, 2, 1)
+    ref_2point_quat = command.vr_3point_body_quat_w[:, :2]
+    ref_2point_root = quat_mul(quat_inv(ref_root_quat), ref_2point_quat)
+    return ref_2point_root.view(env.num_envs, -1)
+
+
+def vr_head_local_orn_target(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get head orientation in reference anchor local frame.
+
+    Returns:
+        torch.Tensor: Head quaternion in anchor frame, shape (num_envs, 4).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    ref_root_quat = command.anchor_quat_w
+    head_quat = command.vr_3point_body_quat_w[:, 2]
+    ref_head_quat_root = quat_mul(quat_inv(ref_root_quat), head_quat)
+    return ref_head_quat_root.view(env.num_envs, -1)
+
+
+def get_command_obs(env: ManagerBasedEnv, command_name: str, obs_name: str) -> torch.Tensor:
+    """Get an arbitrary attribute from the tracking command by name.
+
+    Generic accessor for command observations not covered by dedicated functions.
+
+    Args:
+        command_name: Name of the tracking command term.
+        obs_name: Attribute name on the command object.
+
+    Returns:
+        torch.Tensor: The requested observation tensor.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    return getattr(command, obs_name)
+
+
+def smpl_pose_multi_future_select_joints(
+    env: ManagerBasedEnv, command_name: str, joints_idx: list
+) -> torch.Tensor:
+    """Extract SMPL joint axis-angle poses for selected joints across future frames.
+
+    Args:
+        command_name: Name of the tracking command term.
+        joints_idx: Indices of SMPL joints to select.
+
+    Returns:
+        torch.Tensor: Selected joint poses, shape (..., len(joints_idx) * 3).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    smpl_pose = command.smpl_pose_multi_future
+    smpl_pose_selected = smpl_pose.view(*smpl_pose.shape[:-1], -1, 3)[..., joints_idx, :].view(
+        *smpl_pose.shape[:-1], -1
+    )
+    return smpl_pose_selected
+
+
+def joint_pos_multi_future_select_joints(
+    env: ManagerBasedEnv, command_name: str, joints_idx: list, non_flatten=True
+) -> torch.Tensor:
+    """Extract reference joint positions for selected joints across future frames.
+
+    Args:
+        command_name: Name of the tracking command term.
+        joints_idx: Indices of joints to select from the full joint position vector.
+        non_flatten: If True, return shape (num_envs, num_future_frames, len(joints_idx)).
+
+    Returns:
+        torch.Tensor: Selected joint positions.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    joint_pos = command.joint_pos_multi_future
+    # Flatten case: (num_envs, num_future_frames * num_joints)
+    joint_pos_reshaped = joint_pos.view(env.num_envs, command.num_future_frames, -1)
+    joint_pos_selected = joint_pos_reshaped[..., joints_idx]
+    if not non_flatten:
+        joint_pos_selected = joint_pos_selected.view(env.num_envs, -1)
+    return joint_pos_selected
+
+
+def joint_pos_multi_future_select_joints_for_smpl(
+    env: ManagerBasedEnv, command_name: str, joints_idx: list
+) -> torch.Tensor:
+    """Extract reference joint positions for selected joints across SMPL-aligned future frames.
+
+    Uses smpl_num_future_frames (which may differ from num_future_frames) for
+    alignment with SMPL motion data.
+
+    Args:
+        command_name: Name of the tracking command term.
+        joints_idx: Indices of joints to select.
+
+    Returns:
+        torch.Tensor: Selected joint positions,
+            shape (num_envs, smpl_num_future_frames, len(joints_idx)).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    joint_pos = command.joint_pos_multi_future_for_smpl
+    # Flatten case: (num_envs, num_future_frames * num_joints)
+    joint_pos_reshaped = joint_pos.view(env.num_envs, command.smpl_num_future_frames, -1)
+    joint_pos_selected = joint_pos_reshaped[..., joints_idx]
+    return joint_pos_selected
+
+
+def joint_pos_multi_future_select_joints_for_smpl(  # noqa: F811
+    env: ManagerBasedEnv, command_name: str, joints_idx: list
+) -> torch.Tensor:
+    """Extract reference joint positions for selected joints across SMPL-aligned future frames.
+
+    NOTE: This is a duplicate definition that shadows the previous one. Only this
+    version is active at runtime.
+
+    Args:
+        command_name: Name of the tracking command term.
+        joints_idx: Indices of joints to select.
+
+    Returns:
+        torch.Tensor: Selected joint positions,
+            shape (num_envs, smpl_num_future_frames, len(joints_idx)).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    joint_pos = command.joint_pos_multi_future_for_smpl
+    # Flatten case: (num_envs, num_future_frames * num_joints)
+    joint_pos_reshaped = joint_pos.view(env.num_envs, command.smpl_num_future_frames, -1)
+    joint_pos_selected = joint_pos_reshaped[..., joints_idx]
+    return joint_pos_selected
+
+
+def smpl_elbow_wrist_pose_multi_future(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Extract SMPL elbow and wrist joint poses for multiple future frames.
+
+    Selects indices 54:66 from the SMPL pose vector, corresponding to
+    left elbow, left wrist, right elbow, right wrist (4 joints x 3 axis-angle).
+
+    Returns:
+        torch.Tensor: Elbow/wrist poses, shape (..., 12).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    smpl_pose = command.smpl_pose_multi_future
+    elbow_wrist_pose = smpl_pose[..., 54:66]
+    return elbow_wrist_pose
+
+
+def smpl_root_ori_b(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get SMPL root orientation relative to robot anchor as 6D rotation.
+
+    Returns:
+        torch.Tensor: 6D rotation representation, shape (num_envs, 6).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+
+    _, ori = subtract_frame_transforms(
+        command.robot_anchor_pos_w,
+        command.robot_anchor_quat_w,
+        None,
+        command.smpl_root_quat_w,
+    )
+    # diff = quat_mul(command.smpl_root_quat_w, quat_inv(command.anchor_quat_w))
+    mat = matrix_from_quat(ori)
+    return mat[..., :2].reshape(mat.shape[0], -1)
+
+
+def smpl_root_ori_b_mf(env: ManagerBasedEnv, command_name: str, non_flatten=True) -> torch.Tensor:
+    """Get SMPL root orientation difference in body-local frame for multiple future frames.
+
+    Args:
+        command_name: Name of the tracking command term.
+        non_flatten: If True, return shape (num_envs, smpl_num_future_frames, 6).
+
+    Returns:
+        torch.Tensor: 6D rotation differences.
+    """
+    # ZL: non-flatten set to true temporarily to match previous jobs, will be set to false in the future.
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    if non_flatten:
+        rot = command.smpl_root_quat_w_dif_l_multi_future.reshape(
+            env.num_envs, command.smpl_num_future_frames, -1
+        )
+    else:
+        rot = command.smpl_root_quat_w_dif_l_multi_future.reshape(env.num_envs, -1)
+    return rot
+
+
+def smpl_root_ori_refheading_mf(
+    env: ManagerBasedEnv, command_name: str, non_flatten=True
+) -> torch.Tensor:
+    """SMPL root orientation canonicalized by the first SMPL future frame's heading.
+
+    Uses the heading of the first SMPL future frame instead of the robot's orientation
+    for canonicalization.
+
+    Returns:
+        torch.Tensor: 6D rotation matrix representation,
+            shape (num_envs, smpl_num_future_frames, 6) if non_flatten else
+            (num_envs, smpl_num_future_frames * 6)
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    if non_flatten:
+        return command.smpl_root_quat_w_dif_refheading_multi_future.reshape(
+            env.num_envs, command.smpl_num_future_frames, -1
+        )
+    else:
+        return command.smpl_root_quat_w_dif_refheading_multi_future.reshape(env.num_envs, -1)
+
+
+def smpl_root_ori_heading_mf(
+    env: ManagerBasedEnv, command_name: str, non_flatten=True
+) -> torch.Tensor:
+    """SMPL root orientation canonicalized by robot heading (yaw) only.
+
+    Uses the robot's heading (yaw) instead of full orientation for canonicalization,
+    preserving the SMPL root's pitch/roll relative to gravity.
+
+    Returns:
+        torch.Tensor: 6D rotation matrix representation,
+            shape (num_envs, smpl_num_future_frames, 6) if non_flatten else
+            (num_envs, smpl_num_future_frames * 6)
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    if non_flatten:
+        return command.smpl_root_quat_w_dif_heading_multi_future.reshape(
+            env.num_envs, command.smpl_num_future_frames, -1
+        )
+    else:
+        return command.smpl_root_quat_w_dif_heading_multi_future.reshape(env.num_envs, -1)
+
+
+def smpl_joints_multi_future(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Get SMPL joint positions relative to SMPL root for multiple future frames.
+
+    Canonicalizes joint positions using the first SMPL frame's root orientation.
+
+    Args:
+        command_name: Name of the tracking command term.
+        non_flatten: If True, return shape (num_envs, smpl_num_future_frames, num_joints * 3).
+
+    Returns:
+        torch.Tensor: Root-relative SMPL joint positions.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    ref_joints = command.smpl_joints_multi_future
+    ref_root_quat = command.smpl_root_quat_w.view(env.num_envs, 1, 1, 4).repeat(
+        1, command.num_future_frames, ref_joints.shape[-2], 1
+    )
+    ref_joints_root = quat_apply(quat_inv(ref_root_quat), ref_joints)
+    if non_flatten:
+        return ref_joints_root.reshape(env.num_envs, command.smpl_num_future_frames, -1)
+    else:
+        return ref_joints_root.view(env.num_envs, -1)
+
+
+def smpl_joints_multi_future_local(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False, joints_idx=None
+) -> torch.Tensor:
+    """Get SMPL joint positions relative to each frame's own root orientation.
+
+    Unlike smpl_joints_multi_future which uses only the first frame's root,
+    this version canonicalizes each future frame independently using its own
+    root orientation. Optionally select a subset of joints.
+
+    Args:
+        command_name: Name of the tracking command term.
+        non_flatten: If True, return shape (num_envs, smpl_num_future_frames, num_joints * 3).
+        joints_idx: Optional list of joint indices to select.
+
+    Returns:
+        torch.Tensor: Per-frame root-relative SMPL joint positions.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+
+    ref_joints = command.smpl_joints_multi_future
+    ref_root_quat = command.smpl_root_quat_w_multi_future.unsqueeze(-2).repeat(
+        1, 1, ref_joints.shape[-2], 1
+    )
+    ref_joints_root = quat_apply(quat_inv(ref_root_quat), ref_joints)
+    if joints_idx is not None:
+        ref_joints_root = ref_joints_root[..., joints_idx, :]
+    if non_flatten:
+        return ref_joints_root.reshape(env.num_envs, command.smpl_num_future_frames, -1)
+    else:
+        return ref_joints_root.view(env.num_envs, -1)
+
+
+def smpl_joints_lower_multi_future_local(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Get lower-body SMPL joint positions relative to each frame's own root orientation.
+
+    Selects joints [0,1,2,4,5,7,8,10,11] (hips, knees, ankles) and canonicalizes
+    each future frame using its own root orientation.
+
+    Args:
+        command_name: Name of the tracking command term.
+        non_flatten: If True, return shape (num_envs, smpl_num_future_frames, 9 * 3).
+
+    Returns:
+        torch.Tensor: Per-frame root-relative lower-body SMPL joint positions.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+
+    ref_joints = command.smpl_joints_multi_future[
+        :, :, [0, 1, 2, 4, 5, 7, 8, 10, 11], :
+    ]  # only first 12 joints (lower body)
+    ref_root_quat = command.smpl_root_quat_w_multi_future.unsqueeze(-2).repeat(
+        1, 1, ref_joints.shape[-2], 1
+    )
+    ref_joints_root = quat_apply(quat_inv(ref_root_quat), ref_joints)
+    if non_flatten:
+        return ref_joints_root.reshape(env.num_envs, command.smpl_num_future_frames, -1)
+    else:
+        return ref_joints_root.view(env.num_envs, -1)
+
+
+def smpl_transl_z_multi_future(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """Get SMPL root translation z-height for multiple future frames.
+
+    Args:
+        command_name: Name of the tracking command term.
+        non_flatten: If True, preserve per-frame dimension.
+
+    Returns:
+        torch.Tensor: Z-heights, shape (num_envs, smpl_num_future_frames) when flattened.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    ref_transl_z = command.smpl_transl_z_multi_future
+    if non_flatten:
+        return ref_transl_z
+    else:
+        return ref_transl_z.view(env.num_envs, -1)
+
+
+def soma_joints_multi_future_local(
+    env: ManagerBasedEnv, command_name: str, non_flatten=False
+) -> torch.Tensor:
+    """SOMA skeleton joint positions canonicalized to each frame's root orientation.
+
+    Mirrors smpl_joints_multi_future_local but uses SOMA skeleton (26 joints).
+    Root quat has Y→Z up conversion and base rotation removal applied.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+
+    ref_joints = command.soma_joints_multi_future  # (batch, num_frames, 26, 3)
+    ref_root_quat = command.soma_root_quat_w_multi_future.unsqueeze(-2).repeat(
+        1, 1, ref_joints.shape[-2], 1  # (batch, num_frames, 26, 4)
+    )
+    ref_joints_root = quat_apply(quat_inv(ref_root_quat), ref_joints)
+    if non_flatten:
+        return ref_joints_root.reshape(env.num_envs, command.smpl_num_future_frames, -1)
+    else:
+        return ref_joints_root.view(env.num_envs, -1)
+
+
+def soma_root_ori_b_mf(env: ManagerBasedEnv, command_name: str, non_flatten=True):
+    """SOMA root orientation relative to robot anchor, as 6D rotation matrix.
+
+    Returns:
+        (num_envs, soma_num_future_frames, 6) if non_flatten
+        (num_envs, soma_num_future_frames * 6) otherwise
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    if non_flatten:
+        rot = command.soma_root_quat_w_dif_l_multi_future.reshape(
+            env.num_envs, command.smpl_num_future_frames, -1
+        )
+    else:
+        rot = command.soma_root_quat_w_dif_l_multi_future.reshape(env.num_envs, -1)
+    return rot
+
+
+def compliance(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Get end-effector compliance (stiffness) scaled for observation input.
+
+    Returns:
+        torch.Tensor: Scaled stiffness values, shape (num_envs, num_eef).
+    """
+    command: commands.ForceTrackingCommand = env.command_manager.get_term(command_name)
+    return command.eef_stiffness_buf * 10.0  # rescaling observation
+
+
+def ext_forces(
+    env: ManagerBasedEnv, force_command_name: str, motion_command_name: str
+) -> torch.Tensor:
+    """Get last-applied external forces in heading-local frame.
+
+    Transforms the world-frame external forces into a heading-invariant frame
+    by removing the robot's yaw rotation.
+
+    Returns:
+        torch.Tensor: Flattened forces, shape (num_envs, num_force_bodies * 3).
+    """
+    force_command: commands.ForceTrackingCommand = env.command_manager.get_term(force_command_name)
+    motion_command: commands.TrackingCommand = env.command_manager.get_term(motion_command_name)
+    ext_force_w = force_command.last_force_applied
+    root_quat = motion_command.robot_anchor_quat_w.view(env.num_envs, 1, 4).repeat(
+        1, ext_force_w.shape[1], 1
+    )
+    deheaded_force = quat_apply_yaw(quat_inv(root_quat), ext_force_w)
+    return deheaded_force.view(env.num_envs, -1)
+
+
+# =============================================================================
+# Body-only observation functions (for pre-trained action_transform_module)
+# Uses get_body_joint_indices from joint_utils.py
+
+
+def joint_pos_wo_hand(env: ManagerBasedEnv, asset_cfg) -> torch.Tensor:
+    """Get joint positions excluding hand joints (29 DOF body only)."""
+    asset = env.scene[asset_cfg.name]
+    body_indices = joint_utils.get_body_joint_indices(asset)
+    return asset.data.joint_pos[:, body_indices] - asset.data.default_joint_pos[:, body_indices]
+
+
+def joint_vel_wo_hand(env: ManagerBasedEnv, asset_cfg) -> torch.Tensor:
+    """Get joint velocities excluding hand joints (29 DOF body only)."""
+    asset = env.scene[asset_cfg.name]
+    body_indices = joint_utils.get_body_joint_indices(asset)
+    return asset.data.joint_vel[:, body_indices] - asset.data.default_joint_vel[:, body_indices]
+
+
+def last_action_wo_hand(env: ManagerBasedEnv, asset_cfg) -> torch.Tensor:
+    """Get last actions excluding hand joints."""
+    asset = env.scene[asset_cfg.name]
+    body_indices = joint_utils.get_body_joint_indices(asset)
+    return env.action_manager.action[:, body_indices]
+
+
+def last_meta_action(env: ManagerBasedEnv) -> torch.Tensor:
+    """Get last meta action (policy output: latent residual + finger primitives).
+
+    This returns the policy's output from the previous step, not the joint-level
+    actions that were applied to the simulation. For a student policy using
+    latent residual mode, this is typically:
+    - 64 dims: latent residual (tokenizer space)
+    - 2 dims: finger primitive actions (left + right hand)
+    Total: 66 dims
+
+    The buffer is stored on env by the ManagerEnvWrapper and initialized to zeros
+    on reset.
+    """
+    if hasattr(env, "_last_meta_action"):
+        return env._last_meta_action  # noqa: SLF001
+    else:
+        # Fallback: return zeros if buffer not initialized (shouldn't happen)
+        meta_action_dim = 66  # Default: 64 latent + 2 primitives
+        return torch.zeros(env.num_envs, meta_action_dim, dtype=torch.float32, device=env.device)
+
+
+def ref_root_pos_future_b(env, command_name: str, flatten: bool = False) -> torch.Tensor:
+    """Get reference root positions for future frames in robot anchor frame.
+
+    Args:
+        command_name: Name of the tracking command term.
+        flatten: If True, return shape (num_envs, num_future_frames * 3).
+
+    Returns:
+        torch.Tensor: Reference root positions,
+            shape (num_envs, num_future_frames, 3) or flattened.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    ref_root_pos_future_w = command.anchor_pos_w_multi_future.view(
+        command.num_envs, command.num_future_frames, -1
+    )
+    robot_root_pos_w = command.robot_anchor_pos_w.unsqueeze(1)
+    robot_root_quat_w = command.robot_anchor_quat_w.unsqueeze(1)
+    robot_root_quat_w = robot_root_quat_w.expand(-1, command.num_future_frames, -1)
+
+    ref_root_pos_future_b = quat_apply_inverse(
+        robot_root_quat_w, ref_root_pos_future_w - robot_root_pos_w
+    )
+    if flatten:
+        return ref_root_pos_future_b.reshape(command.num_envs, -1)
+    return ref_root_pos_future_b.reshape(command.num_envs, command.num_future_frames, -1)
+
+
+def ref_root_ori_future_b(env, command_name: str, flatten: bool = False) -> torch.Tensor:
+    """Get reference root orientations for future frames in robot anchor frame as 6D rotation.
+
+    Args:
+        command_name: Name of the tracking command term.
+        flatten: If True, return shape (num_envs, num_future_frames * 6).
+
+    Returns:
+        torch.Tensor: 6D rotation representations,
+            shape (num_envs, num_future_frames, 6) or flattened.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    ref_root_quat_future_w = command.anchor_quat_w_multi_future.view(
+        command.num_envs, command.num_future_frames, -1
+    )
+    robot_root_quat_w = command.robot_anchor_quat_w.unsqueeze(1)
+    robot_root_quat_w = robot_root_quat_w.expand(-1, command.num_future_frames, -1)
+
+    ref_root_quat_future_b = quat_mul(
+        quat_conjugate(robot_root_quat_w),
+        ref_root_quat_future_w,
+    )
+    ref_root_ori_future_b = matrix_from_quat(ref_root_quat_future_b)
+    ref_root_ori_future_b = ref_root_ori_future_b[:, :, :2, :]
+    if flatten:
+        return ref_root_ori_future_b.reshape(command.num_envs, -1)
+    return ref_root_ori_future_b.reshape(command.num_envs, command.num_future_frames, -1)
+
+
+def diff_body_pos_future_local(env, command_name: str, flatten: bool = False) -> torch.Tensor:
+    """Compute reference-minus-robot body position differences in their respective heading frames.
+
+    Reference body positions are expressed in the reference motion's heading (yaw-only,
+    z=0 projected) root frame; robot body positions are expressed in the robot's heading
+    root frame. The difference captures the tracking error per body per future frame.
+
+    Args:
+        command_name: Name of the tracking command term.
+        flatten: If True, return shape (num_envs, num_future_frames * num_bodies * 3).
+
+    Returns:
+        torch.Tensor: Position differences,
+            shape (num_envs, num_future_frames, num_bodies * 3) or flattened.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    ref_body_pos_future_w = command.body_pos_w_multi_future.view(
+        command.num_envs, command.num_future_frames, command.num_bodies, -1
+    )
+
+    ref_root_pos_w = command.anchor_pos_w.unsqueeze(1).unsqueeze(2)
+    # shape: (num_envs, 1, 1, 3)
+    ref_root_quat_w = command.anchor_quat_w.unsqueeze(1).unsqueeze(2)
+    # shape: (num_envs, 1, 1, 4)
+    ref_root_pos_w = ref_root_pos_w.clone()
+    ref_root_pos_w[..., 2] = 0.0
+    ref_root_quat_w = torch_transform.get_heading_q(ref_root_quat_w)
+    ref_root_quat_w = ref_root_quat_w.expand(
+        command.num_envs, command.num_future_frames, len(command.cfg.body_names), -1
+    )
+
+    robot_body_pos_w = command.robot_body_pos_w.view(command.num_envs, command.num_bodies, -1)
+
+    robot_root_pos_w = command.robot_anchor_pos_w.unsqueeze(1)
+    robot_root_quat_w = command.robot_anchor_quat_w.unsqueeze(1)
+    robot_root_pos_w = robot_root_pos_w.clone()
+    robot_root_pos_w[..., 2] = 0.0
+    robot_root_quat_w = torch_transform.get_heading_q(robot_root_quat_w)
+    robot_root_quat_w = robot_root_quat_w.expand(command.num_envs, len(command.cfg.body_names), -1)
+
+    robot_body_pos_local = quat_apply_inverse(
+        robot_root_quat_w, robot_body_pos_w - robot_root_pos_w
+    )
+    ref_body_pos_future_local = quat_apply_inverse(
+        ref_root_quat_w, ref_body_pos_future_w - ref_root_pos_w
+    )
+
+    diff = ref_body_pos_future_local - robot_body_pos_local.unsqueeze(1)
+    if flatten:
+        return diff.reshape(command.num_envs, -1)
+    return diff.reshape(command.num_envs, command.num_future_frames, -1)
+
+
+def diff_body_ori_future_local(env, command_name: str, flatten: bool = False) -> torch.Tensor:
+    """Compute reference-minus-robot body orientation differences in their respective heading frames.
+
+    Both reference and robot body orientations are first canonicalized by their respective
+    heading (yaw-only) root frames, then the relative rotation between them is computed
+    as a 6D rotation representation.
+
+    Args:
+        command_name: Name of the tracking command term.
+        flatten: If True, return shape (num_envs, num_future_frames * num_bodies * 6).
+
+    Returns:
+        torch.Tensor: 6D orientation differences,
+            shape (num_envs, num_future_frames, num_bodies * 6) or flattened.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    ref_body_quat_future_w = command.body_quat_w_multi_future.view(
+        command.num_envs, command.num_future_frames, command.num_bodies, -1
+    )
+
+    ref_root_quat_w = command.anchor_quat_w.unsqueeze(1).unsqueeze(2)
+    ref_root_quat_w = torch_transform.get_heading_q(ref_root_quat_w)
+    ref_root_quat_w = ref_root_quat_w.expand(
+        command.num_envs, command.num_future_frames, len(command.cfg.body_names), -1
+    )
+
+    robot_body_quat_w = command.robot_body_quat_w.view(command.num_envs, command.num_bodies, -1)
+
+    robot_root_quat_w = command.robot_anchor_quat_w.unsqueeze(1)
+    robot_root_quat_w = torch_transform.get_heading_q(robot_root_quat_w)
+    robot_root_quat_w = robot_root_quat_w.expand(command.num_envs, len(command.cfg.body_names), -1)
+
+    robot_body_quat_local = quat_mul(
+        quat_conjugate(robot_root_quat_w),
+        robot_body_quat_w,
+    ).unsqueeze(1)
+    ref_body_quat_future_local = quat_mul(
+        quat_conjugate(ref_root_quat_w),
+        ref_body_quat_future_w,
+    )
+
+    diff_body_quat_future = quat_mul(
+        quat_conjugate(robot_body_quat_local).expand_as(ref_body_quat_future_local),
+        ref_body_quat_future_local,
+    )
+    diff_body_ori_future_local = matrix_from_quat(diff_body_quat_future)
+    diff = diff_body_ori_future_local[:, :, :, :2, :]
+    if flatten:
+        return diff.reshape(command.num_envs, -1)
+    return diff.reshape(command.num_envs, command.num_future_frames, -1)
+
+
+def diff_body_lin_vel_future_local(env, command_name: str, flatten: bool = False) -> torch.Tensor:
+    """Compute reference-minus-robot body linear velocity differences in heading frames.
+
+    Both reference and robot velocities are expressed in their respective heading
+    (yaw-only) root frames. Output is clamped to [-25, 25] for numerical stability.
+
+    Args:
+        command_name: Name of the tracking command term.
+        flatten: If True, return shape (num_envs, num_future_frames * num_bodies * 3).
+
+    Returns:
+        torch.Tensor: Velocity differences,
+            shape (num_envs, num_future_frames, num_bodies * 3) or flattened.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    ref_body_lin_vel_future_w = command.motion_lib.get_body_lin_vel_w(
+        command.future_motion_ids, command.future_time_steps
+    ).view(command.num_envs, command.num_future_frames, -1, 3)
+
+    ref_root_quat_w = command.anchor_quat_w.unsqueeze(1).unsqueeze(2)
+    # shape: (num_envs, 1, 1, 4)
+    ref_root_quat_w = torch_transform.get_heading_q(ref_root_quat_w)
+    ref_root_quat_w = ref_root_quat_w.expand(
+        command.num_envs, command.num_future_frames, command.num_bodies, -1
+    )
+    ref_body_lin_vel_future_local = quat_apply_inverse(ref_root_quat_w, ref_body_lin_vel_future_w)
+
+    robot_body_lin_vel_w = command.robot_body_lin_vel_w
+    robot_root_quat_w = command.robot_anchor_quat_w.unsqueeze(1)
+    robot_root_quat_w = torch_transform.get_heading_q(robot_root_quat_w).expand(
+        command.num_envs, command.num_bodies, -1
+    )
+    robot_body_lin_vel_local = quat_apply_inverse(robot_root_quat_w, robot_body_lin_vel_w)
+
+    diff = ref_body_lin_vel_future_local - robot_body_lin_vel_local.unsqueeze(1)
+    diff.clamp_(min=-25, max=25)
+    if flatten:
+        return diff.reshape(command.num_envs, -1)
+    return diff.reshape(command.num_envs, command.num_future_frames, -1)
+
+
+def diff_body_ang_vel_future_local(env, command_name: str, flatten: bool = False) -> torch.Tensor:
+    """Compute reference-minus-robot body angular velocity differences in heading frames.
+
+    Both reference and robot angular velocities are expressed in their respective heading
+    (yaw-only) root frames. Output is clamped to [-25, 25] for numerical stability.
+
+    Args:
+        command_name: Name of the tracking command term.
+        flatten: If True, return shape (num_envs, num_future_frames * num_bodies * 3).
+
+    Returns:
+        torch.Tensor: Angular velocity differences,
+            shape (num_envs, num_future_frames, num_bodies * 3) or flattened.
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    ref_body_ang_vel_future_w = command.motion_lib.get_body_ang_vel_w(
+        command.future_motion_ids, command.future_time_steps
+    ).view(command.num_envs, command.num_future_frames, -1, 3)
+
+    ref_root_quat_w = command.anchor_quat_w.unsqueeze(1).unsqueeze(2)
+    # shape: (num_envs, 1, 1, 4)
+    ref_root_quat_w = torch_transform.get_heading_q(ref_root_quat_w)
+    ref_root_quat_w = ref_root_quat_w.expand(
+        command.num_envs, command.num_future_frames, command.num_bodies, -1
+    )
+    ref_body_ang_vel_future_local = quat_apply_inverse(ref_root_quat_w, ref_body_ang_vel_future_w)
+
+    robot_body_ang_vel_w = command.robot_body_ang_vel_w
+    robot_root_quat_w = command.robot_anchor_quat_w.unsqueeze(1)
+    robot_root_quat_w = torch_transform.get_heading_q(robot_root_quat_w).expand(
+        command.num_envs, command.num_bodies, -1
+    )
+    robot_body_ang_vel_local = quat_apply_inverse(robot_root_quat_w, robot_body_ang_vel_w)
+
+    diff = ref_body_ang_vel_future_local - robot_body_ang_vel_local.unsqueeze(1)
+    diff.clamp_(min=-25, max=25)
+    if flatten:
+        return diff.reshape(command.num_envs, -1)
+    return diff.reshape(command.num_envs, command.num_future_frames, -1)
+
+
+def height_map(env: ManagerBasedEnv, command_name, random=False) -> torch.Tensor:
+    """Get height map observation from the environment."""
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+
+    if not hasattr(command, "scan_dot_pos_w"):
+        # Height map disabled - return zeros with expected shape
+        n = int(command.cfg.height_map_size / command.cfg.height_map_resolution) + 1
+        return torch.zeros(command.num_envs, n, n, 3, device=command.device)
+
+    if not hasattr(command, "scan_dot_pos_w"):
+        # Height map disabled - return zeros with expected shape
+        n = int(command.cfg.height_map_size / command.cfg.height_map_resolution) + 1
+        return torch.zeros(command.num_envs, n, n, 3, device=command.device)
+
+    robot_root_pos_w, robot_root_quat_w = (
+        command.robot.data.root_pos_w,
+        command.robot.data.root_quat_w,
+    )
+    scan_dot_pos_w = command.scan_dot_pos_w
+
+    robot_root_quat_w_yaw = (
+        torch_transform.get_heading_q(robot_root_quat_w)
+        .unsqueeze(1)
+        .unsqueeze(2)
+        .expand(-1, command.num_rays_x, command.num_rays_y, -1)
+    )
+    obs = quat_apply_inverse(
+        robot_root_quat_w_yaw, scan_dot_pos_w - robot_root_pos_w.unsqueeze(1).unsqueeze(2)
+    )
+    obs.nan_to_num_(0.0)
+    if random:
+        obs = torch.rand_like(obs)
+    return obs
+
+
+def height_map_flat(env: ManagerBasedEnv, command_name="motion", random=False) -> torch.Tensor:
+    """Flattened height map for concatenation with 1D policy observations."""
+    hmap = height_map(env, command_name=command_name, random=random)
+    return hmap.reshape(hmap.shape[0], -1)  # (B, n*n*3)
+
+
+def residual_joint_pos_action(env: ManagerBasedEnv, command_name: str) -> torch.Tensor:
+    """Compute the action-space representation of reference joint positions.
+
+    Converts the reference motion's joint positions into the normalized action space
+    using the action manager's offset and scale.
+
+    Returns:
+        torch.Tensor: Normalized residual actions, shape (num_envs, num_action_joints).
+    """
+    command: commands.TrackingCommand = env.command_manager.get_term(command_name)
+    action_manager = env.action_manager.get_term("joint_pos")
+    action_offset = action_manager._offset  # noqa: SLF001
+    action_scale = action_manager._scale  # noqa: SLF001
+
+    motion_joint_pos = command.joint_pos
+    # motion_joint_names = command.robot.joint_names
+    # action_joint_names = action_manager._joint_names
+    action_joint_pos = motion_joint_pos[:, action_manager._joint_ids]  # noqa: SLF001
+    residual_action = (action_joint_pos - action_offset) / action_scale
+    return residual_action
+
+
+# =============================================================================
+# Task stage observation (for staged training)
+
+
+def get_task_stage(env: ManagerBasedEnv) -> torch.Tensor:
+    """Get current task stage."""
+    if hasattr(env, "task_stage"):
+        return env.task_stage.float().unsqueeze(-1)
+    return torch.zeros(env.num_envs, 1, dtype=torch.float, device=env.device)
+
+
+def get_tiled_camera_image(
+    env: ManagerBasedEnv,
+    camera_cfg: SceneEntityCfg,
+    normalize: bool = True,
+    normalize_mean: list[float] = [0.485, 0.456, 0.406],  # ImageNet mean (RGB)
+    normalize_std: list[float] = [0.229, 0.224, 0.225],  # ImageNet std (RGB)
+    debug_visualize: bool = False,  # Enable real-time visualization with cv2
+    debug_env_idx: int = 0,  # Which environment's image to visualize
+    debug_show_predicted: bool = False,  # Also show predicted object position (in red)
+) -> torch.Tensor:
+    """Get RGB image from tiled camera sensor.
+
+    Returns normalized RGB image of shape (num_envs, H, W, 3) - NOT flattened.
+    Image values can be normalized with configurable mean and std (default: ImageNet stats).
+
+    Args:
+        env: The environment object
+        camera_cfg: Camera configuration with sensor name
+        normalize: Whether to apply mean/std normalization (default: True)
+        normalize_mean: RGB mean values for normalization (default: ImageNet [0.485, 0.456, 0.406])
+        normalize_std: RGB std values for normalization (default: ImageNet [0.229, 0.224, 0.225])
+        debug_visualize: Whether to display the image in real-time using cv2 (default: False)
+        debug_env_idx: Which environment's image to visualize (default: 0)
+        debug_show_predicted: Whether to also show predicted object position in red (default: False)
+
+    Returns:
+        RGB image tensor of shape (num_envs, H, W, 3) - channels last format
+    """
+    camera = env.scene[camera_cfg.name]
+
+    # Get RGB image from camera - shape is (num_envs, H, W, 4) with RGBA
+    rgb_image = camera.data.output["rgb"]
+
+    # Take only RGB channels (drop alpha if present)
+    if rgb_image.shape[-1] == 4:
+        rgb_image = rgb_image[..., :3]
+
+    # Debug visualization (before normalization)
+    if debug_visualize:
+        utils.debug_visualize_object_projection(
+            env, camera, rgb_image, debug_env_idx, show_predicted=debug_show_predicted
+        )
+
+    # Convert to float if needed
+    if rgb_image.dtype != torch.float:
+        rgb_image = rgb_image.float()
+
+    # Normalize to [0, 1] if in [0, 255] range
+    if rgb_image.max() > 1.0:
+        rgb_image = rgb_image / 255.0
+
+    # Apply mean/std normalization if enabled
+    if normalize:
+        mean = torch.tensor(normalize_mean, device=rgb_image.device, dtype=rgb_image.dtype)
+        std = torch.tensor(normalize_std, device=rgb_image.device, dtype=rgb_image.dtype)
+        rgb_image = (rgb_image - mean) / std
+
+    # Return shape: (num_envs, H, W, 3) - do NOT flatten
+    # The vision encoder will handle permuting to (B, C, H, W) if needed
+    return rgb_image
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/recorders.py b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/recorders.py
new file mode 100644
index 0000000000000000000000000000000000000000..7f99e2ee84515dd1abbc6264dd9fcfbc1dcd8681
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/recorders.py
@@ -0,0 +1,393 @@
+"""Custom recorder terms for the manager environment MDP."""
+
+from __future__ import annotations
+
+import json
+import os
+import pickle
+from typing import TYPE_CHECKING
+
+import cv2
+import imageio
+from isaaclab.managers import manager_term_cfg, recorder_manager
+from isaaclab.utils import configclass
+from loguru import logger
+import numpy as np
+import torch
+from tqdm import tqdm
+
+if TYPE_CHECKING:
+    from isaaclab import envs
+
+
+@configclass
+class RecordersCfg(recorder_manager.RecorderManagerBaseCfg):
+    """Recorders terms for the MDP."""
+
+    render_envs = None
+    running_ref_root_height = None
+    trajectory = None
+
+
+class RenderEnvsRecorderTerm(recorder_manager.RecorderTerm):
+    """Recorder term for rendering environments with advanced features like text overlay and frame skipping."""
+
+    cfg: RenderEnvsRecorderCfg
+
+    def __init__(self, cfg: RenderEnvsRecorderCfg, env: envs.ManagerBasedEnv):
+        super().__init__(cfg, env)
+        self.cfg = cfg
+        self.env = env
+
+        # Determine save directory (backward compatibility)
+        self.save_dir = self.cfg.video_save_path
+        logger.info(f"=== Start recording video to {self.save_dir} ===")
+
+        # Create directory if it doesn't exist
+        os.makedirs(self.save_dir, exist_ok=True)
+        self.video_writers = []
+        self._writers_closed = False
+        self.frame_id = 0
+        self.first_render = True
+        self._fixed_eye = None
+        self._fixed_target = None
+
+    def _initialize_writers(self):
+        """Initialize video writers for each environment."""
+        logger.info(f"Saving rendering to {self.save_dir}")
+        # Get configuration parameters with defaults
+        self.group_camera = self.env.wrapper.config.get("group_camera", False)
+        self.max_render_envs = self.env.wrapper.config.get("max_render_envs", self.env.num_envs)
+        if self.group_camera:
+            self.max_render_envs = 1  # single video from overview_camera
+        self.render_frame_skip = self.env.wrapper.config.get("render_frame_skip", 2)
+        self.start_idx = self.env.wrapper.start_idx
+
+        for i in range(self.max_render_envs):
+            file_name = f"{self.save_dir}/{self.start_idx+i:06d}.mp4"
+            fps = 1 / (self.env.step_dt * self.render_frame_skip)
+            writer = imageio.get_writer(
+                file_name,
+                fps=fps,
+                codec="libx264",
+                quality=self.cfg.video_quality,
+                pixelformat="yuv420p",
+            )
+            self.video_writers.append(writer)
+
+    def record_post_step(self) -> tuple[str | None, torch.Tensor | dict | None]:
+        """Record video frames after each step with frame skipping and text overlay support."""
+        if len(self.video_writers) == 0:
+            self._initialize_writers()
+
+        # Check if we should render this frame
+        if self.frame_id % self.render_frame_skip != 0:
+            self.frame_id += 1
+            return "record_post_step", torch.ones(self.env.num_envs, 1, device=self.env.device)
+
+        # Set camera position based on robot root position
+        root_pos = self.env.command_manager.get_term("motion").robot_body_pos_w[:, 0]
+        camera_offset = self.env.wrapper.config.get("eval_camera_offset", [2, 2, 1])
+        fix_camera = self.env.wrapper.config.get("fix_camera_after_first_frame", False)
+        cam = self.env.scene["eval_camera"]
+
+        if fix_camera and self._fixed_eye is not None:
+            # Reuse the camera position from the first frame
+            eye, target = self._fixed_eye, self._fixed_target
+        elif self.group_camera:
+            center = root_pos.mean(dim=0, keepdim=True).expand_as(root_pos)
+            eye = center + torch.tensor(camera_offset, device=self.env.device)
+            target = center
+            if fix_camera:
+                self._fixed_eye = eye.clone()
+                self._fixed_target = center.clone()
+        else:
+            eye = root_pos + torch.tensor(camera_offset, device=self.env.device)
+            target = root_pos
+            if fix_camera:
+                self._fixed_eye = eye.clone()
+                self._fixed_target = root_pos.clone()
+
+        # Write world poses to Fabric AND sync to USD so both renderer paths see it
+        cam._view._sync_usd_on_fabric_write = True  # noqa: SLF001
+        cam.set_world_poses_from_view(eye, target)
+
+        # Two render calls: 1st flushes pose to render pipeline, 2nd captures at new pose
+        if hasattr(self.env, "sim"):
+            self.env.sim.render()
+            self.env.sim.render()
+
+        # Mark sensor as outdated so update actually re-reads the annotator buffers
+        cam._is_outdated[:] = True  # noqa: SLF001
+        cam.update(dt=0.0, force_recompute=True)
+
+        # Get RGB data
+        rgb_viewer = cam.data.output["rgb"].clone()
+
+        # Get render info if available
+        cur_render_info = None
+        if self.env.wrapper.config.get("render_info", None) is not None:
+            end_idx = self.start_idx + self.max_render_envs
+            cur_render_info = self.env.wrapper.config.render_info[self.start_idx : end_idx]
+
+        # Process each environment, loop over the video writers
+        if self.frame_id >= 1:
+            loop = (
+                tqdm(range(self.max_render_envs))
+                if self.first_render
+                else range(self.max_render_envs)
+            )
+            for i in loop:
+                frame = rgb_viewer[i].cpu().numpy()
+
+                # Add text overlay if render info is provided
+                if cur_render_info is not None and i < len(cur_render_info):
+                    for j, text in enumerate(cur_render_info[i]):
+                        frame = cv2.putText(
+                            frame,
+                            str(text),
+                            (10, 30 + j * 25),
+                            cv2.FONT_HERSHEY_SIMPLEX,
+                            0.6,
+                            (0, 0, 0),
+                            1,
+                        )
+
+                self.video_writers[i].append_data(frame)
+            self.first_render = False
+
+        self.frame_id += 1
+        return "record_post_step", torch.ones(self.env.num_envs, 1, device=self.env.device)
+
+    def close_writers(self):
+        """Explicitly close all video writers."""
+        if not self._writers_closed:
+            for i, writer in enumerate(self.video_writers):
+                try:
+                    writer.close()
+                    logger.info(f"Closed video writer {i}")
+                except Exception as e:  # noqa: BLE001
+                    logger.info(f"Error closing video writer {i}: {e}")
+            self.video_writers.clear()
+            self._writers_closed = True
+            self.frame_id = 0
+            self.first_render = True
+            self._fixed_eye = None
+            self._fixed_target = None
+            logger.info("=== All video writers closed ===")
+
+    def __del__(self):
+        """Ensure writers are closed when object is destroyed."""
+        self.close_writers()
+
+
+@configclass
+class RenderEnvsRecorderCfg(manager_term_cfg.RecorderTermCfg):
+    """Configuration for environment rendering recorder with advanced features."""
+
+    class_type = RenderEnvsRecorderTerm
+    video_save_path: str = None
+    video_quality: int = 5
+
+
+class TrajectoryRecorderTerm(recorder_manager.RecorderTerm):
+    """Recorder term that saves per-environment trajectory data (joint positions, root pose, object/table state).
+
+    Saves .trajectory.pkl files alongside the video output, enabling kinematic replay
+    in multi-scene composite renders.
+    """
+
+    cfg: TrajectoryRecorderCfg
+
+    def __init__(self, cfg: TrajectoryRecorderCfg, env: envs.ManagerBasedEnv):
+        super().__init__(cfg, env)
+        self.cfg = cfg
+        self.env = env
+
+        self.save_dir = self.cfg.save_path
+        os.makedirs(self.save_dir, exist_ok=True)
+        logger.info(f"=== TrajectoryRecorder: saving to {self.save_dir} ===")
+
+        self._initialized = False
+        self._closed = False
+        self._frame_data: dict[int, dict] = {}  # env_idx -> {field: [frames]}
+        self.frame_id = 0
+
+    def _initialize(self):
+        """Initialize per-env data buffers after environment is ready."""
+        self.num_record_envs = self.env.num_envs
+        self.start_idx = (
+            getattr(self.env.wrapper, "start_idx", 0) if hasattr(self.env, "wrapper") else 0
+        )
+
+        # Match video recorder's frame skip to keep trajectory in sync with video
+        if hasattr(self.env, "wrapper"):
+            self.render_frame_skip = self.env.wrapper.config.get("render_frame_skip", 2)
+        else:
+            self.render_frame_skip = 2
+
+        # Detect available scene entities
+        self._has_object = "object" in self.env.scene.rigid_objects
+        self._has_table = "table" in self.env.scene.rigid_objects
+
+        # Get motion command for root pose
+        try:
+            self._motion_cmd = self.env.command_manager.get_term("motion")
+        except Exception:  # noqa: BLE001
+            self._motion_cmd = None
+
+        for i in range(self.num_record_envs):
+            self._frame_data[i] = self._create_empty_data()
+
+        self._initialized = True
+
+    def _create_empty_data(self) -> dict:
+        data = {
+            "dof_pos": [],
+            "root_pos_w": [],
+            "root_quat_w": [],
+        }
+        if self._has_object:
+            data["object_pos_w"] = []
+            data["object_quat_w"] = []
+        if self._has_table:
+            data["table_pos_w"] = []
+            data["table_quat_w"] = []
+        return data
+
+    def record_post_step(self) -> tuple[str | None, torch.Tensor | dict | None]:
+        """Record trajectory state after each step, synced with video frame skip."""
+        if not self._initialized:
+            self._initialize()
+
+        # Skip frames to match video recorder cadence
+        if self.frame_id % self.render_frame_skip != 0:
+            self.frame_id += 1
+            return "trajectory_record", torch.ones(self.env.num_envs, 1, device=self.env.device)
+
+        robot = self.env.scene["robot"]
+        env_origins = self.env.scene.env_origins
+
+        for i in range(self.num_record_envs):
+            # Joint positions
+            joint_pos = robot.data.joint_pos[i].cpu().numpy().copy()
+            self._frame_data[i]["dof_pos"].append(joint_pos)
+
+            # Root position (relative to env origin)
+            if self._motion_cmd is not None:
+                root_pos = self._motion_cmd.robot_body_pos_w[i, 0].cpu().numpy().copy()
+            else:
+                root_pos = robot.data.root_pos_w[i].cpu().numpy().copy()
+            root_pos_rel = root_pos - env_origins[i].cpu().numpy()
+            self._frame_data[i]["root_pos_w"].append(root_pos_rel)
+
+            # Root quaternion (wxyz)
+            root_quat = robot.data.root_quat_w[i].cpu().numpy().copy()
+            self._frame_data[i]["root_quat_w"].append(root_quat)
+
+            # Object state
+            if self._has_object:
+                obj = self.env.scene["object"]
+                obj_pos = obj.data.root_pos_w[i].cpu().numpy().copy()
+                obj_pos_rel = obj_pos - env_origins[i].cpu().numpy()
+                obj_quat = obj.data.root_quat_w[i].cpu().numpy().copy()
+                self._frame_data[i]["object_pos_w"].append(obj_pos_rel)
+                self._frame_data[i]["object_quat_w"].append(obj_quat)
+
+            # Table state
+            if self._has_table:
+                table = self.env.scene["table"]
+                table_pos = table.data.root_pos_w[i].cpu().numpy().copy()
+                table_pos_rel = table_pos - env_origins[i].cpu().numpy()
+                table_quat = table.data.root_quat_w[i].cpu().numpy().copy()
+                self._frame_data[i]["table_pos_w"].append(table_pos_rel)
+                self._frame_data[i]["table_quat_w"].append(table_quat)
+
+        self.frame_id += 1
+        return "trajectory_record", torch.ones(self.env.num_envs, 1, device=self.env.device)
+
+    def close_writers(self):
+        """Save all trajectory data to pkl files."""
+        if self._closed or not self._initialized:
+            return
+        self._closed = True
+
+        # FPS matches the video (after frame skip)
+        effective_fps = 1.0 / (self.env.step_dt * self.render_frame_skip)
+
+        scene_metadata = {}
+
+        for i in range(self.num_record_envs):
+            env_idx = self.start_idx + i
+            data = self._frame_data[i]
+
+            if not data["dof_pos"]:
+                continue
+
+            # Stack frame arrays
+            trajectory = {
+                "dof_pos": np.array(data["dof_pos"]),
+                "root_pos_w": np.array(data["root_pos_w"]),
+                "root_quat_w": np.array(data["root_quat_w"]),
+                "quat_format": "wxyz",
+                "fps": effective_fps,
+                "num_joints": data["dof_pos"][0].shape[0],
+                "total_frames": len(data["dof_pos"]),
+            }
+
+            if data.get("object_pos_w"):
+                trajectory["object_pos_w"] = np.array(data["object_pos_w"])
+                trajectory["object_quat_w"] = np.array(data["object_quat_w"])
+            else:
+                trajectory["object_pos_w"] = None
+                trajectory["object_quat_w"] = None
+
+            if data.get("table_pos_w"):
+                trajectory["table_pos_w"] = np.array(data["table_pos_w"])
+                trajectory["table_quat_w"] = np.array(data["table_quat_w"])
+            else:
+                trajectory["table_pos_w"] = None
+                trajectory["table_quat_w"] = None
+
+            # Save pkl
+            pkl_path = os.path.join(self.save_dir, f"{env_idx:06d}.trajectory.pkl")
+            with open(pkl_path, "wb") as f:
+                pickle.dump(trajectory, f, protocol=pickle.HIGHEST_PROTOCOL)
+            logger.info(f"Saved trajectory: {pkl_path} ({trajectory['total_frames']} frames)")
+
+            # Build metadata entry
+            meta = {
+                "trajectory_file": f"{env_idx:06d}.trajectory.pkl",
+                "video_file": f"{env_idx:06d}.mp4",
+                "num_frames": trajectory["total_frames"],
+                "num_joints": trajectory["num_joints"],
+                "fps": effective_fps,
+                "has_object": trajectory["object_pos_w"] is not None,
+                "has_table": trajectory["table_pos_w"] is not None,
+            }
+
+            # Add object USD path if available from config
+            if hasattr(self.env, "wrapper"):
+                obj_usd = self.env.wrapper.config.get("object_usd_path", None)
+                if obj_usd:
+                    meta["object_usd_path"] = obj_usd
+
+            scene_metadata[str(env_idx)] = meta
+
+        # Save scene metadata JSON
+        meta_path = os.path.join(self.save_dir, "scene_metadata.json")
+        with open(meta_path, "w") as f:
+            json.dump(scene_metadata, f, indent=2)
+        logger.info(f"Saved scene metadata: {meta_path}")
+        logger.info("=== TrajectoryRecorder: all data saved ===")
+
+    def __del__(self):
+        self.close_writers()
+
+
+@configclass
+class TrajectoryRecorderCfg(manager_term_cfg.RecorderTermCfg):
+    """Configuration for trajectory recording alongside video."""
+
+    class_type = TrajectoryRecorderTerm
+    save_path: str = None
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/rewards.py b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/rewards.py
new file mode 100644
index 0000000000000000000000000000000000000000..2dbd614bf615128fac5bbc16251b3ad431bc496e
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/rewards.py
@@ -0,0 +1,607 @@
+"""Reward functions for the manager-based RL environment MDP."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from isaaclab.managers import SceneEntityCfg
+from isaaclab.utils import configclass
+from isaaclab.utils.math import (
+    quat_apply,
+    quat_error_magnitude,
+    quat_inv,
+    quat_mul,
+)
+import torch
+
+from gear_sonic.envs.manager_env.mdp.commands import (
+    ForceTrackingCommand,
+    TrackingCommand,
+    _get_body_indexes,
+)
+
+if TYPE_CHECKING:
+    from isaaclab.envs import ManagerBasedRLEnv
+
+
+@configclass
+class RewardsCfg:
+    """Reward terms for the MDP."""
+
+    tracking_anchor_pos = None
+    tracking_anchor_ori = None
+    tracking_relative_body_pos = None
+    tracking_relative_body_ori = None
+    tracking_relative_body_ori_weighted = None
+    tracking_body_linvel = None
+    tracking_body_angvel = None
+    action_rate_l2 = None
+    joint_limit = None
+    undesired_contacts = None
+    undesired_contacts_no_hands = None
+    undesired_contacts_no_ankle_hand = None
+    tracking_body_pos = None
+    tracking_body_ori = None
+    tracking_vr_3point_global = None
+    tracking_vr_3point_local = None
+    tracking_vr_3point_force = None
+    tracking_vr_2wrists_ori_tight = None
+    tracking_vr_2wrists_local_ori = None
+    tracking_head_local_ori = None
+    anti_shake_ang_vel = None
+    tracking_vr_5point_local = None
+    motion_5point_local_pos = None
+    feet_acc = None
+    energy_consumption = None
+    is_terminated = None
+    upright_penalty = None
+
+
+def tracking_anchor_pos_error(
+    env: ManagerBasedRLEnv, command_name: str, std: float
+) -> torch.Tensor:
+    """Compute anchor position tracking reward using a Gaussian kernel.
+
+    Encourages the robot's anchor (root) position to match the reference motion anchor.
+
+    Args:
+        env: The environment.
+        command_name: Name of the tracking command term.
+        std: Standard deviation for the Gaussian kernel. Smaller values produce
+            sharper falloff and stricter tracking.
+
+    Returns:
+        Reward tensor of shape (num_envs,) in [0, 1].
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    diff = command.anchor_pos_w - command.robot_anchor_pos_w
+    sq_dist = (diff * diff).sum(dim=-1)
+    return torch.exp(-sq_dist / (std * std))
+
+
+def tracking_anchor_ori_error(
+    env: ManagerBasedRLEnv, command_name: str, std: float
+) -> torch.Tensor:
+    """Compute anchor orientation tracking reward using a Gaussian kernel.
+
+    Encourages the robot's anchor (root) orientation to match the reference motion anchor.
+
+    Args:
+        env: The environment.
+        command_name: Name of the tracking command term.
+        std: Standard deviation for the Gaussian kernel on the angular error.
+
+    Returns:
+        Reward tensor of shape (num_envs,) in [0, 1].
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    angular_err = quat_error_magnitude(command.anchor_quat_w, command.robot_anchor_quat_w)
+    return torch.exp(-angular_err.square() / (std * std))
+
+
+def upright_penalty(
+    env: ManagerBasedRLEnv,
+    command_name: str,
+    body_name: str | None = None,
+    body_names: list[str] | None = None,
+) -> torch.Tensor:
+    """Penalize tilt of bodies away from upright.
+
+    Compute the squared magnitude of the x/y components of the gravity vector
+    in each body's local frame, summed across all specified bodies. When a body
+    is perfectly upright the local gravity is [0, 0, -1] and the penalty is 0.
+
+    Args:
+        env: The environment.
+        command_name: Name of the tracking command term.
+        body_name: Single body name (for backwards compatibility).
+        body_names: List of body names. If both are None, defaults to ["pelvis"].
+
+    Returns:
+        Penalty tensor of shape (num_envs,). Zero when upright, positive when tilted.
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    robot = env.scene["robot"]
+
+    if body_names is None:
+        body_names = [body_name] if body_name else ["pelvis"]
+
+    total_penalty = torch.zeros(env.num_envs, device=env.device)
+    for name in body_names:
+        body_idx = robot.body_names.index(name)
+        body_quat = robot.data.body_quat_w[:, body_idx]
+        g_local = quat_apply(quat_inv(body_quat), command.down_dir)
+        total_penalty += g_local[:, 0] ** 2 + g_local[:, 1] ** 2
+
+    return total_penalty
+
+
+def tracking_body_pos_error(
+    env: ManagerBasedRLEnv, command_name: str, std: float, body_names: list[str] | None = None
+) -> torch.Tensor:
+    """Compute body position tracking reward in world frame using a Gaussian kernel.
+
+    Encourages tracked body positions to match the reference motion. The reward is
+    the mean squared distance across all tracked bodies, passed through an exponential.
+
+    Args:
+        env: The environment.
+        command_name: Name of the tracking command term.
+        std: Standard deviation for the Gaussian kernel.
+        body_names: Subset of bodies to track. If None, uses all tracked bodies.
+
+    Returns:
+        Reward tensor of shape (num_envs,) in [0, 1].
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    tracked = _get_body_indexes(command, body_names)
+    pos_diff = command.body_pos_w[:, tracked] - command.robot_body_pos_w[:, tracked]
+    per_body_err = (pos_diff * pos_diff).sum(dim=-1)
+    return torch.exp(-per_body_err.mean(dim=-1) / (std * std))
+
+
+def tracking_vr_3point_error(env: ManagerBasedRLEnv, command_name: str, std: float):
+    """Compute VR 3-point tracking reward in world frame using a Gaussian kernel.
+
+    Encourages the robot's 3 VR tracking points (typically left wrist, right wrist,
+    head) to match their reference positions in world frame.
+
+    Args:
+        env: The environment.
+        command_name: Name of the tracking command term.
+        std: Standard deviation for the Gaussian kernel.
+
+    Returns:
+        Reward tensor of shape (num_envs,) in [0, 1].
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    pos_diff = command.robot_vr_3point_pos_w - command.vr_3point_body_pos_w
+    per_point_err = (pos_diff * pos_diff).sum(dim=-1)
+    return torch.exp(-per_point_err.mean(dim=-1) / (std * std))
+
+
+def tracking_vr_2wrists_ori_error(
+    env: ManagerBasedRLEnv, command_name: str, std: float, body_names: list[str] | None = None
+) -> torch.Tensor:
+    """Compute wrist orientation tracking reward in world frame.
+
+    Measure the orientation error of 2 wrist bodies against the reference motion,
+    similar to tracking_relative_body_ori_error but restricted to wrist links.
+
+    NOTE: The rigid extension defined in vr_3point_body_offset can be skipped for
+    orientation error since it does not affect rotations.
+
+    Args:
+        env: The environment.
+        command_name: Name of the tracking command term.
+        std: Standard deviation for the Gaussian kernel on the angular error.
+        body_names: List of wrist body names (must be provided).
+
+    Returns:
+        Reward tensor of shape (num_envs,) in [0, 1].
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    assert body_names is not None, "body_names must be provided"
+    tracked = _get_body_indexes(command, body_names)
+    angular_err = quat_error_magnitude(
+        command.body_quat_w[:, tracked], command.robot_body_quat_w[:, tracked]
+    )
+    return torch.exp(-angular_err.square().mean(dim=-1) / (std * std))
+
+
+def tracking_local_vr_2wrists_ori_error(
+    env: ManagerBasedRLEnv, command_name: str, std: float, body_names: list[str] | None = None
+) -> torch.Tensor:
+    """Compute wrist orientation tracking reward in the anchor's local frame.
+
+    Transform both reference and robot wrist orientations into the anchor (root)
+    frame before computing the angular error. This makes the reward invariant to
+    global root orientation.
+
+    Args:
+        env: The environment.
+        command_name: Name of the tracking command term.
+        std: Standard deviation for the Gaussian kernel on the angular error.
+        body_names: List of wrist body names (must be provided).
+
+    Returns:
+        Reward tensor of shape (num_envs,) in [0, 1].
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    assert body_names is not None, "body_names must be provided"
+    body_indexes = _get_body_indexes(command, body_names)
+    num_bodies = len(body_indexes)
+
+    # reference motion
+    ref_wrist_quat_w = command.body_quat_w[:, body_indexes]
+    ref_anchor_quat_w = command.anchor_quat_w.view(env.num_envs, 1, 4).repeat(1, num_bodies, 1)
+    ref_wrist_quat_local = quat_mul(quat_inv(ref_anchor_quat_w), ref_wrist_quat_w)
+
+    # robot
+    robot_wrist_quat_w = command.robot_body_quat_w[:, body_indexes]
+    robot_anchor_quat_w = command.robot_anchor_quat_w.view(env.num_envs, 1, 4).repeat(
+        1, num_bodies, 1
+    )
+    robot_wrist_quat_local = quat_mul(quat_inv(robot_anchor_quat_w), robot_wrist_quat_w)
+
+    error = quat_error_magnitude(ref_wrist_quat_local, robot_wrist_quat_local) ** 2
+    return torch.exp(-error.mean(-1) / std**2)
+
+
+def energy_consumption(
+    env: ManagerBasedRLEnv,
+    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
+) -> torch.Tensor:
+    """Penalize instantaneous mechanical power across the robot joints."""
+    if isinstance(asset_cfg, dict):
+        robot_name = asset_cfg.get("name", "robot")
+    else:
+        robot_name = getattr(asset_cfg, "name", "robot")
+    robot = env.scene[robot_name]
+    return torch.abs(robot.data.applied_torque * robot.data.joint_vel).sum(dim=-1)
+
+
+def tracking_local_head_ori_error(
+    env: ManagerBasedRLEnv, command_name: str, std: float
+) -> torch.Tensor:
+    """Compute head orientation tracking reward in the anchor's local frame.
+
+    Transform the head (torso_link) orientation into the anchor's local frame for
+    both the reference motion and the robot, then compute the angular error. This
+    encourages the robot to match the head-to-root relative orientation.
+
+    Args:
+        env: The environment.
+        command_name: Name of the tracking command term.
+        std: Standard deviation for the Gaussian kernel on the angular error.
+
+    Returns:
+        Reward tensor of shape (num_envs,) in [0, 1].
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+
+    # Get the head body index (torso_link)
+    head_body_names = ["torso_link"]
+    body_indexes = _get_body_indexes(command, head_body_names)
+
+    # reference motion: head orientation in world frame, transformed to anchor's local frame
+    ref_head_quat_w = command.body_quat_w[:, body_indexes]  # [num_envs, 1, 4]
+    ref_anchor_quat_w = command.anchor_quat_w.view(env.num_envs, 1, 4)
+    ref_head_quat_local = quat_mul(quat_inv(ref_anchor_quat_w), ref_head_quat_w)
+
+    # robot: head orientation in world frame, transformed to anchor's local frame
+    robot_head_quat_w = command.robot_body_quat_w[:, body_indexes]  # [num_envs, 1, 4]
+    robot_anchor_quat_w = command.robot_anchor_quat_w.view(env.num_envs, 1, 4)
+    robot_head_quat_local = quat_mul(quat_inv(robot_anchor_quat_w), robot_head_quat_w)
+
+    error = quat_error_magnitude(ref_head_quat_local, robot_head_quat_local) ** 2
+    return torch.exp(-error.squeeze(-1) / std**2)
+
+
+def tracking_local_vr_3point_error(
+    env: ManagerBasedRLEnv,
+    command_name: str,
+    std: float,
+    point_weights: list[float] | None = None,
+):
+    """Compute VR 3-point tracking reward in the anchor's local frame.
+
+    Transform tracking points into the anchor (root) local frame before computing
+    position error, making the reward invariant to global root position/orientation.
+    Supports optional per-point weighting.
+
+    Args:
+        env: The environment.
+        command_name: Name of the tracking command term.
+        std: Standard deviation for the Gaussian kernel.
+        point_weights: Optional weights for each tracking point. Order matches
+            vr_3point_body config (typically [left_wrist, right_wrist, head]).
+            If None, all points are weighted equally.
+            Example: [2, 2, 1] gives wrists 2x importance vs head.
+
+    Returns:
+        Reward tensor of shape (num_envs,) in [0, 1].
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    ref_3point_diff = command.vr_3point_body_pos_w - command.anchor_pos_w[:, None, :]
+    ref_root_quat = command.anchor_quat_w.view(env.num_envs, 1, 4).repeat(
+        1, len(command.cfg.vr_3point_body), 1
+    )
+    ref_3point_pos = quat_apply(quat_inv(ref_root_quat), ref_3point_diff)
+    robot_root_quat = command.robot_anchor_quat_w.view(env.num_envs, 1, 4).repeat(
+        1, len(command.cfg.vr_3point_body), 1
+    )
+    robot_3point_diff = command.robot_vr_3point_pos_w - command.robot_anchor_pos_w[:, None, :]
+    robot_3point_pos = quat_apply(quat_inv(robot_root_quat), robot_3point_diff)
+    diff = robot_3point_pos - ref_3point_pos
+    error = torch.sum(torch.square(diff), dim=-1)  # [num_envs, num_points]
+
+    if point_weights is not None:
+        # Weighted mean: sum(w_i * e_i) / sum(w_i)
+        weights = torch.tensor(point_weights, dtype=error.dtype, device=error.device)
+        weighted_error = (error * weights).sum(dim=-1) / weights.sum()
+    else:
+        # Simple mean (equal weights)
+        weighted_error = error.mean(dim=-1)
+
+    return torch.exp(-weighted_error / std**2)
+
+
+def tracking_local_vr_5point_error(env: ManagerBasedRLEnv, command_name: str, std: float):
+    """Compute VR 5-point tracking reward in the anchor's local frame.
+
+    Same approach as tracking_local_vr_3point_error but with 5 tracking points
+    (e.g., 2 wrists + head + 2 feet).
+
+    Args:
+        env: The environment.
+        command_name: Name of the tracking command term.
+        std: Standard deviation for the Gaussian kernel.
+
+    Returns:
+        Reward tensor of shape (num_envs,) in [0, 1].
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    ref_5point_diff = command.reward_point_body_pos_w - command.anchor_pos_w[:, None, :]
+    ref_root_quat = command.anchor_quat_w.view(env.num_envs, 1, 4).repeat(
+        1, len(command.cfg.reward_point_body), 1
+    )
+    ref_5point_pos = quat_apply(quat_inv(ref_root_quat), ref_5point_diff)
+    robot_root_quat = command.robot_anchor_quat_w.view(env.num_envs, 1, 4).repeat(
+        1, len(command.cfg.reward_point_body), 1
+    )
+    robot_5point_diff = (
+        command.robot_reward_point_body_pos_w - command.robot_anchor_pos_w[:, None, :]
+    )
+    robot_5point_pos = quat_apply(quat_inv(robot_root_quat), robot_5point_diff)
+    diff = robot_5point_pos - ref_5point_pos
+    error = torch.sum(torch.square(diff), dim=-1)
+    return torch.exp(-error.mean(-1) / std**2)
+
+
+def tracking_vr_3point_error_pos_force(
+    env: ManagerBasedRLEnv, motion_command_name: str, force_command_name: str, std: float
+):
+    """Compute VR 3-point tracking reward with force-based compliance correction.
+
+    Add a force-proportional offset to the wrist tracking error so that applied
+    external forces shift the tracking target, enabling compliant behavior under
+    force perturbations.
+
+    Args:
+        env: The environment.
+        motion_command_name: Name of the motion tracking command term.
+        force_command_name: Name of the force tracking command term.
+        std: Standard deviation for the Gaussian kernel.
+
+    Returns:
+        Reward tensor of shape (num_envs,) in [0, 1].
+    """
+    motion_command: TrackingCommand = env.command_manager.get_term(motion_command_name)
+    force_command: ForceTrackingCommand = env.command_manager.get_term(force_command_name)
+    diff = motion_command.robot_vr_3point_pos_w - motion_command.vr_3point_body_pos_w
+    force_error_wrists = (
+        force_command.last_force_applied * force_command.eef_stiffness_buf[:, :, None]
+    )
+    diff[:, :2] += force_error_wrists
+    error = torch.sum(torch.square(diff), dim=-1)
+    return torch.exp(-error.mean(-1) / std**2)
+
+
+def tracking_body_ori_error(
+    env: ManagerBasedRLEnv, command_name: str, std: float, body_names: list[str] | None = None
+) -> torch.Tensor:
+    """Compute body orientation tracking reward in world frame using a Gaussian kernel.
+
+    Encourages tracked body orientations to match the reference motion.
+
+    Args:
+        env: The environment.
+        command_name: Name of the tracking command term.
+        std: Standard deviation for the Gaussian kernel on the angular error.
+        body_names: Subset of bodies to track. If None, uses all tracked bodies.
+
+    Returns:
+        Reward tensor of shape (num_envs,) in [0, 1].
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    tracked = _get_body_indexes(command, body_names)
+    angular_err = quat_error_magnitude(
+        command.body_quat_w[:, tracked], command.robot_body_quat_w[:, tracked]
+    )
+    return torch.exp(-angular_err.square().mean(dim=-1) / (std * std))
+
+
+def tracking_relative_body_pos_error(
+    env: ManagerBasedRLEnv, command_name: str, std: float, body_names: list[str] | None = None
+) -> torch.Tensor:
+    """Compute body position tracking reward using anchor-relative reference positions.
+
+    Use reference body positions that have been shifted to share the robot's anchor
+    (root) position, so only the relative pose matters rather than absolute position.
+
+    Args:
+        env: The environment.
+        command_name: Name of the tracking command term.
+        std: Standard deviation for the Gaussian kernel.
+        body_names: Subset of bodies to track. If None, uses all tracked bodies.
+
+    Returns:
+        Reward tensor of shape (num_envs,) in [0, 1].
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    tracked = _get_body_indexes(command, body_names)
+    pos_diff = command.body_pos_relative_w[:, tracked] - command.robot_body_pos_w[:, tracked]
+    per_body_err = (pos_diff * pos_diff).sum(dim=-1)
+    return torch.exp(-per_body_err.mean(dim=-1) / (std * std))
+
+
+def tracking_relative_body_ori_error(
+    env: ManagerBasedRLEnv, command_name: str, std: float, body_names: list[str] | None = None
+) -> torch.Tensor:
+    """Compute body orientation tracking reward using anchor-relative reference orientations.
+
+    Use reference body orientations that have been transformed to share the robot's
+    anchor (root) orientation, making the reward invariant to global heading.
+
+    Args:
+        env: The environment.
+        command_name: Name of the tracking command term.
+        std: Standard deviation for the Gaussian kernel on the angular error.
+        body_names: Subset of bodies to track. If None, uses all tracked bodies.
+
+    Returns:
+        Reward tensor of shape (num_envs,) in [0, 1].
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    tracked = _get_body_indexes(command, body_names)
+    angular_err = quat_error_magnitude(
+        command.body_quat_relative_w[:, tracked],
+        command.robot_body_quat_w[:, tracked],
+    )
+    return torch.exp(-angular_err.square().mean(dim=-1) / (std * std))
+
+
+def tracking_relative_body_ori_weighted_error(
+    env: ManagerBasedRLEnv,
+    command_name: str,
+    std: float,
+    body_names: list[str] | None = None,
+    body_weights: dict[str, float] | None = None,
+) -> torch.Tensor:
+    """Compute anchor-relative body orientation tracking reward with per-body weights.
+
+    Same as tracking_relative_body_ori_error but allows different bodies to contribute
+    differently to the mean error. Useful for relaxing tracking on certain joints
+    (e.g., wrists during manipulation).
+
+    Args:
+        env: The environment.
+        command_name: Name of the tracking command term.
+        std: Standard deviation for the Gaussian kernel on the angular error.
+        body_names: Subset of bodies to track. If None, uses all tracked bodies.
+        body_weights: Dict mapping body name to weight multiplier. Bodies not listed
+            default to 1.0. E.g. {"left_wrist_yaw_link": 0.1} to relax wrist tracking.
+
+    Returns:
+        Reward tensor of shape (num_envs,) in [0, 1].
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    body_indexes = _get_body_indexes(command, body_names)
+    error = (
+        quat_error_magnitude(
+            command.body_quat_relative_w[:, body_indexes],
+            command.robot_body_quat_w[:, body_indexes],
+        )
+        ** 2
+    )
+    if body_weights is not None:
+        tracked_names = [command.cfg.body_names[i] for i in body_indexes]
+        weights = torch.tensor(
+            [body_weights.get(name, 1.0) for name in tracked_names],
+            device=error.device,
+            dtype=error.dtype,
+        )
+        weighted_error = (error * weights).sum(-1) / weights.sum()
+    else:
+        weighted_error = error.mean(-1)
+    return torch.exp(-weighted_error / std**2)
+
+
+def tracking_body_linvel_error(
+    env: ManagerBasedRLEnv, command_name: str, std: float, body_names: list[str] | None = None
+) -> torch.Tensor:
+    """Compute body linear velocity tracking reward using a Gaussian kernel.
+
+    Encourages tracked body linear velocities to match the reference motion.
+
+    Args:
+        env: The environment.
+        command_name: Name of the tracking command term.
+        std: Standard deviation for the Gaussian kernel.
+        body_names: Subset of bodies to track. If None, uses all tracked bodies.
+
+    Returns:
+        Reward tensor of shape (num_envs,) in [0, 1].
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    tracked = _get_body_indexes(command, body_names)
+    vel_diff = command.body_lin_vel_w[:, tracked] - command.robot_body_lin_vel_w[:, tracked]
+    per_body_err = (vel_diff * vel_diff).sum(dim=-1)
+    return torch.exp(-per_body_err.mean(dim=-1) / (std * std))
+
+
+def tracking_body_angvel_error(
+    env: ManagerBasedRLEnv, command_name: str, std: float, body_names: list[str] | None = None
+) -> torch.Tensor:
+    """Compute body angular velocity tracking reward using a Gaussian kernel.
+
+    Encourages tracked body angular velocities to match the reference motion.
+
+    Args:
+        env: The environment.
+        command_name: Name of the tracking command term.
+        std: Standard deviation for the Gaussian kernel.
+        body_names: Subset of bodies to track. If None, uses all tracked bodies.
+
+    Returns:
+        Reward tensor of shape (num_envs,) in [0, 1].
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    tracked = _get_body_indexes(command, body_names)
+    vel_diff = command.body_ang_vel_w[:, tracked] - command.robot_body_ang_vel_w[:, tracked]
+    per_body_err = (vel_diff * vel_diff).sum(dim=-1)
+    return torch.exp(-per_body_err.mean(dim=-1) / (std * std))
+
+
+def anti_shake_ang_vel_l2(
+    env: ManagerBasedRLEnv,
+    command_name: str,
+    threshold: float = 1.5,
+    body_names: list[str] | None = None,
+) -> torch.Tensor:
+    """Penalize excessive angular velocity on selected bodies with a deadzone.
+
+    Discourage high-frequency jitter on small links (wrists, head) while allowing
+    normal intentional motion within the threshold. Speeds below the threshold
+    incur zero penalty.
+
+    Args:
+        env: The environment.
+        command_name: Name of the tracking command term.
+        threshold: Angular velocity deadzone (rad/s). No penalty below this.
+        body_names: Bodies to penalize. If None, uses all tracked bodies.
+
+    Returns:
+        Penalty tensor of shape (num_envs,). Positive values (use negative weight).
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    body_indexes = _get_body_indexes(command, body_names)
+    # [E, B, 3]
+    ang_vel = command.robot_body_ang_vel_w[:, body_indexes]
+    # magnitude per body: [E, B]
+    speed = torch.linalg.norm(ang_vel, dim=-1)
+    # deadzone then square: [E, B]
+    excess = torch.relu(speed - threshold)
+    penalty = (excess * excess).mean(dim=-1)
+    return penalty
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/terminations.py b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/terminations.py
new file mode 100644
index 0000000000000000000000000000000000000000..55671b6dbd890f3f6d5f7d17d5deaa75e369fe96
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/terminations.py
@@ -0,0 +1,489 @@
+"""Termination conditions for motion-tracking and HOI reinforcement learning environments."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+import re
+from typing import TYPE_CHECKING
+
+import isaaclab.utils.math as math_utils
+import torch
+
+if TYPE_CHECKING:
+    from isaaclab.envs import ManagerBasedRLEnv
+
+from isaaclab.assets import Articulation, RigidObject
+from isaaclab.managers import ManagerTermBase, SceneEntityCfg, TerminationTermCfg
+from isaaclab.utils import configclass
+from isaaclab.utils.math import (
+    axis_angle_from_quat,
+    quat_apply_inverse,
+    quat_conjugate,
+    quat_error_magnitude,
+    quat_mul,
+)
+
+from gear_sonic.envs.manager_env.mdp.commands import TrackingCommand, _get_body_indexes
+from gear_sonic.trl.utils.torch_transform import get_heading_q
+
+
+@configclass
+class TerminationsCfg:
+    """Termination terms for the MDP."""
+
+    time_out = None
+    anchor_pos = None
+    anchor_ori = None
+    anchor_ori_full = None
+    ee_body_pos = None
+    anchor_pos_xy = None
+    foot_pos_xyz = None
+    cumm_body_pos_error = None
+    cumm_body_ori_error = None
+    cumm_body_pos_error_local = None
+    cumm_body_ori_error_local = None
+    # HOI terminations
+    object_pos_deviation = None
+    object_z_pos_deviation = None
+    object_not_lifted = None
+    robot_table_contact_before_object = None
+    hand_table_contact_termination = None
+    grasp_failure_after_contact = None
+
+
+def exceeded_anchor_pos(
+    env: ManagerBasedRLEnv, command_name: str, threshold: float
+) -> torch.Tensor:
+    """Terminate if anchor (root) position error exceeds a distance threshold.
+
+    Compute the L2 norm between the reference and robot anchor positions in world
+    frame and flag environments where the error exceeds ``threshold``.
+
+    Args:
+        env: The manager-based RL environment.
+        command_name: Name of the tracking command term.
+        threshold: Maximum allowed position error in meters.
+
+    Returns:
+        Boolean tensor of shape ``(num_envs,)``.
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    pos_diff = command.anchor_pos_w - command.robot_anchor_pos_w
+    return pos_diff.norm(dim=1).gt(threshold)
+
+
+def exceeded_anchor_pos_xy(
+    env: ManagerBasedRLEnv, command_name: str, threshold: float, body_names: list[str] | None = None
+) -> torch.Tensor:
+    """Terminate if anchor XY (horizontal) position error exceeds a distance threshold.
+
+    Only the X and Y components of the world-frame anchor position are compared,
+    ignoring vertical drift.
+
+    Args:
+        env: The manager-based RL environment.
+        command_name: Name of the tracking command term.
+        threshold: Maximum allowed horizontal position error in meters.
+        body_names: Unused; kept for config compatibility.
+
+    Returns:
+        Boolean tensor of shape ``(num_envs,)``.
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    xy_diff = command.anchor_pos_w[:, :2] - command.robot_anchor_pos_w[:, :2]
+    return xy_diff.norm(dim=1).gt(threshold)
+
+
+def exceeded_anchor_height(
+    env: ManagerBasedRLEnv,
+    command_name: str,
+    threshold: float,
+    threshold_adaptive: bool = False,
+    down_threshold: float = 0.5,
+    root_height_threshold: float = 1.0,
+) -> torch.Tensor:
+    """Terminate if anchor Z-height error exceeds a threshold.
+
+    When ``threshold_adaptive`` is True, use a looser ``down_threshold`` for
+    environments whose reference root height is below ``root_height_threshold``
+    (e.g. crouching or sitting motions).
+
+    Args:
+        env: The manager-based RL environment.
+        command_name: Name of the tracking command term.
+        threshold: Default maximum allowed height error in meters.
+        threshold_adaptive: Enable per-env adaptive thresholding.
+        down_threshold: Looser threshold applied when the reference root is low.
+        root_height_threshold: Height below which ``down_threshold`` is used.
+
+    Returns:
+        Boolean tensor of shape ``(num_envs,)``.
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    height_diff = (command.anchor_pos_w[:, 2] - command.robot_anchor_pos_w[:, 2]).abs()
+    if threshold_adaptive:
+        thresh = torch.full_like(height_diff, threshold)
+        thresh[command.running_ref_root_height < root_height_threshold] = down_threshold
+        return height_diff.gt(thresh)
+    return height_diff.gt(threshold)
+
+
+def exceeded_anchor_tilt(
+    env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, command_name: str, threshold: float
+) -> torch.Tensor:
+    """Terminate if the anchor tilt deviates from the reference gravity projection.
+
+    Compare the Z-component of the gravity vector rotated into the reference and
+    robot anchor frames. A large difference indicates excessive torso tilt.
+
+    Args:
+        env: The manager-based RL environment.
+        asset_cfg: Scene entity config used to obtain the gravity vector.
+        command_name: Name of the tracking command term.
+        threshold: Maximum allowed absolute difference in projected gravity Z.
+
+    Returns:
+        Boolean tensor of shape ``(num_envs,)``.
+    """
+    asset: RigidObject | Articulation = env.scene[asset_cfg.name]
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    quat_apply_fn = (
+        math_utils.quat_apply_inverse
+        if hasattr(math_utils, "quat_apply_inverse")
+        else math_utils.quat_rotate_inverse
+    )
+    ref_grav = quat_apply_fn(command.anchor_quat_w, asset.data.GRAVITY_VEC_W)
+    robot_grav = quat_apply_fn(command.robot_anchor_quat_w, asset.data.GRAVITY_VEC_W)
+    return (ref_grav[:, 2] - robot_grav[:, 2]).abs().gt(threshold)
+
+
+def exceeded_anchor_ori(
+    env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, command_name: str, threshold: float
+) -> torch.Tensor:
+    """Terminate if the squared anchor orientation error exceeds a threshold.
+
+    Compute the full quaternion error magnitude between reference and robot
+    anchor orientations, then compare the squared value against ``threshold``.
+
+    Args:
+        env: The manager-based RL environment.
+        asset_cfg: Scene entity config (unused but required by termination API).
+        command_name: Name of the tracking command term.
+        threshold: Maximum allowed squared orientation error (radians^2).
+
+    Returns:
+        Boolean tensor of shape ``(num_envs,)``.
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    angular_err = quat_error_magnitude(command.anchor_quat_w, command.robot_anchor_quat_w)
+    return angular_err.square().gt(threshold)
+
+
+def exceeded_body_pos(
+    env: ManagerBasedRLEnv, command_name: str, threshold: float, body_names: list[str] | None = None
+) -> torch.Tensor:
+    """Terminate if any tracked body position error exceeds a distance threshold.
+
+    Compare per-body world-frame positions between the reference motion and the
+    robot, and terminate if *any* body exceeds ``threshold``.
+
+    Args:
+        env: The manager-based RL environment.
+        command_name: Name of the tracking command term.
+        threshold: Maximum allowed position error per body in meters.
+        body_names: Optional list of body names to check. If ``None``, all
+            tracked bodies from the command are used.
+
+    Returns:
+        Boolean tensor of shape ``(num_envs,)``.
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    tracked = _get_body_indexes(command, body_names)
+    pos_diff = command.body_pos_relative_w[:, tracked] - command.robot_body_pos_w[:, tracked]
+    return pos_diff.norm(dim=-1).gt(threshold).any(dim=-1)
+
+
+def exceeded_body_height(
+    env: ManagerBasedRLEnv,
+    command_name: str,
+    threshold: float,
+    threshold_adaptive: bool = False,
+    down_threshold: float = 0.5,
+    body_names: list[str] | None = None,
+    root_height_threshold: float = 0.5,
+) -> torch.Tensor:
+    """Terminate if any tracked body Z-height error exceeds a threshold.
+
+    When ``threshold_adaptive`` is True, use a looser ``down_threshold`` for
+    environments whose reference root height is below ``root_height_threshold``.
+
+    Args:
+        env: The manager-based RL environment.
+        command_name: Name of the tracking command term.
+        threshold: Default maximum allowed height error per body in meters.
+        threshold_adaptive: Enable per-env adaptive thresholding.
+        down_threshold: Looser threshold applied when the reference root is low.
+        body_names: Optional list of body names to check. If ``None``, all
+            tracked bodies from the command are used.
+        root_height_threshold: Height below which ``down_threshold`` is used.
+
+    Returns:
+        Boolean tensor of shape ``(num_envs,)``.
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    tracked = _get_body_indexes(command, body_names)
+    height_err = (
+        command.body_pos_relative_w[:, tracked, 2] - command.robot_body_pos_w[:, tracked, 2]
+    ).abs()
+    if threshold_adaptive:
+        thresh = torch.full_like(height_err, threshold)
+        thresh[command.running_ref_root_height < root_height_threshold] = down_threshold
+        return height_err.gt(thresh).any(dim=-1)
+    return height_err.gt(threshold).any(dim=-1)
+
+
+def tracking_time_out(env: ManagerBasedRLEnv, command_name: str) -> torch.Tensor:
+    """Terminate when the motion clip has been fully played.
+
+    Compare the elapsed simulation steps (including the motion start offset)
+    against the total length of each environment's assigned motion clip.
+
+    Args:
+        env: The manager-based RL environment.
+        command_name: Name of the tracking command term.
+
+    Returns:
+        Boolean tensor of shape ``(num_envs,)``.
+    """
+    command: TrackingCommand = env.command_manager.get_term(command_name)
+    elapsed = command.time_steps + command.motion_start_time_steps + 1
+    total = command.motion_lib.get_time_step_total(command.motion_ids)
+    return elapsed >= total
+
+
+def _resolve_matching_names(patterns: str | Sequence[str], candidates: Sequence[str]) -> list[str]:
+    """Resolve body name patterns (exact or regex) against available names."""
+    if isinstance(patterns, str | bytes):
+        patterns = [patterns]
+    resolved: list[str] = []
+    for pattern in patterns:
+        if pattern == ".*":
+            resolved.extend(candidates)
+            continue
+        regex = re.compile(pattern)
+        resolved.extend([name for name in candidates if regex.fullmatch(name)])
+    # preserve order and drop duplicates
+    seen = set()
+    ordered_unique = []
+    for name in resolved:
+        if name not in seen:
+            ordered_unique.append(name)
+            seen.add(name)
+    return ordered_unique
+
+
+class _CummErrorMixin(ManagerTermBase):
+    """Shared logic for cumulative error-based terminations."""
+
+    def __init__(self, cfg: TerminationTermCfg, env):
+        """Initialize cumulative error tracking buffers.
+
+        Args:
+            cfg: Termination term config. Expected ``params`` keys:
+                ``min_steps``, ``threshold``, ``command_name``.
+            env: The manager-based RL environment.
+        """
+        super().__init__(cfg=cfg, env=env)
+        self.min_steps: int = cfg.params.get("min_steps")
+        self.threshold: float = cfg.params.get("threshold")
+        self.command_name = cfg.params.get("command_name")
+        self.command: TrackingCommand = env.command_manager.get_term(self.command_name)
+
+        device = self.command.device
+        self.error = torch.zeros(env.num_envs, device=device)
+        self._cum_steps = torch.zeros(env.num_envs, dtype=torch.int32, device=device)
+
+    def _update_counters(self) -> torch.Tensor:
+        """Accumulate consecutive steps above the threshold and return done mask."""
+        exceeded = self.error >= self.threshold
+        self._cum_steps[exceeded] += 1
+        self._cum_steps[~exceeded] = 0
+        return self._cum_steps >= self.min_steps
+
+    def reset(self, env_ids: Sequence[int] | None = None):
+        """Reset cumulative step counters for the given environments.
+
+        Args:
+            env_ids: Environment indices to reset, or ``None`` for all.
+        """
+        if env_ids is None:
+            env_ids = slice(None)
+        self._cum_steps[env_ids] = 0
+
+
+class CummBodyPosError(_CummErrorMixin):
+    """Terminate if body position error in world frame exceeds threshold for min_steps."""
+
+    def __init__(self, cfg: TerminationTermCfg, env):
+        """Initialize body index mapping for world-frame position error tracking.
+
+        Args:
+            cfg: Termination term config. Additional ``params`` key:
+                ``body_names`` -- regex or list of body names to track
+                (default ``".*"`` for all).
+            env: The manager-based RL environment.
+        """
+        super().__init__(cfg=cfg, env=env)
+        motion_names = self.command.cfg.body_names
+        body_names = cfg.params.get("body_names", ".*")
+        selected = _resolve_matching_names(body_names, motion_names)
+        self.motion_body_indices = [motion_names.index(name) for name in selected]
+        # command.body_indexes maps motion body order to robot body indices
+        self.robot_body_indices = self.motion_body_indices
+
+    def __call__(self, env, body_names=None, min_steps=None, threshold=None, command_name=None):
+        """Compute max world-frame body position error and check cumulative threshold.
+
+        Returns:
+            Boolean tensor of shape ``(num_envs,)``.
+        """
+        ref_body_pos = self.command.body_pos_w[:, self.motion_body_indices]
+        robot_body_pos = self.command.robot_body_pos_w[:, self.robot_body_indices]
+        body_pos_error = (ref_body_pos - robot_body_pos).norm(dim=-1)
+        self.error[:] = body_pos_error.max(dim=1).values
+        return self._update_counters()
+
+
+class CummBodyOriError(_CummErrorMixin):
+    """Terminate if body orientation error in world frame exceeds threshold for min_steps."""
+
+    def __init__(self, cfg: TerminationTermCfg, env):
+        """Initialize body index mapping for world-frame orientation error tracking.
+
+        Args:
+            cfg: Termination term config. Additional ``params`` key:
+                ``body_names`` -- regex or list of body names to track
+                (default ``".*"`` for all).
+            env: The manager-based RL environment.
+        """
+        super().__init__(cfg=cfg, env=env)
+        motion_names = self.command.cfg.body_names
+        body_names = cfg.params.get("body_names", ".*")
+        selected = _resolve_matching_names(body_names, motion_names)
+        self.motion_body_indices = [motion_names.index(name) for name in selected]
+        # self.robot_body_indices = self.command.body_indexes[self.motion_body_indices].tolist()
+        self.robot_body_indices = self.motion_body_indices
+
+    def __call__(self, env, body_names=None, min_steps=None, threshold=None, command_name=None):
+        """Compute max world-frame body orientation error and check cumulative threshold.
+
+        Returns:
+            Boolean tensor of shape ``(num_envs,)``.
+        """
+        ref_body_quat = self.command.body_quat_w[:, self.motion_body_indices]
+        robot_body_quat = self.command.robot_body_quat_w[:, self.robot_body_indices]
+        quat_diff = quat_mul(quat_conjugate(ref_body_quat), robot_body_quat)
+        body_ori_error = axis_angle_from_quat(quat_diff).norm(dim=-1)
+        self.error[:] = body_ori_error.max(dim=1).values
+        return self._update_counters()
+
+
+class CummBodyPosErrorLocal(_CummErrorMixin):
+    """Terminate if body position error in the root-yaw frame exceeds threshold for min_steps."""
+
+    def __init__(self, cfg, env):
+        """Initialize body index mapping for root-local position error tracking.
+
+        Args:
+            cfg: Termination term config. Additional ``params`` key:
+                ``body_names`` -- regex or list of body names to track
+                (default ``".*"`` for all).
+            env: The manager-based RL environment.
+        """
+        super().__init__(cfg=cfg, env=env)
+        body_names = cfg.params.get("body_names", ".*")
+        motion_names = self.command.cfg.body_names
+        selected = _resolve_matching_names(body_names, motion_names)
+        self.motion_body_indices = [motion_names.index(name) for name in selected]
+        # self.robot_body_indices = self.command.body_indexes[self.motion_body_indices].tolist()
+        self.robot_body_indices = self.motion_body_indices
+
+    def __call__(self, env, body_names=None, min_steps=None, threshold=None, command_name=None):
+        """Compute max root-yaw-local body position error and check cumulative threshold.
+
+        Transform body positions into the heading-aligned root frame before
+        computing errors, making the check invariant to global position and yaw.
+
+        Returns:
+            Boolean tensor of shape ``(num_envs,)``.
+        """
+        ref_body_pos = self.command.body_pos_w[:, self.motion_body_indices]
+        robot_body_pos = self.command.robot_body_pos_w[:, self.robot_body_indices]
+
+        ref_root_pos = self.command.anchor_pos_w.view(self.num_envs, 1, 3).clone()
+        robot_root_pos = self.command.robot_anchor_pos_w.view(self.num_envs, 1, 3).clone()
+        ref_root_pos[..., 2] = 0.0
+        robot_root_pos[..., 2] = 0.0
+
+        ref_root_quat = get_heading_q(self.command.anchor_quat_w.view(self.num_envs, 1, 4))
+        robot_root_quat = get_heading_q(self.command.robot_anchor_quat_w.view(self.num_envs, 1, 4))
+        # expand root quaternions to match per-body vectors for quat_apply_inverse
+        ref_root_quat = ref_root_quat.expand(ref_body_pos.shape[0], ref_body_pos.shape[1], -1)
+        robot_root_quat = robot_root_quat.expand(
+            robot_body_pos.shape[0], robot_body_pos.shape[1], -1
+        )
+
+        ref_body_local = quat_apply_inverse(ref_root_quat, ref_body_pos - ref_root_pos)
+        robot_body_local = quat_apply_inverse(robot_root_quat, robot_body_pos - robot_root_pos)
+
+        body_pos_error = (ref_body_local - robot_body_local).norm(dim=-1)
+        self.error[:] = body_pos_error.max(dim=1).values
+        return self._update_counters()
+
+
+class CummBodyOriErrorLocal(_CummErrorMixin):
+    """Terminate if body orientation error in the root-yaw frame exceeds threshold for min_steps."""
+
+    def __init__(self, cfg: TerminationTermCfg, env):
+        """Initialize body index mapping for root-local orientation error tracking.
+
+        Args:
+            cfg: Termination term config. Additional ``params`` key:
+                ``body_names`` -- regex or list of body names to track
+                (default ``".*"`` for all).
+            env: The manager-based RL environment.
+        """
+        super().__init__(cfg=cfg, env=env)
+        motion_names = self.command.cfg.body_names
+        body_names = cfg.params.get("body_names", ".*")
+        selected = _resolve_matching_names(body_names, motion_names)
+        self.motion_body_indices = [motion_names.index(name) for name in selected]
+        # self.robot_body_indices = self.command.body_indexes[self.motion_body_indices].tolist()
+        self.robot_body_indices = self.motion_body_indices
+
+    def __call__(self, env, body_names=None, min_steps=None, threshold=None, command_name=None):
+        """Compute max root-yaw-local body orientation error and check cumulative threshold.
+
+        Transform body quaternions into the heading-aligned root frame before
+        computing axis-angle errors, making the check invariant to global yaw.
+
+        Returns:
+            Boolean tensor of shape ``(num_envs,)``.
+        """
+        ref_body_quat = self.command.body_quat_w[:, self.motion_body_indices]
+        robot_body_quat = self.command.robot_body_quat_w[:, self.robot_body_indices]
+
+        ref_root_quat = get_heading_q(self.command.anchor_quat_w.view(self.num_envs, 1, 4))
+        robot_root_quat = get_heading_q(self.command.robot_anchor_quat_w.view(self.num_envs, 1, 4))
+        ref_root_quat = ref_root_quat.expand_as(ref_body_quat)
+        robot_root_quat = robot_root_quat.expand_as(robot_body_quat)
+
+        ref_body_quat_local = quat_mul(quat_conjugate(ref_root_quat), ref_body_quat)
+        robot_body_quat_local = quat_mul(quat_conjugate(robot_root_quat), robot_body_quat)
+
+        quat_diff = quat_mul(quat_conjugate(ref_body_quat_local), robot_body_quat_local)
+        body_ori_error = axis_angle_from_quat(quat_diff).norm(dim=-1)
+        self.error[:] = body_ori_error.max(dim=1).values
+        return self._update_counters()
+
+
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/terrain.py b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/terrain.py
new file mode 100644
index 0000000000000000000000000000000000000000..c5e25f75a46377e7206b8fed3f249b7631ba7019
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/terrain.py
@@ -0,0 +1,51 @@
+# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Configuration for custom terrains."""
+
+import isaaclab.terrains as terrain_gen
+from isaaclab.terrains.terrain_generator_cfg import TerrainGeneratorCfg
+
+ROUGH_TERRAINS_CFG = TerrainGeneratorCfg(
+    size=(8.0, 8.0),
+    border_width=20.0,
+    num_rows=20,
+    num_cols=20,
+    horizontal_scale=0.1,
+    vertical_scale=0.005,
+    slope_threshold=0.75,
+    use_cache=False,
+    sub_terrains={
+        # "pyramid_stairs": terrain_gen.MeshPyramidStairsTerrainCfg(
+        #     proportion=0.2,
+        #     step_height_range=(0.05, 0.23),
+        #     step_width=0.3,
+        #     platform_width=3.0,
+        #     border_width=1.0,
+        #     holes=False,
+        # ),
+        # "pyramid_stairs_inv": terrain_gen.MeshInvertedPyramidStairsTerrainCfg(
+        #     proportion=0.2,
+        #     step_height_range=(0.05, 0.23),
+        #     step_width=0.3,
+        #     platform_width=3.0,
+        #     border_width=1.0,
+        #     holes=False,
+        # ),
+        "boxes": terrain_gen.MeshRandomGridTerrainCfg(
+            proportion=0.3, grid_width=0.45, grid_height_range=(0.001, 0.005), platform_width=2.0
+        ),
+        "random_rough": terrain_gen.HfRandomUniformTerrainCfg(
+            proportion=0.05, noise_range=(0.001, 0.005), noise_step=0.02, border_width=0.25
+        ),
+        # "hf_pyramid_slope": terrain_gen.HfPyramidSlopedTerrainCfg(
+        #     proportion=0.1, slope_range=(0.0, 0.4), platform_width=2.0, border_width=0.25
+        # ),
+        # "hf_pyramid_slope_inv": terrain_gen.HfInvertedPyramidSlopedTerrainCfg(
+        #     proportion=0.1, slope_range=(0.0, 0.4), platform_width=2.0, border_width=0.25
+        # ),
+    },
+)
+"""Rough terrains configuration."""
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/utils.py b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..dd7ad6b5c8433a4c4eb4b982b38aa1f3cafcf580
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/mdp/utils.py
@@ -0,0 +1,258 @@
+"""
+Utility functions for manager environment MDP, including debug visualization.
+"""
+
+import torch
+
+
+def debug_visualize_object_projection(
+    env, camera, rgb_image, debug_env_idx: int = 0, show_predicted: bool = False
+):
+    """
+    Debug visualization: Project ground truth and predicted object positions onto camera image.
+
+    This function computes the camera pose from the robot's d435_link and camera offset,
+    transforms the object position to the camera frame, and projects it to 2D pixel coordinates.
+
+    The coordinate transformation pipeline:
+    1. Get object position in robot base (pelvis) frame
+    2. Get d435_link pose in robot base frame
+    3. Apply camera offset (from config) + random extrinsics delta
+    4. Transform object from base frame to camera frame
+    5. Map camera axes to OpenCV convention:
+       - x_cv = -camera_Y  (camera left → OpenCV right)
+       - y_cv = -camera_Z  (camera up → OpenCV down)
+       - z_cv = camera_X   (camera forward → OpenCV depth)
+    6. Project using pinhole camera model: u = fx*x/z + cx, v = fy*y/z + cy
+
+    Args:
+        env: The environment object
+        camera: The TiledCamera sensor
+        rgb_image: RGB image tensor [num_envs, H, W, 3]
+        debug_env_idx: Which environment to visualize
+        show_predicted: Whether to also visualize predicted object position (in red)
+    """
+    import cv2
+    from isaaclab.utils.math import (
+        quat_apply,
+        quat_apply_inverse,
+        quat_conjugate,
+        quat_from_euler_xyz,
+        quat_mul,
+    )
+
+    # Get the image for the specified environment
+    debug_img = rgb_image[debug_env_idx].detach().cpu().numpy()
+    if debug_img.max() <= 1.0:
+        debug_img = (debug_img * 255).astype("uint8")
+    else:
+        debug_img = debug_img.astype("uint8")
+    debug_img_bgr = cv2.cvtColor(debug_img, cv2.COLOR_RGB2BGR)
+    H, W = debug_img_bgr.shape[:2]
+
+    try:
+        # Find object in scene
+        object_name = None
+        if hasattr(env.scene, "rigid_objects"):
+            for name in env.scene.rigid_objects.keys():
+                if "object" in name.lower() or "obj" in name.lower():
+                    object_name = name
+                    break
+
+        if object_name is not None and "d435_link" in env.scene["robot"].body_names:
+            robot = env.scene["robot"]
+            object_asset = env.scene[object_name]
+            device = robot.data.root_pos_w.device
+
+            # Step 1: Object position in robot base frame
+            obj_pos_w = object_asset.data.root_pos_w[debug_env_idx]
+            robot_pos_w = robot.data.root_pos_w[debug_env_idx]
+            robot_quat_w = robot.data.root_quat_w[debug_env_idx]
+            obj_rel_w = obj_pos_w - robot_pos_w
+            obj_pos_base = quat_apply_inverse(robot_quat_w.unsqueeze(0), obj_rel_w.unsqueeze(0))[0]
+
+            # Step 2: d435_link pose in robot base frame
+            link_idx = robot.body_names.index("d435_link")
+            link_pos_w = robot.data.body_link_pos_w[debug_env_idx, link_idx]
+            link_quat_w = robot.data.body_link_quat_w[debug_env_idx, link_idx]
+            link_rel_w = link_pos_w - robot_pos_w
+            link_pos_base = quat_apply_inverse(robot_quat_w.unsqueeze(0), link_rel_w.unsqueeze(0))[
+                0
+            ]
+            link_quat_base = quat_mul(
+                quat_conjugate(robot_quat_w.unsqueeze(0)), link_quat_w.unsqueeze(0)
+            )[0]
+
+            # Step 3: Camera extrinsics (offset from d435_link)
+            base_pos_offset = [0.0, 0.0, 0.0]
+            base_rot_offset_quat = [1.0, 0.0, 0.0, 0.0]
+            if hasattr(camera, "cfg") and hasattr(camera.cfg, "offset"):
+                offset_cfg = camera.cfg.offset
+                if hasattr(offset_cfg, "pos"):
+                    base_pos_offset = list(offset_cfg.pos)
+                if hasattr(offset_cfg, "rot"):
+                    base_rot_offset_quat = list(offset_cfg.rot)
+
+            # Random extrinsics delta from wrapper
+            pos_delta = [0.0, 0.0, 0.0]
+            rot_delta = [0.0, 0.0, 0.0]
+            if hasattr(env, "wrapper") and hasattr(env.wrapper, "_camera_random_deltas"):
+                deltas = env.wrapper._camera_random_deltas.get(debug_env_idx, None)
+                if deltas is not None:
+                    pos_delta = deltas["pos_delta"]
+                    rot_delta = [deltas["roll_delta"], deltas["pitch_delta"], deltas["yaw_delta"]]
+
+            cam_offset_pos = torch.tensor(
+                [
+                    base_pos_offset[0] + pos_delta[0],
+                    base_pos_offset[1] + pos_delta[1],
+                    base_pos_offset[2] + pos_delta[2],
+                ],
+                device=device,
+                dtype=torch.float32,
+            )
+            base_rot_quat = torch.tensor(base_rot_offset_quat, device=device, dtype=torch.float32)
+
+            # Step 4: Camera pose in robot base frame
+            cam_offset_world_base = quat_apply(
+                link_quat_base.unsqueeze(0), cam_offset_pos.unsqueeze(0)
+            )[0]
+            cam_pos_base = link_pos_base + cam_offset_world_base
+            cam_quat_base = quat_mul(link_quat_base.unsqueeze(0), base_rot_quat.unsqueeze(0))[0]
+            if any(r != 0 for r in rot_delta):
+                delta_quat = quat_from_euler_xyz(
+                    torch.tensor([rot_delta[0]], device=device),
+                    torch.tensor([rot_delta[1]], device=device),
+                    torch.tensor([rot_delta[2]], device=device),
+                )[0]
+                cam_quat_base = quat_mul(cam_quat_base.unsqueeze(0), delta_quat.unsqueeze(0))[0]
+
+            # Step 5: Object in camera frame
+            obj_rel_cam = obj_pos_base - cam_pos_base
+            obj_pos_cam = quat_apply_inverse(cam_quat_base.unsqueeze(0), obj_rel_cam.unsqueeze(0))[
+                0
+            ]
+
+            # Step 6: Map to OpenCV convention
+            # Camera frame: X=forward, Y=left, Z=up
+            # OpenCV frame: X=right, Y=down, Z=depth
+            x_cv = -obj_pos_cam[1].item()  # -Y (left→right)
+            y_cv = -obj_pos_cam[2].item()  # -Z (up→down)
+            z_cv = obj_pos_cam[0].item()  # X (forward→depth)
+
+            # Step 7: Project using intrinsics
+            K = camera.data.intrinsic_matrices[debug_env_idx].cpu().numpy()
+            fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2]
+
+            if z_cv > 0.01:  # Object in front of camera
+                u = int(fx * x_cv / z_cv + cx)
+                v = int(fy * y_cv / z_cv + cy)
+
+                if 0 <= u < W and 0 <= v < H:
+                    cv2.circle(debug_img_bgr, (u, v), 12, (0, 255, 0), 2)
+                    cv2.circle(debug_img_bgr, (u, v), 4, (0, 255, 0), -1)
+                    cv2.putText(
+                        debug_img_bgr,
+                        f"GT d={z_cv:.2f}m",
+                        (u + 15, v),
+                        cv2.FONT_HERSHEY_SIMPLEX,
+                        0.5,
+                        (0, 255, 0),
+                        1,
+                    )
+                else:
+                    u_c = max(10, min(W - 10, u))
+                    v_c = max(10, min(H - 10, v))
+                    cv2.circle(debug_img_bgr, (u_c, v_c), 8, (0, 0, 255), -1)
+                    cv2.putText(
+                        debug_img_bgr,
+                        "OOB",
+                        (u_c + 10, v_c),
+                        cv2.FONT_HERSHEY_SIMPLEX,
+                        0.4,
+                        (0, 0, 255),
+                        1,
+                    )
+            else:
+                cv2.putText(
+                    debug_img_bgr,
+                    "Object behind camera",
+                    (10, 30),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    0.6,
+                    (0, 0, 255),
+                    2,
+                )
+
+            # Visualize predicted object position (in RED) if available
+            # NOTE: pred_pos is in robot BASE (pelvis) frame, same as GT during training
+            if show_predicted and hasattr(env, "wrapper"):
+                pred_pos = env.wrapper._last_predicted_object_pos
+                if pred_pos is not None and pred_pos.shape[0] > debug_env_idx:
+                    # pred_pos is in robot BASE (pelvis) frame - need to transform to camera frame
+                    pred_pos_base = pred_pos[debug_env_idx]
+
+                    # Transform from base frame to camera frame (same as GT)
+                    pred_rel_cam = pred_pos_base - cam_pos_base
+                    pred_pos_cam = quat_apply_inverse(
+                        cam_quat_base.unsqueeze(0), pred_rel_cam.unsqueeze(0)
+                    )[0]
+
+                    # Map to OpenCV convention (same as GT)
+                    pred_x_cv = -pred_pos_cam[1].item()  # -Y (left→right)
+                    pred_y_cv = -pred_pos_cam[2].item()  # -Z (up→down)
+                    pred_z_cv = pred_pos_cam[0].item()  # X (forward→depth)
+
+                    if pred_z_cv > 0.01:  # Predicted object in front of camera
+                        pred_u = int(fx * pred_x_cv / pred_z_cv + cx)
+                        pred_v = int(fy * pred_y_cv / pred_z_cv + cy)
+
+                        if 0 <= pred_u < W and 0 <= pred_v < H:
+                            # Draw RED circle for predicted position
+                            cv2.circle(debug_img_bgr, (pred_u, pred_v), 12, (0, 0, 255), 2)
+                            cv2.circle(debug_img_bgr, (pred_u, pred_v), 4, (0, 0, 255), -1)
+                            cv2.putText(
+                                debug_img_bgr,
+                                f"PRED d={pred_z_cv:.2f}m",
+                                (pred_u + 15, pred_v + 20),
+                                cv2.FONT_HERSHEY_SIMPLEX,
+                                0.5,
+                                (0, 0, 255),
+                                1,
+                            )
+                        else:
+                            # Out of bounds - draw at edge
+                            pred_u_c = max(10, min(W - 10, pred_u))
+                            pred_v_c = max(10, min(H - 10, pred_v))
+                            cv2.circle(debug_img_bgr, (pred_u_c, pred_v_c), 8, (255, 0, 255), -1)
+                            cv2.putText(
+                                debug_img_bgr,
+                                "PRED OOB",
+                                (pred_u_c + 10, pred_v_c),
+                                cv2.FONT_HERSHEY_SIMPLEX,
+                                0.4,
+                                (255, 0, 255),
+                                1,
+                            )
+
+    except Exception as e:
+        import traceback
+
+        traceback.print_exc()
+        cv2.putText(
+            debug_img_bgr,
+            f"Error: {str(e)[:50]}",
+            (10, 30),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.4,
+            (0, 0, 255),
+            1,
+        )
+
+    # Save and display
+    cv2.imwrite("/tmp/camera_debug.png", debug_img_bgr)
+    try:
+        cv2.imshow(f"Ego Camera (env {debug_env_idx})", debug_img_bgr)
+        cv2.waitKey(1)
+    except Exception:
+        pass
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/modular_tracking_env_cfg.py b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/modular_tracking_env_cfg.py
new file mode 100644
index 0000000000000000000000000000000000000000..9a6f399a041122edda6e820a230b594e9cd1a3eb
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/modular_tracking_env_cfg.py
@@ -0,0 +1,1041 @@
+from __future__ import annotations
+
+import dataclasses
+import os
+import re
+
+from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg
+from isaaclab.envs import ManagerBasedRLEnvCfg, ViewerCfg
+from isaaclab.scene import InteractiveSceneCfg
+from isaaclab.sensors import CameraCfg, ContactSensorCfg, FrameTransformerCfg, TiledCameraCfg
+import isaaclab.sim as sim_utils
+from isaaclab.sim.utils import clone
+from isaaclab.terrains import TerrainImporterCfg
+from isaaclab.utils import configclass
+import joblib
+import pxr
+
+from gear_sonic.envs.manager_env.mdp import terrain
+from gear_sonic.envs.manager_env.robots import g1, h2
+from gear_sonic.trl.utils import common
+
+
+def _load_opencv_params_from_usd(usd_path: str) -> dict:
+    """Parse a USD file to extract OpenCV lens distortion parameters.
+
+    Args:
+        usd_path: Path to the USD file containing camera with OpenCV distortion.
+
+    Returns:
+        Dictionary with keys: fx, fy, cx, cy, k1-k6, p1, p2, s1-s4
+
+    Raises:
+        FileNotFoundError: If USD file doesn't exist
+        ValueError: If required intrinsics (fx, fy, cx, cy) are missing
+    """
+    with open(usd_path) as f:
+        content = f.read()
+
+    params = {}
+
+    # Parse the USD text format for omni:lensdistortion:opencvPinhole attributes
+    patterns = {
+        "fx": r"omni:lensdistortion:opencvPinhole:fx\s*=\s*([-\d.]+)",
+        "fy": r"omni:lensdistortion:opencvPinhole:fy\s*=\s*([-\d.]+)",
+        "cx": r"omni:lensdistortion:opencvPinhole:cx\s*=\s*([-\d.]+)",
+        "cy": r"omni:lensdistortion:opencvPinhole:cy\s*=\s*([-\d.]+)",
+        "k1": r"omni:lensdistortion:opencvPinhole:k1\s*=\s*([-\d.]+)",
+        "k2": r"omni:lensdistortion:opencvPinhole:k2\s*=\s*([-\d.]+)",
+        "k3": r"omni:lensdistortion:opencvPinhole:k3\s*=\s*([-\d.]+)",
+        "k4": r"omni:lensdistortion:opencvPinhole:k4\s*=\s*([-\d.]+)",
+        "k5": r"omni:lensdistortion:opencvPinhole:k5\s*=\s*([-\d.]+)",
+        "k6": r"omni:lensdistortion:opencvPinhole:k6\s*=\s*([-\d.]+)",
+        "p1": r"omni:lensdistortion:opencvPinhole:p1\s*=\s*([-\d.]+)",
+        "p2": r"omni:lensdistortion:opencvPinhole:p2\s*=\s*([-\d.]+)",
+        "s1": r"omni:lensdistortion:opencvPinhole:s1\s*=\s*([-\d.]+)",
+        "s2": r"omni:lensdistortion:opencvPinhole:s2\s*=\s*([-\d.]+)",
+        "s3": r"omni:lensdistortion:opencvPinhole:s3\s*=\s*([-\d.]+)",
+        "s4": r"omni:lensdistortion:opencvPinhole:s4\s*=\s*([-\d.]+)",
+    }
+
+    for key, pattern in patterns.items():
+        match = re.search(pattern, content)
+        if match:
+            params[key] = float(match.group(1))
+
+    # Assert required intrinsics exist
+    required_intrinsics = ["fx", "fy", "cx", "cy"]
+    missing = [k for k in required_intrinsics if k not in params]
+    if missing:
+        raise ValueError(f"USD file {usd_path} missing required OpenCV intrinsics: {missing}")
+
+    # Assert at least k1 distortion exists (otherwise why use OpenCV mode?)
+    if "k1" not in params:
+        raise ValueError(
+            f"USD file {usd_path} missing distortion coefficient k1 - use pinhole camera instead"
+        )
+
+    print(f"Loaded OpenCV params from {usd_path}: {params}")  # noqa: T201
+    return params
+
+
+@clone
+def spawn_opencv_camera(
+    prim_path: str,
+    cfg,
+    translation: tuple[float, float, float] | None = None,
+    orientation: tuple[float, float, float, float] | None = None,
+    **kwargs,  # noqa: ARG001
+) -> pxr.Usd.Prim:
+    """Create a camera with OpenCV lens distortion applied at spawn time.
+
+    This spawner creates a standard camera and then applies the OmniLensDistortionOpenCvPinholeAPI
+    schema with the distortion parameters from the config. This happens BEFORE environment cloning,
+    so all cloned environments will have the correct distortion properties.
+    """
+    import isaacsim.core.utils.prims as prim_utils
+
+    # Create the camera prim
+    if not prim_utils.is_prim_path_valid(prim_path):
+        prim_utils.create_prim(
+            prim_path, "Camera", translation=translation, orientation=orientation
+        )
+    else:
+        raise ValueError(f"A prim already exists at path: '{prim_path}'.")
+
+    prim = prim_utils.get_prim_at_path(prim_path)
+
+    # Set clipping range (the only standard camera param that matters for OpenCV mode)
+    if hasattr(cfg, "clipping_range") and cfg.clipping_range is not None:
+        prim.GetAttribute("clippingRange").Set(cfg.clipping_range)
+
+    # Apply OmniLensDistortionOpenCvPinholeAPI schema via Sdf layer
+    # Note: This only works with headless=false (use xvfb-run for headless rendering)
+    stage = prim.GetStage()
+    layer = stage.GetRootLayer()
+    prim_spec = layer.GetPrimAtPath(prim_path)
+    if prim_spec is None:
+        prim_spec = pxr.Sdf.CreatePrimInLayer(layer, prim_path)  # noqa: F823
+
+    api_schemas = prim_spec.GetInfo("apiSchemas")
+    if api_schemas is None:
+        api_schemas = pxr.Sdf.TokenListOp()
+    prepend_items = list(api_schemas.prependedItems) if api_schemas.prependedItems else []
+    if "OmniLensDistortionOpenCvPinholeAPI" not in prepend_items:
+        prepend_items.append("OmniLensDistortionOpenCvPinholeAPI")
+        api_schemas.prependedItems = prepend_items
+        prim_spec.SetInfo("apiSchemas", api_schemas)
+
+    # Set the distortion model
+    prim.CreateAttribute("omni:lensdistortion:model", pxr.Sdf.ValueTypeNames.Token).Set(
+        "opencvPinhole"
+    )
+
+    # Set OpenCV distortion parameters from config
+    opencv_attrs = {
+        "omni:lensdistortion:opencvPinhole:fx": ("opencv_fx", pxr.Sdf.ValueTypeNames.Float),
+        "omni:lensdistortion:opencvPinhole:fy": ("opencv_fy", pxr.Sdf.ValueTypeNames.Float),
+        "omni:lensdistortion:opencvPinhole:cx": ("opencv_cx", pxr.Sdf.ValueTypeNames.Float),
+        "omni:lensdistortion:opencvPinhole:cy": ("opencv_cy", pxr.Sdf.ValueTypeNames.Float),
+        "omni:lensdistortion:opencvPinhole:k1": ("opencv_k1", pxr.Sdf.ValueTypeNames.Float),
+        "omni:lensdistortion:opencvPinhole:k2": ("opencv_k2", pxr.Sdf.ValueTypeNames.Float),
+        "omni:lensdistortion:opencvPinhole:k3": ("opencv_k3", pxr.Sdf.ValueTypeNames.Float),
+        "omni:lensdistortion:opencvPinhole:k4": ("opencv_k4", pxr.Sdf.ValueTypeNames.Float),
+        "omni:lensdistortion:opencvPinhole:k5": ("opencv_k5", pxr.Sdf.ValueTypeNames.Float),
+        "omni:lensdistortion:opencvPinhole:k6": ("opencv_k6", pxr.Sdf.ValueTypeNames.Float),
+        "omni:lensdistortion:opencvPinhole:p1": ("opencv_p1", pxr.Sdf.ValueTypeNames.Float),
+        "omni:lensdistortion:opencvPinhole:p2": ("opencv_p2", pxr.Sdf.ValueTypeNames.Float),
+        "omni:lensdistortion:opencvPinhole:s1": ("opencv_s1", pxr.Sdf.ValueTypeNames.Float),
+        "omni:lensdistortion:opencvPinhole:s2": ("opencv_s2", pxr.Sdf.ValueTypeNames.Float),
+        "omni:lensdistortion:opencvPinhole:s3": ("opencv_s3", pxr.Sdf.ValueTypeNames.Float),
+        "omni:lensdistortion:opencvPinhole:s4": ("opencv_s4", pxr.Sdf.ValueTypeNames.Float),
+    }
+
+    for attr_name, (cfg_name, attr_type) in opencv_attrs.items():
+        if hasattr(cfg, cfg_name):
+            value = getattr(cfg, cfg_name)
+            if value is not None:
+                prim.CreateAttribute(attr_name, attr_type).Set(value)
+
+    # Set image size
+    if hasattr(cfg, "opencv_image_size") and cfg.opencv_image_size is not None:
+        import pxr
+
+        prim.CreateAttribute(
+            "omni:lensdistortion:opencvPinhole:imageSize", pxr.Sdf.ValueTypeNames.Int2
+        ).Set(pxr.Gf.Vec2i(cfg.opencv_image_size[0], cfg.opencv_image_size[1]))
+        print(  # noqa: T201
+            f"[DEBUG] spawn_opencv_camera: Set imageSize to {cfg.opencv_image_size}"
+        )  # noqa: T201
+
+    # Debug: print all the values that were set
+    print(f"[DEBUG] spawn_opencv_camera at {prim_path}:")  # noqa: T201
+    print(  # noqa: T201
+        f"[DEBUG]   fx={getattr(cfg, 'opencv_fx', None)}, fy={getattr(cfg, 'opencv_fy', None)}"
+    )  # noqa: T201
+    print(  # noqa: T201
+        f"[DEBUG]   cx={getattr(cfg, 'opencv_cx', None)}, cy={getattr(cfg, 'opencv_cy', None)}"
+    )  # noqa: T201
+    print(f"[DEBUG]   image_size={getattr(cfg, 'opencv_image_size', None)}")  # noqa: T201
+
+    return prim
+
+
+@configclass
+class OpenCVCameraCfg:
+    """Camera config with OpenCV lens distortion parameters."""
+
+    func = spawn_opencv_camera
+    copy_from_source: bool = True  # Required by IsaacLab spawner
+
+    clipping_range: tuple[float, float] = (0.01, 20.0)
+
+    # These must exist or Camera.__init__ crashes, but values don't matter -
+    # OpenCV distortion API (fx/fy/cx/cy) overrides the projection at render time
+    horizontal_aperture: float = 1.0
+    vertical_aperture: float | None = None
+
+    # OpenCV intrinsics (required)
+    opencv_fx: float = None
+    opencv_fy: float = None
+    opencv_cx: float = None
+    opencv_cy: float = None
+    opencv_image_size: tuple[int, int] = None
+
+    # OpenCV distortion coefficients (required - at least k1 should be non-zero)
+    opencv_k1: float = None
+    opencv_k2: float = 0.0
+    opencv_k3: float = 0.0
+    opencv_k4: float = 0.0
+    opencv_k5: float = 0.0
+    opencv_k6: float = 0.0
+    opencv_p1: float = 0.0
+    opencv_p2: float = 0.0
+    opencv_s1: float = 0.0
+    opencv_s2: float = 0.0
+    opencv_s3: float = 0.0
+    opencv_s4: float = 0.0
+
+
+def _resolve_object_usd_paths(usd_path):
+    def expand_regex(path_pattern):
+        if os.path.exists(path_pattern):
+            return [os.path.abspath(path_pattern)]
+        directory = os.path.dirname(path_pattern) or "."
+        if not os.path.isdir(directory):
+            return []
+        base_pattern = os.path.basename(path_pattern)
+        try:
+            regex = re.compile(base_pattern)
+        except re.error as exc:
+            raise ValueError(f"Invalid object_usd_path regex: {base_pattern}") from exc
+        matched = [
+            os.path.abspath(os.path.join(directory, name))
+            for name in os.listdir(directory)
+            if regex.fullmatch(name)
+        ]
+        matched.sort(key=lambda p: os.path.basename(p))
+        return matched
+
+    if isinstance(usd_path, list | tuple):
+        raw_paths = list(usd_path)
+    else:
+        raw_paths = [usd_path]
+
+    expanded_paths = []
+    used_pattern = False
+    for path in raw_paths:
+        if isinstance(path, str):
+            pattern_matches = expand_regex(path)
+            if len(pattern_matches) == 1:
+                expanded_paths.append(os.path.abspath(pattern_matches[0]))
+            elif len(pattern_matches) > 1:
+                used_pattern = True
+                expanded_paths.extend(pattern_matches)
+            else:
+                raise ValueError(f"object_usd_path regex did not match any files: {path}")
+        else:
+            expanded_paths.append(path)
+
+    if not expanded_paths:
+        expanded_paths.append(os.path.abspath(raw_paths[0]))
+    if used_pattern:
+        expanded_paths.sort(key=lambda p: os.path.basename(p))
+    return expanded_paths
+
+
+@configclass
+class MySceneCfg(InteractiveSceneCfg):
+    """Configuration for the terrain scene with a legged robot."""
+
+    def __init__(self, config, **kwargs):  # noqa: ARG002
+        super().__init__()
+
+        self.num_envs = config.get("num_envs", 4096)
+
+        self.env_spacing = config.get("env_spacing", 2.5)
+
+        # Allow config to override replicate_physics (default True)
+        # Set to False for per-environment randomization (e.g., table size)
+        self.replicate_physics = config.get("replicate_physics", True)
+
+        self.eval_camera = None
+        if config.get("render_results", False):
+            self.eval_camera = TiledCameraCfg(
+                prim_path="/World/envs/env_.*/eval_camera",
+                offset=TiledCameraCfg.OffsetCfg(
+                    pos=(0, 0, 0), rot=(1, 0, 0, 0), convention="world"
+                ),
+                data_types=["rgb"],
+                spawn=sim_utils.PinholeCameraCfg(
+                    focal_length=5.0,
+                    focus_distance=50.0,
+                    horizontal_aperture=5,
+                    clipping_range=(0.1, 20.0),
+                ),
+                width=config.get("render_width", 1920),
+                height=config.get("render_height", 1080),
+            )
+
+        # Single global overview camera (not per-env) for replay videos
+        self.overview_camera = None
+        if config.get("overview_camera", False):
+            self.overview_camera = CameraCfg(
+                prim_path="/World/OverviewCamera",
+                offset=CameraCfg.OffsetCfg(pos=(0, 0, 50), rot=(1, 0, 0, 0), convention="world"),
+                data_types=["rgb"],
+                spawn=sim_utils.PinholeCameraCfg(
+                    focal_length=5.0,
+                    focus_distance=100.0,
+                    horizontal_aperture=10.0,  # Wider FOV for overview
+                    clipping_range=(0.1, 500.0),  # Extended for large grids
+                ),
+                width=config.get("render_width", 1920),
+                height=config.get("render_height", 1080),
+            )
+
+        # ground terrain
+        terrain_type = config.get("terrain_type", "plane")
+        if terrain_type == "plane":
+            self.terrain = TerrainImporterCfg(
+                prim_path="/World/ground",
+                terrain_type="plane",
+                collision_group=-1,
+                physics_material=sim_utils.RigidBodyMaterialCfg(
+                    friction_combine_mode="multiply",
+                    restitution_combine_mode="multiply",
+                    static_friction=1.0,
+                    dynamic_friction=1.0,
+                ),
+                visual_material=sim_utils.MdlFileCfg(
+                    mdl_path="{NVIDIA_NUCLEUS_DIR}/Materials/Base/Architecture/Shingles_01.mdl",
+                    project_uvw=True,
+                ),
+            )
+        elif terrain_type == "trimesh":
+            self.terrain = TerrainImporterCfg(
+                prim_path="/World/ground",
+                terrain_type="generator",
+                terrain_generator=terrain.ROUGH_TERRAINS_CFG,
+                max_init_terrain_level=10,
+                collision_group=1,
+                physics_material=sim_utils.RigidBodyMaterialCfg(
+                    friction_combine_mode="multiply",
+                    restitution_combine_mode="multiply",
+                    static_friction=1.0,
+                    dynamic_friction=1.0,
+                ),
+                visual_material=sim_utils.MdlFileCfg(
+                    mdl_path="{NVIDIA_NUCLEUS_DIR}/Materials/Base/Architecture/Shingles_01.mdl",
+                    project_uvw=True,
+                ),
+                debug_vis=False,
+            )
+        else:
+            raise ValueError(f"Unknown terrain type: {terrain_type}")
+
+        # robots
+        self.robot: ArticulationCfg = dataclasses.MISSING
+
+        # lights
+        if not config.get("render_ego_random", False):
+            self.light = AssetBaseCfg(
+                prim_path="/World/light",
+                spawn=sim_utils.DistantLightCfg(color=(0.75, 0.75, 0.75), intensity=3000.0),
+            )
+
+            self.sky_light = AssetBaseCfg(
+                prim_path="/World/skyLight",
+                spawn=sim_utils.DomeLightCfg(color=(0.13, 0.13, 0.13), intensity=1000.0),
+            )
+        else:
+            self.sky_light = AssetBaseCfg(
+                prim_path="/World/skyLight",
+                spawn=sim_utils.DomeLightCfg(color=(0.13, 0.13, 0.13), intensity=1000.0),
+            )
+            # from isaac_playground.env_rand.domelight import RandomDomeLightCfg
+            # self.sky_light = RandomDomeLightCfg(
+            #     prim_path="/World/skyLight",
+            #     texture_file_folder="../rl_data/HDRIs",
+            #     dynamic_randomize_texture=True,
+            #     dynamic_randomize_texture_interval=1.0
+            # )
+
+        self.contact_forces = ContactSensorCfg(
+            prim_path="{ENV_REGEX_NS}/Robot/.*",
+            history_length=3,
+            track_air_time=True,
+            force_threshold=10.0,
+            debug_vis=False,
+        )
+
+        # Check if robot has hands (43 DOF robots) - used for hand-related sensors
+        robot_type = config.get("robot", {}).get("type", "g1")
+        robot_has_hands = "43dof" in robot_type or "hand" in robot_type
+
+        motion_meta_info_path = config.get("motion_meta_info_path", None)
+        motion_meta_info = None
+
+        usd_path = config.get("object_usd_path", "")
+        if motion_meta_info_path is None and isinstance(usd_path, str) and os.path.isfile(usd_path):
+            motion_meta_info_path = usd_path.replace(".usd", ".pkl").replace("object_usd", "meta")
+        if motion_meta_info_path is not None and os.path.isfile(motion_meta_info_path):
+            motion_meta_info = joblib.load(motion_meta_info_path)
+
+        # Scene scale: config override > meta file > default 1.0
+        scene_scale = config.get("scene_scale", 1.0)
+        if motion_meta_info is not None:
+            scene_scale = motion_meta_info.get("scene_scale", 1.0)
+
+        # Ensure scene_scale is never None
+        if scene_scale is None:
+            scene_scale = 1.0
+
+        # Object with rigid body (optional, can be configured via config)
+        # Supports: single file, directory, list, or regex patterns for object_usd_path
+        #
+        # multi_object_per_env=True:  All objects spawned in every env (one active at a time)
+        # multi_object_per_env=False: One object per env (MultiUsdFileCfg cycles through)
+        if config.get("add_object", False):
+            usd_path = config.get("object_usd_path", f"{os.getcwd()}/data/wheelchair.usd")
+            object_is_dynamic = config.get("object_is_dynamic", False)
+            object_collision_enabled = config.get("object_collision_enabled", True)
+            multi_object_per_env = config.get("multi_object_per_env", False)
+            object_color = config.get("object_color", None)  # e.g. [0.6, 0.4, 0.2] for wood-brown
+
+            # --- Step 1: Resolve usd_path to a list of absolute paths ---
+            if isinstance(usd_path, list):
+                resolved_paths = _resolve_object_usd_paths(usd_path)
+            elif os.path.isdir(usd_path):
+                resolved_paths = sorted(
+                    os.path.abspath(os.path.join(usd_path, f))
+                    for f in os.listdir(usd_path)
+                    if f.endswith(".usd")
+                )
+            else:
+                resolved_paths = [os.path.abspath(usd_path)]
+
+            # --- Step 2: Spawn based on multi_object_per_env flag ---
+            if multi_object_per_env:
+                # MULTI-OBJECT MODE: One RigidObjectCfg per USD, all present in every env.
+                # commands.py detects this by checking for object_* entries in scene.rigid_objects.
+                #
+                # CRITICAL: Initial positions must be spread apart to avoid collision pairs!
+                # If all objects spawn at (0,0,0), PhysX creates O(N²) pairs at scene creation.
+                # Spread in Z (vertical) since envs only vary in X,Y.
+                z_spacing = 10.0
+                for idx, path in enumerate(resolved_paths):
+                    obj_name = os.path.splitext(os.path.basename(path))[0]
+                    obj_name_safe = obj_name.replace("-", "_")
+                    init_z = -100.0 - idx * z_spacing
+                    setattr(
+                        self,
+                        f"object_{obj_name_safe}",
+                        RigidObjectCfg(
+                            prim_path=f"{{ENV_REGEX_NS}}/Object_{obj_name_safe}",
+                            spawn=sim_utils.UsdFileCfg(
+                                usd_path=path,
+                                activate_contact_sensors=True,
+                                rigid_props=sim_utils.RigidBodyPropertiesCfg(
+                                    kinematic_enabled=not object_is_dynamic,
+                                    max_depenetration_velocity=1.0,
+                                ),
+                                collision_props=sim_utils.CollisionPropertiesCfg(
+                                    collision_enabled=object_collision_enabled,
+                                ),
+                            ),
+                            init_state=RigidObjectCfg.InitialStateCfg(pos=(1000.0, 0.0, init_z)),
+                        ),
+                    )
+                print(  # noqa: T201
+                    f"[Multi-Object Mode] Spawned {len(resolved_paths)} objects at spread Z positions"
+                )
+            elif len(resolved_paths) == 1:
+                # SINGLE OBJECT MODE
+                object_mass = config.get("object_mass", None)
+                mass_props = (
+                    sim_utils.MassPropertiesCfg(mass=object_mass)
+                    if object_mass is not None
+                    else None
+                )
+                object_opacity = config.get("object_opacity", 1.0)
+                if object_color is not None:
+                    visual_material = sim_utils.PreviewSurfaceCfg(
+                        diffuse_color=tuple(object_color[:3]),
+                        opacity=object_opacity,
+                        metallic=object_color[3] if len(object_color) > 3 else 0.0,
+                        roughness=object_color[4] if len(object_color) > 4 else 0.5,
+                    )
+                elif object_opacity < 1.0:
+                    visual_material = sim_utils.PreviewSurfaceCfg(opacity=object_opacity)
+                else:
+                    visual_material = None
+                self.object = RigidObjectCfg(
+                    prim_path="{ENV_REGEX_NS}/Object",
+                    spawn=sim_utils.UsdFileCfg(
+                        usd_path=resolved_paths[0],
+                        activate_contact_sensors=True,
+                        rigid_props=sim_utils.RigidBodyPropertiesCfg(
+                            kinematic_enabled=not object_is_dynamic,
+                            max_depenetration_velocity=1.0,
+                        ),
+                        mass_props=mass_props,
+                        collision_props=sim_utils.CollisionPropertiesCfg(
+                            collision_enabled=object_collision_enabled,
+                        ),
+                        visual_material=visual_material,
+                        scale=(scene_scale, scene_scale, scene_scale),
+                    ),
+                    init_state=RigidObjectCfg.InitialStateCfg(
+                        pos=tuple(config.get("object_position", [2.0, 0.0, 0.0]))
+                    ),
+                )
+            else:
+                # ONE-PER-ENV MODE: Different object per env via MultiUsdFileCfg
+                self.replicate_physics = False
+                self.object = RigidObjectCfg(
+                    prim_path="{ENV_REGEX_NS}/Object",
+                    spawn=sim_utils.MultiUsdFileCfg(
+                        usd_path=resolved_paths,
+                        random_choice=False,
+                        activate_contact_sensors=True,
+                        rigid_props=sim_utils.RigidBodyPropertiesCfg(
+                            kinematic_enabled=not object_is_dynamic,
+                            max_depenetration_velocity=1.0,
+                        ),
+                        collision_props=sim_utils.CollisionPropertiesCfg(
+                            collision_enabled=object_collision_enabled,
+                        ),
+                    ),
+                    init_state=RigidObjectCfg.InitialStateCfg(
+                        pos=tuple(config.get("object_position", [2.0, 0.0, 0.0]))
+                    ),
+                )
+
+            if robot_has_hands:
+                # Frame transformer for hand-object tracking
+                self.object_to_hand_frame_transformer = FrameTransformerCfg(
+                    prim_path="{ENV_REGEX_NS}/Object",  # Source: Object
+                    target_frames=[
+                        FrameTransformerCfg.FrameCfg(
+                            name="right_hand_thumb",
+                            prim_path="{ENV_REGEX_NS}/Robot/right_hand_thumb_2_link",
+                        ),
+                        FrameTransformerCfg.FrameCfg(
+                            name="right_hand_index",
+                            prim_path="{ENV_REGEX_NS}/Robot/right_hand_index_1_link",
+                        ),
+                        FrameTransformerCfg.FrameCfg(
+                            name="right_hand_middle",
+                            prim_path="{ENV_REGEX_NS}/Robot/right_hand_middle_1_link",
+                        ),
+                    ],
+                )
+
+                # Object-to-hand contact sensor for grasp detection (only for robots with hands)
+                # force_matrix_w gives shape [num_envs, 1, N, 3] - force on Object from each finger link
+                # Configurable via contact_sensor_finger_links (list of link names without Robot/ prefix)
+                # Default: tips + palm (4 links). For fuller grasp detection, use all finger segments:
+                #   [right_hand_palm_link, right_hand_thumb_0_link, right_hand_thumb_1_link,
+                #    right_hand_thumb_2_link, right_hand_index_0_link, right_hand_index_1_link,
+                #    right_hand_middle_0_link, right_hand_middle_1_link]
+                custom_finger_links = config.get("contact_sensor_finger_links", None)
+                if custom_finger_links is not None:
+                    right_finger_tip_bodies = [
+                        f"{{ENV_REGEX_NS}}/Robot/{link}" for link in custom_finger_links
+                    ]
+                else:
+                    # Default: fingertips + palm (4 links)
+                    right_finger_tip_bodies = [
+                        "{ENV_REGEX_NS}/Robot/right_hand_thumb_2_link",
+                        "{ENV_REGEX_NS}/Robot/right_hand_index_1_link",
+                        "{ENV_REGEX_NS}/Robot/right_hand_middle_1_link",
+                        "{ENV_REGEX_NS}/Robot/right_hand_palm_link",
+                    ]
+                self.object_to_hand_contact_sensor = ContactSensorCfg(
+                    prim_path="{ENV_REGEX_NS}/Object",
+                    filter_prim_paths_expr=right_finger_tip_bodies,
+                    history_length=2,
+                    track_air_time=False,
+                )
+
+                left_finger_tip_bodies = [
+                    "{ENV_REGEX_NS}/Robot/left_hand_thumb_2_link",
+                    "{ENV_REGEX_NS}/Robot/left_hand_index_1_link",
+                    "{ENV_REGEX_NS}/Robot/left_hand_middle_1_link",
+                    "{ENV_REGEX_NS}/Robot/left_hand_palm_link",
+                ]
+                self.object_to_left_hand_contact_sensor = ContactSensorCfg(
+                    prim_path="{ENV_REGEX_NS}/Object",
+                    filter_prim_paths_expr=left_finger_tip_bodies,
+                    history_length=2,
+                    track_air_time=False,
+                )
+
+        if config.get("add_table", False):
+            # Table initial position and orientation from config or meta
+            # (per-env update in commands.py for multi-motion)
+            if "table_position" in config:
+                table_pos = config["table_position"]
+            elif motion_meta_info is not None and "table_pos" in motion_meta_info:
+                table_pos = motion_meta_info["table_pos"]
+            else:
+                table_pos = [0.0, 0.0, 0.8]
+
+            # Apply table_offset if configured (additive offset to table position)
+            table_offset = config.get("table_offset", None)
+            if table_offset is not None:
+                table_pos[0] += table_offset[0]
+                table_pos[1] += table_offset[1]
+                table_pos[2] += table_offset[2]
+
+            # Table quaternion (w, x, y, z) - GRAB data uses 90° X-rotation
+            if "table_quat" in config:
+                table_quat = tuple(config["table_quat"])
+            elif motion_meta_info is not None and "table_quat" in motion_meta_info:
+                table_quat = tuple(motion_meta_info["table_quat"])
+            else:
+                # Default for GRAB table USD: 90° rotation around X-axis
+                table_quat = (0.7071068, -0.7071068, 0.0, 0.0)
+
+            # Use table USD when: scene_scale is explicitly set OR table_usd_path is provided OR meta info exists
+            use_table_usd = config.get("table_usd_path") is not None
+
+            if not use_table_usd:
+                # Use simple cuboid table for GeniHOI or default scenarios
+                # Table size can be configured via config or meta file
+                # For GeniHOI: table_size = [width, depth, thickness] from meta
+                # Default: 1.0 x 0.6 x 0.04 (reasonable for kitchen/desk scenarios)
+
+                if "table_size" in config:
+                    table_size = config["table_size"]
+                    table_width = table_size[0]
+                    table_depth = table_size[1]
+                    table_thickness = table_size[2]
+                elif motion_meta_info is not None and "table_size" in motion_meta_info:
+                    table_size = motion_meta_info["table_size"]
+                    table_width = table_size[0]
+                    table_depth = table_size[1]
+                    table_thickness = table_size[2]
+                else:
+                    # Default cuboid size - large enough for most GeniHOI scenarios
+                    table_width = 1.0
+                    table_depth = 0.6
+                    table_thickness = 0.04
+
+                # For cuboid tables, use identity quaternion (no rotation needed)
+                table_quat_cuboid = (1.0, 0.0, 0.0, 0.0)
+
+                self.table = RigidObjectCfg(
+                    prim_path="{ENV_REGEX_NS}/Table",
+                    spawn=sim_utils.CuboidCfg(
+                        activate_contact_sensors=True,
+                        size=(table_width, table_depth, table_thickness),
+                        rigid_props=sim_utils.RigidBodyPropertiesCfg(
+                            kinematic_enabled=True,
+                            max_depenetration_velocity=1.0,
+                        ),
+                        mass_props=sim_utils.MassPropertiesCfg(density=500.0),
+                        collision_props=sim_utils.CollisionPropertiesCfg(
+                            collision_enabled=True,
+                        ),
+                        visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.5, 0.35, 0.2)),
+                    ),
+                    init_state=RigidObjectCfg.InitialStateCfg(
+                        pos=tuple(table_pos), rot=table_quat_cuboid
+                    ),
+                )
+            else:
+                # Use table USD with scene_scale and proper rotation
+                self.table = RigidObjectCfg(
+                    prim_path="{ENV_REGEX_NS}/Table",
+                    spawn=sim_utils.UsdFileCfg(
+                        activate_contact_sensors=True,
+                        usd_path=config.get(
+                            "table_usd_path", "data/motion_lib_grab/common/table.usda"
+                        ),
+                        rigid_props=sim_utils.RigidBodyPropertiesCfg(
+                            kinematic_enabled=True,
+                            max_depenetration_velocity=1.0,
+                        ),
+                        collision_props=sim_utils.CollisionPropertiesCfg(
+                            collision_enabled=True,
+                        ),
+                        scale=(scene_scale, scene_scale, scene_scale),
+                    ),
+                    init_state=RigidObjectCfg.InitialStateCfg(pos=tuple(table_pos), rot=table_quat),
+                )
+
+            # Table-to-hand contact sensor for walk_stand_grasp rewards (only for robots with hands)
+            if robot_has_hands:
+                # Use Table as prim_path and explicitly list hand bodies in filter_prim_paths_expr
+                table_contact_bodies = [
+                    # Right wrist
+                    "{ENV_REGEX_NS}/Robot/right_wrist_roll_link",
+                    "{ENV_REGEX_NS}/Robot/right_wrist_pitch_link",
+                    "{ENV_REGEX_NS}/Robot/right_wrist_yaw_link",
+                    # Right hand
+                    "{ENV_REGEX_NS}/Robot/right_hand_palm_link",
+                    "{ENV_REGEX_NS}/Robot/right_hand_index_0_link",
+                    "{ENV_REGEX_NS}/Robot/right_hand_index_1_link",
+                    "{ENV_REGEX_NS}/Robot/right_hand_middle_0_link",
+                    "{ENV_REGEX_NS}/Robot/right_hand_middle_1_link",
+                    "{ENV_REGEX_NS}/Robot/right_hand_thumb_0_link",
+                    "{ENV_REGEX_NS}/Robot/right_hand_thumb_1_link",
+                    "{ENV_REGEX_NS}/Robot/right_hand_thumb_2_link",
+                ]
+                self.table_to_hand_contact_sensor = ContactSensorCfg(
+                    prim_path="{ENV_REGEX_NS}/Table",
+                    filter_prim_paths_expr=table_contact_bodies,
+                    history_length=2,
+                    track_air_time=False,
+                )
+
+            # Object-to-table contact sensor for detecting object-table contact forces
+            # Sensor is attached to Object, filters for contact with Table
+            # force_matrix_w gives shape [num_envs, 1, 1, 3] - force on Object from Table
+            if config.get("add_object", False):
+                self.object_to_table_contact_sensor = ContactSensorCfg(
+                    prim_path="{ENV_REGEX_NS}/Object",
+                    filter_prim_paths_expr=["{ENV_REGEX_NS}/Table"],
+                    history_length=2,
+                    track_air_time=False,
+                )
+
+                # Object-to-robot contact sensor: detect when right hand/wrist touches object
+                # Used for termination logic (robot must touch object before table)
+                # NOTE: IsaacLab ContactSensor only supports one-to-many filtering,
+                # so prim_path must be single body (Object), filter can be multiple (robot links)
+                right_hand_wrist_links = [
+                    # Right wrist
+                    "{ENV_REGEX_NS}/Robot/right_wrist_roll_link",
+                    "{ENV_REGEX_NS}/Robot/right_wrist_pitch_link",
+                    "{ENV_REGEX_NS}/Robot/right_wrist_yaw_link",
+                    # Right hand
+                    "{ENV_REGEX_NS}/Robot/right_hand_palm_link",
+                    "{ENV_REGEX_NS}/Robot/right_hand_index_0_link",
+                    "{ENV_REGEX_NS}/Robot/right_hand_index_1_link",
+                    "{ENV_REGEX_NS}/Robot/right_hand_middle_0_link",
+                    "{ENV_REGEX_NS}/Robot/right_hand_middle_1_link",
+                    "{ENV_REGEX_NS}/Robot/right_hand_thumb_0_link",
+                    "{ENV_REGEX_NS}/Robot/right_hand_thumb_1_link",
+                    "{ENV_REGEX_NS}/Robot/right_hand_thumb_2_link",
+                ]
+                self.object_to_robot_contact_sensor = ContactSensorCfg(
+                    prim_path="{ENV_REGEX_NS}/Object",
+                    filter_prim_paths_expr=right_hand_wrist_links,
+                    history_length=2,
+                    track_air_time=False,
+                )
+
+            # Table-to-robot contact sensor: detect when right hand/wrist touches table
+            # Used for termination logic (terminate if table contact before object contact)
+            self.table_to_robot_contact_sensor = ContactSensorCfg(
+                prim_path="{ENV_REGEX_NS}/Table",
+                filter_prim_paths_expr=right_hand_wrist_links,
+                history_length=2,
+                track_air_time=False,
+            )
+
+        # # TODO: Do this better.  # noqa: TD002, TD003
+        # self.ego_camera = TiledCameraCfg(
+        #     prim_path="{ENV_REGEX_NS}/Robot/head_link/HeadCamera",
+        #     # parent=SceneEntityCfg(
+        #     #     name="robot",
+        #     #     body_names=["head_link"],
+        #     # ),
+        #     # offset=TiledCameraCfg.OffsetCfg(pos=camera_pos_offset, rot=camera_rot_offset, convention="world"),
+        #     data_types=['rgb'],
+        #     spawn=sim_utils.PinholeCameraCfg(focal_length=5.0, focus_distance=50.0, horizontal_aperture=5, clipping_range=(0.1, 20.0)),  # noqa: E501
+        #     width=384,
+        #     height=384,
+        #     debug_vis=True,
+        # )
+
+        # TODO: Do this better.  # noqa: TD002, TD003
+        # Copied from gear_sonic/config/simulator/isaacsim.yaml
+        # enable_cameras flag creates the ego camera for vision-based policies
+        if config.get("enable_cameras", False) or config.get("render_ego", False):
+            # Get camera config from nested cameras dict
+            cameras_cfg = config.get("cameras", {})
+
+            # Choose camera type: fisheye, pinhole, or opencv (with lens distortion)
+            if config.get("render_ego_opencv_usd", False):
+                # Load OpenCV distortion parameters from USD file
+                opencv_usd_path = cameras_cfg.get(
+                    "opencv_usd_path", "runs/oak_camera.usda"  # Default path
+                )
+                print(  # noqa: T201
+                    f"USING OPENCV CAMERA - loading distortion from: {opencv_usd_path}"
+                )  # noqa: T201
+
+                # Parse the USD file to extract distortion parameters
+                opencv_params = _load_opencv_params_from_usd(opencv_usd_path)
+
+                # Camera resolution [H, W] - required, no fallback
+                camera_resolution = cameras_cfg["camera_resolution"]
+
+                # Use custom OpenCV spawner that applies distortion at spawn time (before cloning)
+                print(  # noqa: T201
+                    f"[DEBUG] Creating OpenCVCameraCfg with resolution={camera_resolution} -> image_size={(camera_resolution[1], camera_resolution[0])}"  # noqa: E501
+                )
+                print(  # noqa: T201
+                    f"[DEBUG] fx={opencv_params['fx']}, fy={opencv_params['fy']}, cx={opencv_params['cx']}, cy={opencv_params['cy']}"  # noqa: E501
+                )
+                print(f"[DEBUG] k1={opencv_params['k1']}")  # noqa: T201
+                camera_spawn_cfg = OpenCVCameraCfg(
+                    clipping_range=(0.01, 20.0),
+                    opencv_fx=opencv_params["fx"],
+                    opencv_fy=opencv_params["fy"],
+                    opencv_cx=opencv_params["cx"],
+                    opencv_cy=opencv_params["cy"],
+                    opencv_k1=opencv_params["k1"],  # Required - validated by parser
+                    opencv_k2=opencv_params.get("k2", 0.0),
+                    opencv_k3=opencv_params.get("k3", 0.0),
+                    opencv_k4=opencv_params.get("k4", 0.0),
+                    opencv_k5=opencv_params.get("k5", 0.0),
+                    opencv_k6=opencv_params.get("k6", 0.0),
+                    opencv_p1=opencv_params.get("p1", 0.0),
+                    opencv_p2=opencv_params.get("p2", 0.0),
+                    opencv_s1=opencv_params.get("s1", 0.0),
+                    opencv_s2=opencv_params.get("s2", 0.0),
+                    opencv_s3=opencv_params.get("s3", 0.0),
+                    opencv_s4=opencv_params.get("s4", 0.0),
+                    opencv_image_size=(camera_resolution[1], camera_resolution[0]),  # (W, H)
+                )
+            elif config.get("render_ego_fisheye", False):
+                print("USING FISHEYE CAMERA" * 1000)  # noqa: T201
+                camera_spawn_cfg = sim_utils.FisheyeCameraCfg(
+                    projection_type=config.get("render_ego_fisheye_projection", "fisheyeSpherical"),
+                    fisheye_max_fov=config.get("render_ego_fisheye_fov", 180.0),
+                    focus_distance=0.5,
+                    clipping_range=(0.1, 20.0),
+                )
+            elif config.get("render_ego_fisheye_polynomial", False):
+                camera_spawn_cfg = sim_utils.FisheyeCameraCfg(
+                    projection_type="fisheyePolynomial",
+                    fisheye_max_fov=360.0,
+                    focus_distance=0.5,
+                    clipping_range=(0.1, 20.0),
+                )
+            else:
+                # Convert clipping_range to tuple if it's a list or string (from YAML config)
+                clipping_range = cameras_cfg.get("camera_clipping_range", (0.1, 20.0))
+                if isinstance(clipping_range, str):
+                    # Parse string like "(0.01,20.0)" to tuple
+                    clipping_range = tuple(map(float, clipping_range.strip("()").split(",")))
+                elif isinstance(clipping_range, list):
+                    clipping_range = tuple(clipping_range)
+                camera_spawn_cfg = sim_utils.PinholeCameraCfg(
+                    focal_length=cameras_cfg.get("camera_focal_length", 1.88),
+                    focus_distance=cameras_cfg.get("camera_focus_distance", 0.5),
+                    horizontal_aperture=cameras_cfg.get("camera_horizontal_aperture", 2.6035),
+                    vertical_aperture=cameras_cfg.get("camera_vertical_aperture", 1.4621),
+                    clipping_range=clipping_range,
+                )
+
+            # Camera attachment link (e.g., "d435_link" for RealSense D435)
+            # If not specified, camera is attached to robot root
+            camera_attached_link = cameras_cfg.get("camera_attached_link", "")
+            if camera_attached_link:
+                camera_prim_path = f"{{ENV_REGEX_NS}}/Robot/{camera_attached_link}/ego_camera"
+            else:
+                camera_prim_path = "{ENV_REGEX_NS}/Robot/ego_camera"
+
+            # Camera position and rotation offsets (relative to attached link)
+            camera_pos_offset = tuple(cameras_cfg.get("camera_pos_offset", [0.0, 0.0, 0.0]))
+            camera_rot_offset = tuple(
+                cameras_cfg.get("camera_rot_offset", [1.0, 0.0, 0.0, 0.0])
+            )  # wxyz quaternion
+
+            # Camera data types (e.g., ["rgb"], ["rgb", "depth"])
+            camera_data_types = cameras_cfg.get("camera_data_types", ["rgb"])
+
+            # Camera resolution [H, W]
+            camera_resolution = cameras_cfg.get("camera_resolution", [108, 192])
+
+            self.ego_camera = TiledCameraCfg(
+                prim_path=camera_prim_path,
+                offset=TiledCameraCfg.OffsetCfg(
+                    pos=camera_pos_offset, rot=camera_rot_offset, convention="world"
+                ),
+                data_types=camera_data_types,
+                spawn=camera_spawn_cfg,
+                height=camera_resolution[0],
+                width=camera_resolution[1],
+                debug_vis=True,
+                update_period=0.0,
+                update_latest_camera_pose=True,
+            )
+
+
+@configclass
+class ModularTrackingEnvCfg(ManagerBasedRLEnvCfg):
+    """Modular configuration for the tracking environment that uses Hydra composition."""
+
+    def __init__(  # noqa: D417
+        self,
+        config,
+        actions,
+        observations,
+        rewards,
+        terminations,
+        commands,
+        events,
+        curriculum,
+        recorders,
+        **kwargs,  # noqa: ARG002
+    ):
+        """Initialize the configuration with Hydra config.
+
+        Args:
+            hydra_cfg: Hydra configuration containing component specifications
+        """
+        super().__init__()
+        self._setup_from_hydra(
+            config,
+            actions,
+            observations,
+            rewards,
+            terminations,
+            commands,
+            events,
+            curriculum,
+            recorders,
+        )
+        self.override_settings()
+
+    def _setup_from_hydra(
+        self,
+        config,
+        actions,
+        observations,
+        rewards,
+        terminations,
+        commands,
+        events,
+        curriculum,
+        recorders,
+    ):
+        """Setup configuration from Hydra config."""
+        self.config = config
+        # Scene settings
+        self.scene = MySceneCfg(config=config)
+        # Instantiate components using Hydra
+        self.actions = common.custom_instantiate(actions, _recursive=True)
+        self.observations = common.custom_instantiate(observations, _recursive=True)
+        self.rewards = common.custom_instantiate(rewards, _recursive=True)
+        self.terminations = common.custom_instantiate(terminations, _recursive=True)
+        self.commands = common.custom_instantiate(commands, _recursive=True)
+        self.events = common.custom_instantiate(events, _recursive=True)
+        # Debug: Print instantiated events
+        print(f"[DEBUG] Instantiated events: {self.events}")  # noqa: T201
+        if hasattr(self.events, "__dict__"):
+            for event_name, event_cfg in self.events.__dict__.items():
+                if event_cfg is not None and not event_name.startswith("_"):
+                    print(  # noqa: T201
+                        f"[DEBUG]   Event '{event_name}': {type(event_cfg).__name__}"
+                    )  # noqa: T201
+        self.curriculum = common.custom_instantiate(curriculum, _recursive=True)
+        self.recorders = common.custom_instantiate(recorders, _recursive=True)
+
+    def override_settings(self):
+        config = self.config
+        # General settings
+        self.decimation = config.get("decimation", 4)
+        self.episode_length_s = config.get("episode_length_s", 10.0)
+
+        # Simulation settings
+        self.sim.dt = config.get("sim_dt", 0.005)
+        self.sim.render_interval = self.decimation
+        self.sim.physics_material = self.scene.terrain.physics_material
+        self.sim.physx.gpu_max_rigid_patch_count = 10 * 2**15
+
+        # Increase collision stack size for scenes with complex collision meshes (e.g. staircases)
+        gpu_collision_stack_size_exp = config.get("gpu_collision_stack_size_exp", 26)
+        self.sim.physx.gpu_collision_stack_size = 2**gpu_collision_stack_size_exp
+
+        # Increase PhysX GPU memory only for multi-object scenes
+        # These prevent "totalAggregatePairsCapacity" errors when many objects are spawned
+        # Check if object_usd_path is a directory (multi-object mode)
+        object_usd_path = config.get("object_usd_path", "")
+        if config.get("add_object", False) and (
+            isinstance(object_usd_path, list) or os.path.isdir(object_usd_path)
+        ):
+            # With proper Z-spacing of initial positions, collision pairs should be minimal
+            # These are moderate values that should work for 1000+ envs
+            self.sim.physx.gpu_found_lost_pairs_capacity = 2**24  # ~16M
+            self.sim.physx.gpu_found_lost_aggregate_pairs_capacity = 2**24
+            self.sim.physx.gpu_total_aggregate_pairs_capacity = 2**21  # ~2M
+
+        # Viewer settings
+        viewer_config = config.get("viewer", {})
+        self.viewer = ViewerCfg(
+            eye=viewer_config.get("eye", [4.5, 0.0, 4.0]),
+            lookat=viewer_config.get("lookat", [0.0, 0.0, 0.0]),
+        )
+
+        robot_mapping = {
+            "g1_model_12_dex": {
+                "robot_cfg": g1.G1_CYLINDER_MODEL_12_DEX_CFG,
+                "action_scale": g1.G1_MODEL_12_ACTION_SCALE,
+                "isaaclab_to_mujoco_mapping": g1.G1_ISAACLAB_TO_MUJOCO_MAPPING,
+            },
+            "h2": {
+                "robot_cfg": h2.H2_CFG,
+                "action_scale": h2.H2_ACTION_SCALE,
+                "isaaclab_to_mujoco_mapping": h2.H2_ISAACLAB_TO_MUJOCO_MAPPING,
+            },
+        }
+
+        robot_type = config["robot"].get("type", "g1")
+        self.scene.robot = robot_mapping[robot_type]["robot_cfg"].replace(
+            prim_path="{ENV_REGEX_NS}/Robot"
+        )
+        self.actions.joint_pos.scale = robot_mapping[config["robot"].get("type", "g1")][
+            "action_scale"
+        ]
+        self.isaaclab_to_mujoco_mapping = robot_mapping[config["robot"].get("type", "g1")][
+            "isaaclab_to_mujoco_mapping"
+        ]
+
+        # curriculum? WARNING HARDCODED
+        import importlib
+
+        if (
+            hasattr(self.curriculum, "force_push_curriculum")
+            and self.curriculum.force_push_curriculum
+        ):
+            module = importlib.import_module("gear_sonic.envs.manager_env.mdp")
+            self.curriculum.force_push_curriculum.params["modify_fn"] = getattr(
+                module, "step_curriculum"
+            )
+
+        if (
+            hasattr(self.curriculum, "force_push_linear_curriculum")
+            and self.curriculum.force_push_linear_curriculum
+        ):
+            module = importlib.import_module("gear_sonic.envs.manager_env.mdp")
+            self.curriculum.force_push_linear_curriculum.params["modify_fn"] = getattr(
+                module, "linear_curriculum"
+            )
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__init__.py b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..dc1a8efbc5a3bb28e5305e13de3a3ef9f166bc2a
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__init__.py
@@ -0,0 +1,4 @@
+# Re-export robot configs for backward compatibility.
+# Import from gear_sonic.envs.manager_env.robots.g1 or .h2 directly for new code.
+from gear_sonic.envs.manager_env.robots.g1 import *  # noqa: F401,F403
+from gear_sonic.envs.manager_env.robots.h2 import *  # noqa: F401,F403
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__pycache__/__init__.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e9dcea0aaf00d009cb6cc761e67ac908a595ce1d
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__pycache__/__init__.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__pycache__/__init__.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f7882307ac7b7dafc22fa76d9529cecb8089ba81
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__pycache__/__init__.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__pycache__/g1.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__pycache__/g1.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4e71e3a06e9388f0b637468ad134bc840d1a1625
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__pycache__/g1.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__pycache__/g1.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__pycache__/g1.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0a122d0f421c35e590e3e967b7296ede539493a5
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__pycache__/g1.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__pycache__/h2.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__pycache__/h2.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d700ce787f4ba6ed063bed758ef57d437e5ddf06
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__pycache__/h2.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__pycache__/h2.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__pycache__/h2.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c1a09efa03925dc3fb4defd1ecf1cb281cdd4dde
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/__pycache__/h2.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/g1.py b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/g1.py
new file mode 100644
index 0000000000000000000000000000000000000000..5066546859b11f49faeae77efa6a6e617641ef03
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/g1.py
@@ -0,0 +1,371 @@
+# Robot configuration adapted from the BeyondMimic project.
+# See: https://github.com/HybridRobotics/whole_body_tracking
+
+from isaaclab.actuators import ImplicitActuatorCfg
+from isaaclab.assets.articulation import ArticulationCfg
+import isaaclab.sim as sim_utils
+
+ASSET_DIR = "gear_sonic/data/assets"
+
+ARMATURE_5020 = 0.003609725
+ARMATURE_7520_14 = 0.010177520
+ARMATURE_7520_22 = 0.025101925
+ARMATURE_4010 = 0.00425
+
+NATURAL_FREQ = 10 * 2.0 * 3.1415926535  # 10Hz
+DAMPING_RATIO = 2.0
+
+STIFFNESS_5020 = ARMATURE_5020 * NATURAL_FREQ**2
+STIFFNESS_7520_14 = ARMATURE_7520_14 * NATURAL_FREQ**2
+STIFFNESS_7520_22 = ARMATURE_7520_22 * NATURAL_FREQ**2
+STIFFNESS_4010 = ARMATURE_4010 * NATURAL_FREQ**2
+
+DAMPING_5020 = 2.0 * DAMPING_RATIO * ARMATURE_5020 * NATURAL_FREQ
+DAMPING_7520_14 = 2.0 * DAMPING_RATIO * ARMATURE_7520_14 * NATURAL_FREQ
+DAMPING_7520_22 = 2.0 * DAMPING_RATIO * ARMATURE_7520_22 * NATURAL_FREQ
+DAMPING_4010 = 2.0 * DAMPING_RATIO * ARMATURE_4010 * NATURAL_FREQ
+
+G1_ISAACLAB_JOINTS = [
+    "pelvis",
+    "left_hip_pitch_link",
+    "right_hip_pitch_link",
+    "waist_yaw_link",
+    "left_hip_roll_link",
+    "right_hip_roll_link",
+    "waist_roll_link",
+    "left_hip_yaw_link",
+    "right_hip_yaw_link",
+    "torso_link",
+    "left_knee_link",
+    "right_knee_link",
+    "left_shoulder_pitch_link",
+    "right_shoulder_pitch_link",
+    "left_ankle_pitch_link",
+    "right_ankle_pitch_link",
+    "left_shoulder_roll_link",
+    "right_shoulder_roll_link",
+    "left_ankle_roll_link",
+    "right_ankle_roll_link",
+    "left_shoulder_yaw_link",
+    "right_shoulder_yaw_link",
+    "left_elbow_link",
+    "right_elbow_link",
+    "left_wrist_roll_link",
+    "right_wrist_roll_link",
+    "left_wrist_pitch_link",
+    "right_wrist_pitch_link",
+    "left_wrist_yaw_link",
+    "right_wrist_yaw_link",
+]
+
+G1_ISAACLAB_TO_MUJOCO_DOF = [
+    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,
+]
+
+G1_MUJOCO_TO_ISAACLAB_DOF = [
+    0,
+    6,
+    12,
+    1,
+    7,
+    13,
+    2,
+    8,
+    14,
+    3,
+    9,
+    15,
+    22,
+    4,
+    10,
+    16,
+    23,
+    5,
+    11,
+    17,
+    24,
+    18,
+    25,
+    19,
+    26,
+    20,
+    27,
+    21,
+    28,
+]
+
+G1_ISAACLAB_TO_MUJOCO_BODY = [
+    0,
+    1,
+    4,
+    7,
+    10,
+    14,
+    18,
+    2,
+    5,
+    8,
+    11,
+    15,
+    19,
+    3,
+    6,
+    9,
+    12,
+    16,
+    20,
+    22,
+    24,
+    26,
+    28,
+    13,
+    17,
+    21,
+    23,
+    25,
+    27,
+    29,
+]
+
+G1_MUJOCO_TO_ISAACLAB_BODY = [
+    0,
+    1,
+    7,
+    13,
+    2,
+    8,
+    14,
+    3,
+    9,
+    15,
+    4,
+    10,
+    16,
+    23,
+    5,
+    11,
+    17,
+    24,
+    6,
+    12,
+    18,
+    25,
+    19,
+    26,
+    20,
+    27,
+    21,
+    28,
+    22,
+    29,
+]
+
+G1_ISAACLAB_TO_MUJOCO_MAPPING = {
+    "isaaclab_joints": G1_ISAACLAB_JOINTS,
+    "isaaclab_to_mujoco_dof": G1_ISAACLAB_TO_MUJOCO_DOF,
+    "mujoco_to_isaaclab_dof": G1_MUJOCO_TO_ISAACLAB_DOF,
+    "isaaclab_to_mujoco_body": G1_ISAACLAB_TO_MUJOCO_BODY,
+    "mujoco_to_isaaclab_body": G1_MUJOCO_TO_ISAACLAB_BODY,
+}
+
+G1_CYLINDER_MODEL_12_DEX_CFG = ArticulationCfg(
+    spawn=sim_utils.UrdfFileCfg(
+        fix_base=False,
+        replace_cylinders_with_capsules=True,
+        asset_path=f"{ASSET_DIR}/robot_description/urdf/g1/main.urdf",
+        activate_contact_sensors=True,
+        rigid_props=sim_utils.RigidBodyPropertiesCfg(
+            disable_gravity=False,
+            retain_accelerations=False,
+            linear_damping=0.0,
+            angular_damping=0.0,
+            max_linear_velocity=1000.0,
+            max_angular_velocity=1000.0,
+            max_depenetration_velocity=1.0,
+        ),
+        articulation_props=sim_utils.ArticulationRootPropertiesCfg(
+            enabled_self_collisions=True,
+            solver_position_iteration_count=8,
+            solver_velocity_iteration_count=4,
+        ),
+        joint_drive=sim_utils.UrdfConverterCfg.JointDriveCfg(
+            gains=sim_utils.UrdfConverterCfg.JointDriveCfg.PDGainsCfg(stiffness=0, damping=0)
+        ),
+    ),
+    init_state=ArticulationCfg.InitialStateCfg(
+        pos=(0.0, 0.0, 0.76),
+        joint_pos={
+            ".*_hip_pitch_joint": -0.312,
+            ".*_knee_joint": 0.669,
+            ".*_ankle_pitch_joint": -0.363,
+            ".*_elbow_joint": 0.6,
+            "left_shoulder_roll_joint": 0.2,
+            "left_shoulder_pitch_joint": 0.2,
+            "right_shoulder_roll_joint": -0.2,
+            "right_shoulder_pitch_joint": 0.2,
+        },
+        joint_vel={".*": 0.0},
+    ),
+    soft_joint_pos_limit_factor=0.9,
+    actuators={
+        "legs": ImplicitActuatorCfg(
+            joint_names_expr=[
+                ".*_hip_yaw_joint",
+                ".*_hip_roll_joint",
+                ".*_hip_pitch_joint",
+                ".*_knee_joint",
+            ],
+            effort_limit_sim={
+                ".*_hip_yaw_joint": 88.0,
+                ".*_hip_roll_joint": 139.0,
+                ".*_hip_pitch_joint": 139.0,
+                ".*_knee_joint": 139.0,
+            },
+            velocity_limit_sim={
+                ".*_hip_yaw_joint": 32.0,
+                ".*_hip_roll_joint": 20.0,
+                ".*_hip_pitch_joint": 20.0,
+                ".*_knee_joint": 20.0,
+            },
+            stiffness={
+                ".*_hip_pitch_joint": STIFFNESS_7520_22,
+                ".*_hip_roll_joint": STIFFNESS_7520_22,
+                ".*_hip_yaw_joint": STIFFNESS_7520_14,
+                ".*_knee_joint": STIFFNESS_7520_22,
+            },
+            damping={
+                ".*_hip_pitch_joint": DAMPING_7520_22,
+                ".*_hip_roll_joint": DAMPING_7520_22,
+                ".*_hip_yaw_joint": DAMPING_7520_14,
+                ".*_knee_joint": DAMPING_7520_22,
+            },
+            armature={
+                ".*_hip_pitch_joint": ARMATURE_7520_22,
+                ".*_hip_roll_joint": ARMATURE_7520_22,
+                ".*_hip_yaw_joint": ARMATURE_7520_14,
+                ".*_knee_joint": ARMATURE_7520_22,
+            },
+        ),
+        "feet": ImplicitActuatorCfg(
+            effort_limit_sim=50.0,
+            velocity_limit_sim=37.0,
+            joint_names_expr=[".*_ankle_pitch_joint", ".*_ankle_roll_joint"],
+            stiffness=2.0 * STIFFNESS_5020,
+            damping=2.0 * DAMPING_5020,
+            armature=2.0 * ARMATURE_5020,
+        ),
+        "waist": ImplicitActuatorCfg(
+            effort_limit_sim=50,
+            velocity_limit_sim=37.0,
+            joint_names_expr=["waist_roll_joint", "waist_pitch_joint"],
+            stiffness=2.0 * STIFFNESS_5020,
+            damping=2.0 * DAMPING_5020,
+            armature=2.0 * ARMATURE_5020,
+        ),
+        "waist_yaw": ImplicitActuatorCfg(
+            effort_limit_sim=88,
+            velocity_limit_sim=32.0,
+            joint_names_expr=["waist_yaw_joint"],
+            stiffness=STIFFNESS_7520_14,
+            damping=DAMPING_7520_14,
+            armature=ARMATURE_7520_14,
+        ),
+        "arms": ImplicitActuatorCfg(
+            joint_names_expr=[
+                ".*_shoulder_pitch_joint",
+                ".*_shoulder_roll_joint",
+                ".*_shoulder_yaw_joint",
+                ".*_elbow_joint",
+                ".*_wrist_roll_joint",
+                ".*_wrist_pitch_joint",
+                ".*_wrist_yaw_joint",
+            ],
+            effort_limit_sim={
+                ".*_shoulder_pitch_joint": 25.0,
+                ".*_shoulder_roll_joint": 25.0,
+                ".*_shoulder_yaw_joint": 25.0,
+                ".*_elbow_joint": 25.0,
+                ".*_wrist_roll_joint": 25.0,
+                ".*_wrist_pitch_joint": 5.0,
+                ".*_wrist_yaw_joint": 5.0,
+            },
+            velocity_limit_sim={
+                ".*_shoulder_pitch_joint": 37.0,
+                ".*_shoulder_roll_joint": 37.0,
+                ".*_shoulder_yaw_joint": 37.0,
+                ".*_elbow_joint": 37.0,
+                ".*_wrist_roll_joint": 37.0,
+                ".*_wrist_pitch_joint": 22.0,
+                ".*_wrist_yaw_joint": 22.0,
+            },
+            stiffness={
+                ".*_shoulder_pitch_joint": STIFFNESS_5020,
+                ".*_shoulder_roll_joint": STIFFNESS_5020,
+                ".*_shoulder_yaw_joint": STIFFNESS_5020,
+                ".*_elbow_joint": STIFFNESS_5020,
+                ".*_wrist_roll_joint": STIFFNESS_5020,
+                ".*_wrist_pitch_joint": STIFFNESS_4010,
+                ".*_wrist_yaw_joint": STIFFNESS_4010,
+            },
+            damping={
+                ".*_shoulder_pitch_joint": DAMPING_5020,
+                ".*_shoulder_roll_joint": DAMPING_5020,
+                ".*_shoulder_yaw_joint": DAMPING_5020,
+                ".*_elbow_joint": DAMPING_5020,
+                ".*_wrist_roll_joint": DAMPING_5020,
+                ".*_wrist_pitch_joint": DAMPING_4010,
+                ".*_wrist_yaw_joint": DAMPING_4010,
+            },
+            armature={
+                ".*_shoulder_pitch_joint": ARMATURE_5020,
+                ".*_shoulder_roll_joint": ARMATURE_5020,
+                ".*_shoulder_yaw_joint": ARMATURE_5020,
+                ".*_elbow_joint": ARMATURE_5020,
+                ".*_wrist_roll_joint": ARMATURE_5020,
+                ".*_wrist_pitch_joint": ARMATURE_4010,
+                ".*_wrist_yaw_joint": ARMATURE_4010,
+            },
+        ),
+    },
+)
+
+G1_MODEL_12_ACTION_SCALE = {}
+for a in G1_CYLINDER_MODEL_12_DEX_CFG.actuators.values():
+    e = a.effort_limit_sim
+    s = a.stiffness
+    names = a.joint_names_expr
+    if not isinstance(e, dict):
+        e = dict.fromkeys(names, e)
+    if not isinstance(s, dict):
+        s = dict.fromkeys(names, s)
+    for n in names:
+        if n in e and n in s and s[n]:
+            G1_MODEL_12_ACTION_SCALE[n] = 0.25 * e[n] / s[n]
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/h2.py b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/h2.py
new file mode 100644
index 0000000000000000000000000000000000000000..ae58e3a140377d435dd05622cb6f7fbf8b7baf45
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/envs/manager_env/robots/h2.py
@@ -0,0 +1,387 @@
+from isaaclab.actuators import ImplicitActuatorCfg
+from isaaclab.assets.articulation import ArticulationCfg
+import isaaclab.sim as sim_utils
+
+ASSET_DIR = "gear_sonic/data/assets"
+
+ARMATURE_5020 = 0.003609725
+ARMATURE_7520_14 = 0.010177520
+ARMATURE_7520_22 = 0.025101925
+ARMATURE_4010 = 0.00425
+
+NATURAL_FREQ = 10 * 2.0 * 3.1415926535  # 10Hz
+DAMPING_RATIO = 2.0
+
+STIFFNESS_5020 = ARMATURE_5020 * NATURAL_FREQ**2
+STIFFNESS_7520_14 = ARMATURE_7520_14 * NATURAL_FREQ**2
+STIFFNESS_7520_22 = ARMATURE_7520_22 * NATURAL_FREQ**2
+STIFFNESS_4010 = ARMATURE_4010 * NATURAL_FREQ**2
+
+DAMPING_5020 = 2.0 * DAMPING_RATIO * ARMATURE_5020 * NATURAL_FREQ
+DAMPING_7520_14 = 2.0 * DAMPING_RATIO * ARMATURE_7520_14 * NATURAL_FREQ
+DAMPING_7520_22 = 2.0 * DAMPING_RATIO * ARMATURE_7520_22 * NATURAL_FREQ
+DAMPING_4010 = 2.0 * DAMPING_RATIO * ARMATURE_4010 * NATURAL_FREQ
+
+H2_ISAACLAB_TO_MUJOCO_MAPPING = {}
+
+H2_ISAACLAB_JOINTS = [
+    "pelvis",
+    "left_hip_pitch_link",
+    "right_hip_pitch_link",
+    "waist_yaw_link",
+    "left_hip_roll_link",
+    "right_hip_roll_link",
+    "waist_roll_link",
+    "left_hip_yaw_link",
+    "right_hip_yaw_link",
+    "torso_link",
+    "left_knee_link",
+    "right_knee_link",
+    "head_pitch_link",
+    "left_shoulder_pitch_link",
+    "right_shoulder_pitch_link",
+    "left_ankle_roll_link",
+    "right_ankle_roll_link",
+    "head_yaw_link",
+    "left_shoulder_roll_link",
+    "right_shoulder_roll_link",
+    "left_ankle_pitch_link",
+    "right_ankle_pitch_link",
+    "left_shoulder_yaw_link",
+    "right_shoulder_yaw_link",
+    "left_elbow_link",
+    "right_elbow_link",
+    "left_wrist_roll_link",
+    "right_wrist_roll_link",
+    "left_wrist_pitch_link",
+    "right_wrist_pitch_link",
+    "left_wrist_yaw_link",
+    "right_wrist_yaw_link",
+]
+H2_ISAACLAB_TO_MUJOCO_DOF = [
+    0,
+    3,
+    6,
+    9,
+    14,
+    19,
+    1,
+    4,
+    7,
+    10,
+    15,
+    20,
+    2,
+    5,
+    8,
+    11,
+    16,
+    12,
+    17,
+    21,
+    23,
+    25,
+    27,
+    29,
+    13,
+    18,
+    22,
+    24,
+    26,
+    28,
+    30,
+]
+H2_MUJOCO_TO_ISAACLAB_DOF = [
+    0,
+    6,
+    12,
+    1,
+    7,
+    13,
+    2,
+    8,
+    14,
+    3,
+    9,
+    15,
+    17,
+    24,
+    4,
+    10,
+    16,
+    18,
+    25,
+    5,
+    11,
+    19,
+    26,
+    20,
+    27,
+    21,
+    28,
+    22,
+    29,
+    23,
+    30,
+]
+H2_ISAACLAB_TO_MUJOCO_BODY = [
+    0,
+    1,
+    4,
+    7,
+    10,
+    15,
+    20,
+    2,
+    5,
+    8,
+    11,
+    16,
+    21,
+    3,
+    6,
+    9,
+    12,
+    17,
+    13,
+    18,
+    22,
+    24,
+    26,
+    28,
+    30,
+    14,
+    19,
+    23,
+    25,
+    27,
+    29,
+    31,
+]
+H2_MUJOCO_TO_ISAACLAB_BODY = [
+    0,
+    1,
+    7,
+    13,
+    2,
+    8,
+    14,
+    3,
+    9,
+    15,
+    4,
+    10,
+    16,
+    18,
+    25,
+    5,
+    11,
+    17,
+    19,
+    26,
+    6,
+    12,
+    20,
+    27,
+    21,
+    28,
+    22,
+    29,
+    23,
+    30,
+    24,
+    31,
+]
+
+H2_ISAACLAB_TO_MUJOCO_MAPPING = {
+    "isaaclab_joints": H2_ISAACLAB_JOINTS,
+    "isaaclab_to_mujoco_dof": H2_ISAACLAB_TO_MUJOCO_DOF,
+    "mujoco_to_isaaclab_dof": H2_MUJOCO_TO_ISAACLAB_DOF,
+    "isaaclab_to_mujoco_body": H2_ISAACLAB_TO_MUJOCO_BODY,
+    "mujoco_to_isaaclab_body": H2_MUJOCO_TO_ISAACLAB_BODY,
+}
+
+
+# H2 Robot Configuration
+H2_CFG = ArticulationCfg(
+    spawn=sim_utils.UrdfFileCfg(
+        fix_base=False,
+        replace_cylinders_with_capsules=True,
+        asset_path=f"{ASSET_DIR}/robot_description/urdf/h2/h2.urdf",
+        activate_contact_sensors=True,
+        rigid_props=sim_utils.RigidBodyPropertiesCfg(
+            disable_gravity=False,
+            retain_accelerations=False,
+            linear_damping=0.0,
+            angular_damping=0.0,
+            max_linear_velocity=1000.0,
+            max_angular_velocity=1000.0,
+            max_depenetration_velocity=1.0,
+        ),
+        articulation_props=sim_utils.ArticulationRootPropertiesCfg(
+            enabled_self_collisions=True,
+            solver_position_iteration_count=8,
+            solver_velocity_iteration_count=4,
+        ),
+        joint_drive=sim_utils.UrdfConverterCfg.JointDriveCfg(
+            gains=sim_utils.UrdfConverterCfg.JointDriveCfg.PDGainsCfg(stiffness=0, damping=0)
+        ),
+    ),
+    init_state=ArticulationCfg.InitialStateCfg(
+        pos=(0.0, 0.0, 1.04),
+        joint_pos={
+            ".*_hip_pitch_joint": -0.312,
+            ".*_knee_joint": 0.669,
+            ".*_ankle_pitch_joint": -0.363,
+            ".*_elbow_joint": 0.6,
+            "left_shoulder_roll_joint": 0.2,
+            "left_shoulder_pitch_joint": 0.2,
+            "right_shoulder_roll_joint": -0.2,
+            "right_shoulder_pitch_joint": 0.2,
+        },
+        joint_vel={".*": 0.0},
+    ),
+    soft_joint_pos_limit_factor=0.9,
+    actuators={
+        "legs": ImplicitActuatorCfg(
+            joint_names_expr=[
+                ".*_hip_yaw_joint",
+                ".*_hip_roll_joint",
+                ".*_hip_pitch_joint",
+                ".*_knee_joint",
+            ],
+            effort_limit_sim={
+                ".*_hip_yaw_joint": 264.0,
+                ".*_hip_roll_joint": 417.0,
+                ".*_hip_pitch_joint": 417.0,
+                ".*_knee_joint": 417.0,
+            },
+            velocity_limit_sim={
+                ".*_hip_yaw_joint": 32.0,
+                ".*_hip_roll_joint": 20.0,
+                ".*_hip_pitch_joint": 20.0,
+                ".*_knee_joint": 20.0,
+            },
+            stiffness={
+                ".*_hip_pitch_joint": STIFFNESS_7520_22,
+                ".*_hip_roll_joint": STIFFNESS_7520_22,
+                ".*_hip_yaw_joint": STIFFNESS_7520_14,
+                ".*_knee_joint": STIFFNESS_7520_22,
+            },
+            damping={
+                ".*_hip_pitch_joint": DAMPING_7520_22,
+                ".*_hip_roll_joint": DAMPING_7520_22,
+                ".*_hip_yaw_joint": DAMPING_7520_14,
+                ".*_knee_joint": DAMPING_7520_22,
+            },
+            armature={
+                ".*_hip_pitch_joint": ARMATURE_7520_22,
+                ".*_hip_roll_joint": ARMATURE_7520_22,
+                ".*_hip_yaw_joint": ARMATURE_7520_14,
+                ".*_knee_joint": ARMATURE_7520_22,
+            },
+        ),
+        "feet": ImplicitActuatorCfg(
+            effort_limit_sim=150.0,
+            velocity_limit_sim=37.0,
+            joint_names_expr=[".*_ankle_pitch_joint", ".*_ankle_roll_joint"],
+            stiffness=2.0 * STIFFNESS_5020,
+            damping=2.0 * DAMPING_5020,
+            armature=2.0 * ARMATURE_5020,
+        ),
+        "waist": ImplicitActuatorCfg(
+            effort_limit_sim=150.0,
+            velocity_limit_sim=37.0,
+            joint_names_expr=["waist_roll_joint", "waist_pitch_joint"],
+            stiffness=2.0 * STIFFNESS_5020,
+            damping=2.0 * DAMPING_5020,
+            armature=2.0 * ARMATURE_5020,
+        ),
+        "waist_yaw": ImplicitActuatorCfg(
+            effort_limit_sim=264.0,
+            velocity_limit_sim=32.0,
+            joint_names_expr=["waist_yaw_joint"],
+            stiffness=STIFFNESS_7520_14,
+            damping=DAMPING_7520_14,
+            armature=ARMATURE_7520_14,
+        ),
+        "head": ImplicitActuatorCfg(
+            effort_limit_sim=150.0,
+            velocity_limit_sim=37.0,
+            joint_names_expr=["head_pitch_joint", "head_yaw_joint"],
+            stiffness=2.0 * STIFFNESS_5020,
+            damping=2.0 * DAMPING_5020,
+            armature=2.0 * ARMATURE_5020,
+        ),
+        "arms": ImplicitActuatorCfg(
+            joint_names_expr=[
+                ".*_shoulder_pitch_joint",
+                ".*_shoulder_roll_joint",
+                ".*_shoulder_yaw_joint",
+                ".*_elbow_joint",
+                ".*_wrist_roll_joint",
+                ".*_wrist_pitch_joint",
+                ".*_wrist_yaw_joint",
+            ],
+            effort_limit_sim={
+                ".*_shoulder_pitch_joint": 75.0,
+                ".*_shoulder_roll_joint": 75.0,
+                ".*_shoulder_yaw_joint": 75.0,
+                ".*_elbow_joint": 75.0,
+                ".*_wrist_roll_joint": 75.0,
+                ".*_wrist_pitch_joint": 15.0,
+                ".*_wrist_yaw_joint": 15.0,
+            },
+            velocity_limit_sim={
+                ".*_shoulder_pitch_joint": 37.0,
+                ".*_shoulder_roll_joint": 37.0,
+                ".*_shoulder_yaw_joint": 37.0,
+                ".*_elbow_joint": 37.0,
+                ".*_wrist_roll_joint": 37.0,
+                ".*_wrist_pitch_joint": 22.0,
+                ".*_wrist_yaw_joint": 22.0,
+            },
+            stiffness={
+                ".*_shoulder_pitch_joint": STIFFNESS_5020,
+                ".*_shoulder_roll_joint": STIFFNESS_5020,
+                ".*_shoulder_yaw_joint": STIFFNESS_5020,
+                ".*_elbow_joint": STIFFNESS_5020,
+                ".*_wrist_roll_joint": STIFFNESS_5020,
+                ".*_wrist_pitch_joint": STIFFNESS_4010,
+                ".*_wrist_yaw_joint": STIFFNESS_4010,
+            },
+            damping={
+                ".*_shoulder_pitch_joint": DAMPING_5020,
+                ".*_shoulder_roll_joint": DAMPING_5020,
+                ".*_shoulder_yaw_joint": DAMPING_5020,
+                ".*_elbow_joint": DAMPING_5020,
+                ".*_wrist_roll_joint": DAMPING_5020,
+                ".*_wrist_pitch_joint": DAMPING_4010,
+                ".*_wrist_yaw_joint": DAMPING_4010,
+            },
+            armature={
+                ".*_shoulder_pitch_joint": ARMATURE_5020,
+                ".*_shoulder_roll_joint": ARMATURE_5020,
+                ".*_shoulder_yaw_joint": ARMATURE_5020,
+                ".*_elbow_joint": ARMATURE_5020,
+                ".*_wrist_roll_joint": ARMATURE_5020,
+                ".*_wrist_pitch_joint": ARMATURE_4010,
+                ".*_wrist_yaw_joint": ARMATURE_4010,
+            },
+        ),
+    },
+)
+
+# H2 Action Scale
+H2_ACTION_SCALE = {}
+for a in H2_CFG.actuators.values():
+    e = a.effort_limit_sim
+    s = a.stiffness
+    names = a.joint_names_expr
+    if not isinstance(e, dict):
+        e = dict.fromkeys(names, e)
+    if not isinstance(s, dict):
+        s = dict.fromkeys(names, s)
+    for n in names:
+        if n in e and n in s and s[n]:
+            H2_ACTION_SCALE[n] = 0.25 * e[n] / s[n]
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/wrapper/__init__.py b/GR00T-WholeBodyControl/gear_sonic/envs/wrapper/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/wrapper/__pycache__/__init__.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/wrapper/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9133ad7f7316ee6503587d04a1e58920d6095d79
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/wrapper/__pycache__/__init__.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/wrapper/__pycache__/__init__.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/envs/wrapper/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..38dd66a287be03c8b677e5298df5b2f4920a822d
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/envs/wrapper/__pycache__/__init__.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/envs/wrapper/manager_env_wrapper.py b/GR00T-WholeBodyControl/gear_sonic/envs/wrapper/manager_env_wrapper.py
new file mode 100644
index 0000000000000000000000000000000000000000..506c1fc8f18c9fa55944df2a4525c5d2d20d830e
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/envs/wrapper/manager_env_wrapper.py
@@ -0,0 +1,2100 @@
+from typing import TYPE_CHECKING  # noqa: I001
+
+import numpy as np
+from omegaconf import OmegaConf
+import omni
+from pxr import Gf, UsdGeom
+import torch
+from loguru import logger
+from gear_sonic.trl.utils.common import custom_instantiate
+
+if TYPE_CHECKING:
+    from isaaclab.envs.manager_based_rl_env import ManagerBasedEnv
+
+# Import joint index functions (single source of truth)
+from gear_sonic.envs.env_utils.joint_utils import get_body_joint_indices, get_hand_joint_indices
+
+# Import visualization markers for contact point visualization
+try:
+    from isaaclab.markers import VisualizationMarkers, VisualizationMarkersCfg
+    import isaaclab.sim as sim_utils
+
+    VISUALIZATION_AVAILABLE = True
+except ImportError:
+    VISUALIZATION_AVAILABLE = False
+
+
+class ManagerEnvWrapper:
+    def __init__(self, env: "ManagerBasedEnv", config):
+        env.wrapper = self
+        self.env = env
+        self.config = OmegaConf.create(config)
+        self.device = env.device
+        self.viewer_focused = True
+        self.is_manager_env = True
+        if hasattr(self.env, "num_envs"):
+            self.num_envs = self.env.num_envs
+        else:
+            self.num_envs = self.env.env.unwrapped.num_envs
+        try:
+            self.motion_command = env.command_manager.get_term("motion")
+            self._motion_lib = self.motion_command.motion_lib
+        except:  # noqa: E722
+            logger.info("No motion lib found")
+            self.motion_command = None
+            self._motion_lib = None
+        try:
+            self.force_command = env.command_manager.get_term("force")
+        except:  # noqa: E722
+            self.force_command = None
+        self.is_evaluating = False
+        self.start_idx = 0
+        self._last_predicted_object_pos = (
+            None  # For debug visualization of predicted object position
+        )
+        if not self.config.headless:
+            self.setup_keyboard()
+
+        # Action visualization toggle and state
+        self.turn_on_visualization = bool(self.config.get("turn_on_visualization", False)) and (
+            not self.config.get("headless", False)
+        )
+        self._viz_every_n_steps = int(self.config.get("viz_every_n_steps", 1))
+        self._plot_action_dim = int(self.config.get("action_plot_dim", 29))
+        clip_default = self.config.get("action_clip_value", 1.0)  # noqa: F841
+        self._action_ylim = float(self.config.get("action_plot_ylim", 10.0))
+        self._plot_window = int(self.config.get("action_plot_window", 200))
+        self._step_counter = 0
+        self._action_fig = None
+        self._action_lines = None
+        self._action_hist = None
+        self._hist_idx = 0
+
+        self._blit_background = None
+        self._blit_supported = True
+        self._blit_refresh_interval = int(self.config.get("plot_blit_refresh_interval", 120))
+
+        # Initialize action transform module from config
+        self.action_transform_module = None
+        self._needs_policy_atm = False  # Default: use policy obs directly
+        self._policy_atm_indices = None
+
+        # Initialize joint index mappings (for DOF mismatch handling in replay and step)
+        self._body_joint_indices = None
+        self._hand_joint_indices = None
+        self._setup_replay_joint_indices()  # Setup for replay mode
+
+        # Initialize finger primitive support
+        self._use_finger_primitive = self.config.get("use_finger_primitive", False)
+        self._finger_primitive_map = None
+        if self._use_finger_primitive:
+            self._setup_finger_primitives()
+
+        # Latent residual mode: policy outputs residual added to token latent space
+        self._use_latent_residual = self.config.get("use_latent_residual", False)
+
+        # Latent residual options (only used when use_latent_residual=True)
+        self._latent_residual_mode = self.config.get("latent_residual_mode", "post_quantization")
+        self._latent_residual_scale = self.config.get("latent_residual_scale", 1.0)
+
+        # Student direct latent mode: policy outputs FULL latent (not residual)
+        # This is used for vision student policies that learned to output full latent
+        # directly, bypassing the ATM encoder entirely at inference time.
+        # The 64-dim output goes directly to ATM decoder (no encoding step).
+        self._use_student_direct_latent = self.config.get("use_student_direct_latent", False)
+
+        if self._use_latent_residual:
+            logger.info(
+                f"Latent residual enabled: mode={self._latent_residual_mode}, "
+                f"scale={self._latent_residual_scale}"
+            )
+
+        if self._use_student_direct_latent:
+            logger.info(
+                "Student direct latent mode enabled: policy output goes directly to ATM decoder "
+                "(no encoding step)"
+            )
+
+        # Camera extrinsics randomization state
+        self._camera_extrinsics_randomized = False
+        self._camera_base_transforms = {}  # Store original transforms per env
+
+        action_transform_module_cfg = self.config.get("action_transform_module_cfg", None)
+
+        if action_transform_module_cfg is not None:
+            # Load configs from exported YAML file
+            with open(action_transform_module_cfg) as f:
+                exported_config = OmegaConf.load(f)
+
+            env_config = exported_config.get("env_config", {})
+            algo_config = exported_config.get("algo_config", {})
+
+            self.action_transform_module = custom_instantiate(
+                algo_config.actor, env_config=env_config, algo_config=algo_config, _resolve=False
+            ).to(self.device)
+            logger.info(f"Initialized action_transform_module from config: {action_transform_module_cfg}")
+
+            # Load checkpoint if provided
+            action_transform_module_checkpoint = self.config.get(
+                "action_transform_module_checkpoint", None
+            )
+            if action_transform_module_checkpoint is not None:
+                # Compatibility shim: checkpoints saved with TRL < 0.28.0 reference
+                # trl.trainer.utils.OnlineTrainerState, which was moved in 0.28.0
+                try:
+                    from trl.experimental.ppo.ppo_trainer import OnlineTrainerState, exact_div
+                    import trl.trainer.utils
+
+                    trl.trainer.utils.OnlineTrainerState = OnlineTrainerState
+                    trl.trainer.utils.exact_div = exact_div
+                except ImportError:
+                    pass
+                checkpoint = torch.load(
+                    action_transform_module_checkpoint, map_location=self.device, weights_only=False
+                )
+                self.action_transform_module.load_state_dict(checkpoint["policy_state_dict"])
+                logger.info(
+                    f"Loaded action_transform_module checkpoint: {action_transform_module_checkpoint}"
+                )
+
+            # Precompute tokenizer observation indices for meta_action target
+            self._tokenizer_obs_indices = self._compute_tokenizer_obs_indices()
+
+            # Setup policy_atm observations for action_transform_module if DOF mismatch exists
+            self._setup_policy_atm(env_config, algo_config)
+
+        try:
+            self.viewer_focused = True
+            self.env.viewport_camera_controller.update_view_to_world()
+            self.env.viewport_camera_controller.update_view_to_asset_root("robot")
+        except:  # noqa: E722
+            self.viewer_focused = False
+
+    def _setup_replay_joint_indices(self):
+        """Setup joint indices for replay mode if robot has more DOFs than motion lib (29)."""
+        if self.env.scene["robot"].num_joints > 29:
+            self._setup_action_joint_indices()
+
+    def _compute_tokenizer_obs_indices(self):
+        """Compute start and end indices for each tokenizer observation in the flattened tensor."""
+        if self.action_transform_module is None:
+            return {}
+
+        tokenizer_obs_names = self.action_transform_module.actor_module.tokenizer_obs_names
+        tokenizer_obs_dims = self.action_transform_module.actor_module.tokenizer_obs_dims
+
+        indices = {}
+        current_index = 0
+        for name in tokenizer_obs_names:
+            all_dim = int(np.prod(tokenizer_obs_dims[name]))
+            indices[name] = (current_index, current_index + all_dim)
+            current_index += all_dim
+
+        return indices
+
+    def _setup_policy_atm(self, env_config, algo_config):  # noqa: ARG002
+        """Setup policy_atm for action_transform_module when robot has more DOFs than ATM expects."""
+        atm_num_joints = env_config.get("robot", {}).get("actions_dim") or env_config.get(
+            "robot", {}
+        ).get("num_joints", 29)
+        self._needs_policy_atm = self.config.get("needs_policy_atm", True)
+        self._atm_num_joints = atm_num_joints
+        self._current_num_joints = self.env.scene["robot"].num_joints
+        if self._current_num_joints > atm_num_joints:
+            assert (
+                self._needs_policy_atm
+            ), "Robot has more DOFs than ATM expects, but needs_policy_atm is False"
+
+        has_policy_atm = (
+            hasattr(self.env, "observation_manager")
+            and "policy_atm" in self.env.observation_manager._group_obs_term_names  # noqa: SLF001
+        )
+        self._use_policy_atm_group = self._needs_policy_atm and has_policy_atm
+
+        if self._use_policy_atm_group:
+            self._setup_action_joint_indices()
+
+    def _setup_action_joint_indices(self):
+        """Compute joint indices for mapping body (29 DOF) and hand (14 DOF) actions."""
+        if self._body_joint_indices is not None:
+            return
+
+        robot = self.env.scene["robot"]
+        self._body_joint_indices = get_body_joint_indices(robot)
+        self._hand_joint_indices = get_hand_joint_indices(robot)
+
+    def _setup_finger_primitives(self):
+        """Setup finger primitive action mapping from config.
+
+        Finger primitives allow the policy to output 2 actions (left/right gripper)
+        instead of 14 individual finger joint actions. Each primitive action
+        interpolates between open (pos_0) and closed (pos_1) positions.
+        """
+        primitive_cfg = self.config.get("finger_primitive", {})
+        primitive_action_map = primitive_cfg.get("primitive_action_map", {})
+
+        if not primitive_action_map:
+            logger.info("Warning: use_finger_primitive=True but no primitive_action_map defined")
+            self._use_finger_primitive = False
+            return
+
+        robot = self.env.scene["robot"]
+        joint_names = robot.joint_names
+
+        self._finger_primitive_map = []
+        for action_name in sorted(primitive_action_map.keys()):
+            prim_cfg = primitive_action_map[action_name]
+            dof_names = list(prim_cfg.get("dof_names", []))
+            pos_0 = torch.tensor(prim_cfg.get("pos_0", []), device=self.device, dtype=torch.float32)
+            pos_1 = torch.tensor(prim_cfg.get("pos_1", []), device=self.device, dtype=torch.float32)
+
+            # Find indices in the hand joints
+            dof_idx = []
+            for dof_name in dof_names:
+                if dof_name in joint_names:
+                    dof_idx.append(joint_names.index(dof_name))
+                else:
+                    logger.info(f"Warning: DOF {dof_name} not found in robot joints")
+            # Support both "mode" (new) and "discrete" (legacy) config keys
+            mode = prim_cfg.get("mode", None)
+            if mode is None:
+                # Legacy support: discrete=True → mode="discrete", discrete=False → mode="linear"
+                mode = "discrete" if prim_cfg.get("discrete", False) else "linear"
+
+            self._finger_primitive_map.append(
+                {
+                    "action_name": action_name,
+                    "dof_names": dof_names,
+                    "dof_idx": dof_idx,
+                    "pos_0": pos_0,
+                    "pos_1": pos_1,
+                    "mode": mode,
+                }
+            )
+            logger.info(f"Finger primitive '{action_name}': {len(dof_names)} DOFs, mode={mode}")
+
+        self._num_finger_primitives = len(self._finger_primitive_map)
+        logger.info(f"Initialized {self._num_finger_primitives} finger primitives")
+
+    def _convert_primitive_to_finger_actions(self, primitive_actions: torch.Tensor) -> torch.Tensor:
+        """Convert primitive actions (num_envs, num_primitives) to finger joint targets (num_envs, 14).
+
+        Follows groot_backup implementation:
+        - "linear" mode: maps [-1, 1] → [0, 1] via (x + 1) / 2, then lerps between pos_0 and pos_1
+        - "discrete" mode: action >= 0 → pos_1 (closed), action < 0 → pos_0 (open)
+
+        Args:
+            primitive_actions: Tensor of shape (num_envs, num_primitives), values in [-1, 1]
+
+        Returns:
+            Tensor of shape (num_envs, num_finger_joints) with joint position targets
+        """
+        num_envs = primitive_actions.shape[0]
+        num_finger_joints = (
+            len(self._hand_joint_indices) if self._hand_joint_indices is not None else 14
+        )
+        finger_targets = torch.zeros(
+            num_envs, num_finger_joints, device=self.device, dtype=primitive_actions.dtype
+        )
+
+        for i, prim_cfg in enumerate(self._finger_primitive_map):
+            action = primitive_actions[:, i]  # (num_envs,), values in [-1, 1]
+            p0 = prim_cfg["pos_0"]
+            p1 = prim_cfg["pos_1"]
+            dof_idx = prim_cfg["dof_idx"]
+            mode = prim_cfg.get("mode", "linear")
+            if mode == "linear":
+                # Clamp to [-1, 1] then map to [0, 1]
+                action_clamped = action.clamp(min=-1.0, max=1.0)
+                t = (action_clamped + 1.0) / 2.0  # (num_envs,)
+                joint_targets = torch.lerp(p0.unsqueeze(0), p1.unsqueeze(0), t.unsqueeze(1))
+            elif mode == "discrete":
+                # action >= 0 → closed (pos_1), action < 0 → open (pos_0)
+                joint_targets = torch.where(
+                    action.unsqueeze(1) >= 0, p1.unsqueeze(0), p0.unsqueeze(0)
+                )
+            else:
+                raise ValueError(f"Invalid finger primitive mode: {mode}")
+
+            # Map to the correct indices in finger_targets
+            # dof_idx are absolute joint indices, need to convert to relative hand indices
+            for j, abs_idx in enumerate(dof_idx):
+                if self._hand_joint_indices is not None:
+                    # Convert tensor to list if needed for .index() lookup
+                    hand_indices_list = (
+                        self._hand_joint_indices.tolist()
+                        if isinstance(self._hand_joint_indices, torch.Tensor)
+                        else self._hand_joint_indices
+                    )
+                    rel_idx = hand_indices_list.index(abs_idx)
+                    finger_targets[:, rel_idx] = joint_targets[:, j]
+                else:
+                    # Fallback: assume hand joints are at the end
+                    finger_targets[:, abs_idx - 29] = joint_targets[:, j]
+
+        return finger_targets
+
+    def _prepare_obs_for_action_transform_module(self, obs_dict):
+        """Use policy_atm observations for ATM if DOF mismatch exists, else use policy."""
+        if not self._use_policy_atm_group or "policy_atm" not in obs_dict:
+            atm_obs_dict = obs_dict.copy()
+        else:
+            atm_obs_dict = obs_dict.copy()
+            atm_obs_dict["actor_obs"] = atm_obs_dict["policy_atm"]
+
+        # Ensure all observations have a sequence dimension [num_envs, seq_len, dim]
+        # The action_transform_module expects 3D tensors
+        for k, v in atm_obs_dict.items():
+            if isinstance(v, torch.Tensor) and v.dim() == 2:
+                atm_obs_dict[k] = v.unsqueeze(1)  # Add seq_len=1 dimension
+
+        return atm_obs_dict
+
+    def reset_all(self, global_rank=0):  # noqa: ARG002
+        return self.reset()
+
+    def process_raw_obs(self, obs, flatten_dict_obs):
+        new_obs = {
+            "actor_obs": obs["policy"],
+            "critic_obs": obs["critic"],
+        }
+        for k, v in obs.items():
+            if k not in ["policy", "critic"]:
+                if isinstance(v, dict) and flatten_dict_obs:
+                    if k == "height_map":
+                        # Special case: do not flatten height map
+                        new_obs[k] = v["height_map"]
+                        continue
+                    if k == "camera_rgb":
+                        # Special case: do not flatten camera RGB image
+                        # Keep original shape [B, H, W, C] for vision encoder
+                        new_obs[k] = v["camera_rgb"]
+                        continue
+                    obs_names = self.env.observation_manager._group_obs_term_names[k]  # noqa: SLF001
+                    new_obs[k] = torch.cat(
+                        [v[obs_name].reshape(v[obs_name].shape[0], -1) for obs_name in obs_names],
+                        dim=-1,
+                    )
+                else:
+                    new_obs[k] = v
+        return new_obs
+
+    def reset(self, flatten_dict_obs=True):
+        obs, info = self.env.reset()
+        new_obs = self.process_raw_obs(obs, flatten_dict_obs)
+        # Initialize success_lift to False for all envs after reset (used unconditionally in step())
+        self.env.success_lift = torch.zeros(
+            self.env.num_envs, dtype=torch.bool, device=self.env.device
+        )
+        if self.action_transform_module is not None:
+            # Store obs for action_transform_module when obs_dict is not provided in step()
+            self._last_obs_dict = new_obs
+            # Initialize last meta action buffer (policy output: latent + primitives)
+            # meta_action_dim = tokenizer_action_dim + hand_action_dim (e.g., 64 + 2 = 66)
+            meta_action_dim = self.config.get("meta_action_dim", 66)
+            self.env._last_meta_action = torch.zeros(  # noqa: SLF001
+                self.env.num_envs, meta_action_dim, dtype=torch.float32, device=self.env.device
+            )
+            # Previous meta action buffer for meta_action_rate_l2 reward (token smoothness)
+            self.env._prev_meta_action = torch.zeros(  # noqa: SLF001
+                self.env.num_envs, meta_action_dim, dtype=torch.float32, device=self.env.device
+            )
+            # Full latent buffers for full_latent_rate_l2 reward (decoder input smoothness)
+            # tokenizer_action_dim = latent_dim (e.g., 64 = num_tokens * token_dim)
+            tokenizer_action_dim = self.config.get("tokenizer_action_dim", 64)
+            self.env._full_latent = torch.zeros(  # noqa: SLF001
+                self.env.num_envs, tokenizer_action_dim, dtype=torch.float32, device=self.env.device
+            )
+            self.env._prev_full_latent = torch.zeros(  # noqa: SLF001
+                self.env.num_envs, tokenizer_action_dim, dtype=torch.float32, device=self.env.device
+            )
+
+            # Apply camera extrinsics randomization on every reset
+            self.apply_random_camera_extrinsics()
+
+        return new_obs
+
+    def apply_random_camera_extrinsics(self):
+        """Apply per-environment random camera extrinsics (position and rotation offsets).
+
+        Reads randomization ranges from config:
+        - cameras.camera_extrinsics_randomization: Enable/disable switch (default: False)
+        - cameras.camera_pos_rand_range: ±meters for x, y, z position
+        - cameras.camera_roll_rand_range: ±radians for roll
+        - cameras.camera_pitch_rand_range: ±radians for pitch
+        - cameras.camera_yaw_rand_range: ±radians for yaw
+        """
+        cameras_config = self.config.get("cameras", {})
+
+        # Check if randomization is enabled via switch
+        if not cameras_config.get("camera_extrinsics_randomization", False):
+            return
+
+        # Get randomization ranges
+        pos_range = cameras_config.get("camera_pos_rand_range", 0.0)
+        roll_range = cameras_config.get("camera_roll_rand_range", 0.0)
+        pitch_range = cameras_config.get("camera_pitch_rand_range", 0.0)
+        yaw_range = cameras_config.get("camera_yaw_rand_range", 0.0)
+
+        # Check if any randomization values are non-zero
+        if pos_range == 0 and roll_range == 0 and pitch_range == 0 and yaw_range == 0:
+            return
+
+        # Get camera attached link from config
+        camera_attached_link = cameras_config.get("camera_attached_link", None)
+        if camera_attached_link is None:
+            logger.info("Skipping camera extrinsics randomization: no camera_attached_link configured")
+            return
+
+        # Only print details on first call
+        if not self._camera_extrinsics_randomized:
+            logger.info("Applying random camera extrinsics per environment:")
+            logger.info(f"  Position range: ±{pos_range*100:.1f}cm")
+            logger.info(f"  Roll range: ±{roll_range:.3f} rad ({np.degrees(roll_range):.1f}°)")
+            logger.info(f"  Pitch range: ±{pitch_range:.3f} rad ({np.degrees(pitch_range):.1f}°)")
+            logger.info(f"  Yaw range: ±{yaw_range:.3f} rad ({np.degrees(yaw_range):.1f}°)")
+
+        stage = omni.usd.get_context().get_stage()
+
+        for env_id in range(self.env.scene.num_envs):
+            camera_prim_path = f"/World/envs/env_{env_id}/Robot/{camera_attached_link}/ego_camera"
+            camera_prim = stage.GetPrimAtPath(camera_prim_path)
+
+            if not camera_prim.IsValid():
+                if env_id == 0:
+                    logger.info(f"  Warning: Camera prim not found at {camera_prim_path}")
+                continue
+
+            # Get current camera transform
+            xformable = UsdGeom.Xformable(camera_prim)
+
+            # On first call, store the original/base transforms
+            if env_id not in self._camera_base_transforms:
+                base_translate = None
+                base_orient = None
+                for op in xformable.GetOrderedXformOps():
+                    if op.GetOpType() == UsdGeom.XformOp.TypeTranslate:
+                        base_translate = Gf.Vec3d(op.Get())  # Make a copy
+                    elif op.GetOpType() == UsdGeom.XformOp.TypeOrient:
+                        base_orient = Gf.Quatd(op.Get())  # Make a copy
+                self._camera_base_transforms[env_id] = {
+                    "translate": base_translate,
+                    "orient": base_orient,
+                }
+
+            # Get base transforms
+            base_translate = self._camera_base_transforms[env_id]["translate"]
+            base_orient = self._camera_base_transforms[env_id]["orient"]
+
+            # Sample random deltas
+            if pos_range > 0:  # noqa: SIM108
+                pos_delta = np.random.uniform(-pos_range, pos_range, 3)  # noqa: NPY002
+            else:
+                pos_delta = np.zeros(3)
+
+            # Sample rotation deltas (roll, pitch, yaw)
+            roll_delta = np.random.uniform(-roll_range, roll_range) if roll_range > 0 else 0.0  # noqa: NPY002
+            pitch_delta = np.random.uniform(-pitch_range, pitch_range) if pitch_range > 0 else 0.0  # noqa: NPY002
+            yaw_delta = np.random.uniform(-yaw_range, yaw_range) if yaw_range > 0 else 0.0  # noqa: NPY002
+
+            # Apply position delta relative to BASE (not current)
+            if base_translate is not None:
+                new_pos = Gf.Vec3d(
+                    base_translate[0] + pos_delta[0],
+                    base_translate[1] + pos_delta[1],
+                    base_translate[2] + pos_delta[2],
+                )
+                # Find and update the translate op
+                for op in xformable.GetOrderedXformOps():
+                    if op.GetOpType() == UsdGeom.XformOp.TypeTranslate:
+                        op.Set(new_pos)
+                        break
+
+            # Apply rotation delta relative to BASE (not current)
+            if base_orient is not None and (roll_delta != 0 or pitch_delta != 0 or yaw_delta != 0):
+                # Convert euler deltas to quaternion
+                # Order: roll (X), pitch (Y), yaw (Z)
+                cr, sr = np.cos(roll_delta / 2), np.sin(roll_delta / 2)
+                cp, sp = np.cos(pitch_delta / 2), np.sin(pitch_delta / 2)
+                cy, sy = np.cos(yaw_delta / 2), np.sin(yaw_delta / 2)
+
+                # Quaternion from euler (ZYX convention)
+                delta_quat = Gf.Quatd(
+                    cr * cp * cy + sr * sp * sy,  # w
+                    sr * cp * cy - cr * sp * sy,  # x
+                    cr * sp * cy + sr * cp * sy,  # y
+                    cr * cp * sy - sr * sp * cy,  # z
+                )
+
+                # Compose: new = delta * BASE (not current!)
+                new_orient = delta_quat * base_orient
+
+                # Find and update the orient op
+                for op in xformable.GetOrderedXformOps():
+                    if op.GetOpType() == UsdGeom.XformOp.TypeOrient:
+                        op.Set(new_orient)
+                        break
+
+            # Store the random deltas for each environment (for projection in observations)
+            if not hasattr(self, "_camera_random_deltas"):
+                self._camera_random_deltas = {}
+            self._camera_random_deltas[env_id] = {
+                "pos_delta": pos_delta.tolist() if isinstance(pos_delta, np.ndarray) else [0, 0, 0],
+                "roll_delta": float(roll_delta),
+                "pitch_delta": float(pitch_delta),
+                "yaw_delta": float(yaw_delta),
+            }
+
+            if env_id == 0 and not self._camera_extrinsics_randomized:
+                logger.info(
+                    f"  Env 0: pos_delta={pos_delta}, rot_delta=[{roll_delta:.3f}, {pitch_delta:.3f}, {yaw_delta:.3f}]"  # noqa: E501
+                )
+
+        # Mark as initialized (for print suppression)
+        self._camera_extrinsics_randomized = True
+
+    def _decode_direct_latent(self, full_latent, atm_obs_dict):
+        """Decode full latent directly using ATM decoder (skip encoder).
+        Used for student rollout where policy outputs full latent.
+
+        Args:
+            full_latent: (batch, latent_dim) full latent from student policy
+            atm_obs_dict: observation dict for ATM
+
+        Returns:
+            body_actions: (batch, seq_len, action_dim) decoded actions
+        """  # noqa: D205
+        # Get proprioception for decoder
+        if "policy_atm" in atm_obs_dict:  # noqa: SIM108
+            proprioception = atm_obs_dict["policy_atm"]
+        else:
+            proprioception = atm_obs_dict["actor_obs"]
+
+        # Ensure sequence dimension
+        if proprioception.dim() == 2:
+            proprioception = proprioception.unsqueeze(1)
+
+        atm = self.action_transform_module.actor_module
+        return self._decode_direct_latent_batch(full_latent, proprioception, atm)
+
+    def _decode_direct_latent_batch(self, full_latent, proprioception, atm):
+        """Decode a batch of full latents directly using ATM decoder.
+
+        The student policy outputs pre-quantization values (latent + residual).
+        We apply quantization here before decoding to match the teacher's flow:
+        - Teacher: latent + residual -> quantize -> decode
+        - Student: direct_latent -> quantize -> decode
+
+        Args:
+            full_latent: (batch, latent_dim) full latent (pre-quantization)
+            proprioception: (batch, seq, proprio_dim) proprioception input
+            atm: ATM actor module
+
+        Returns:
+            body_actions: (batch, seq_len, action_dim) decoded actions
+        """
+        batch_size = full_latent.shape[0]
+
+        # Reshape latent for quantization: (batch, latent_dim) -> (batch, num_tokens, token_dim)
+        latent_reshaped = full_latent.view(batch_size, atm.max_num_tokens, atm.token_dim)
+
+        # Apply quantization (same as teacher's flow)
+        # This ensures student inference matches teacher: latent -> quantize -> decode
+        if atm.quantizer is not None:
+            quantized_codes, _ = atm.quantizer(latent_reshaped)
+            tokens_for_decode = quantized_codes.contiguous()
+        else:
+            tokens_for_decode = latent_reshaped
+
+        # Reshape for decoder: (batch, num_tokens, token_dim) -> (batch, 1, num_tokens, token_dim)
+        tokens_reshaped = tokens_for_decode.unsqueeze(1)
+        tokens_flattened = tokens_for_decode.view(batch_size, -1).unsqueeze(1)
+
+        # Prepare decode input
+        decode_input_dict = {
+            "token": tokens_reshaped,
+            "token_flattened": tokens_flattened,  # (batch, 1, latent_dim)
+            "proprioception": proprioception,
+        }
+
+        # Decode directly (skip encoding entirely)
+        decoded_output = atm.decode("g1_dyn", decode_input_dict)
+
+        # Get body actions from decoder output
+        # Note: "meta_action" is for special hierarchical policies, "action" is standard
+        body_actions = decoded_output.get("meta_action", decoded_output.get("action"))
+        if body_actions is None:
+            raise KeyError(
+                f"Decoder output missing 'action' or 'meta_action'. Keys: {decoded_output.keys()}"
+            )
+
+        return body_actions
+
+    def step(self, actions):
+        if self.action_transform_module is not None:
+            # Use provided obs_dict or fall back to stored obs from last reset/step
+            if "obs_dict" in actions:
+                obs_dict = actions["obs_dict"].copy()
+            else:
+                # Fallback for callbacks that don't provide obs_dict (e.g., im_eval)
+                obs_dict = getattr(self, "_last_obs_dict", None)
+                if obs_dict is None:
+                    raise ValueError(
+                        "action_transform_module requires obs_dict but none was provided or stored"
+                    )
+                obs_dict = obs_dict.copy() if isinstance(obs_dict, dict) else obs_dict
+            meta_actions = actions["actions"]
+            # Determine action mode: "direct_latent", "residual", or "mixed"
+            # Priority: 1) explicit action_mode in actions dict, 2) config flag
+            # During training: trainer sets action_mode explicitly
+            # During eval: fallback to config flags
+            action_mode = actions.get("action_mode", None)
+            if action_mode is None:
+                # Fallback for eval scripts that don't set action_mode
+                if self._use_student_direct_latent:
+                    action_mode = "direct_latent"
+                elif self._use_latent_residual:
+                    action_mode = "residual"
+                else:
+                    action_mode = "residual"  # Default to residual if nothing specified
+
+            # Shift meta action buffers for meta_action_rate_l2 reward (token smoothness)
+            self.env._prev_meta_action = self.env._last_meta_action.clone()  # noqa: SLF001
+            # Store meta action for observation (last policy output)
+            self.env._last_meta_action = meta_actions.clone()  # noqa: SLF001
+
+            atm_obs_dict = self._prepare_obs_for_action_transform_module(obs_dict)
+
+            # Split actions: first tokenizer_action_dim for tokenizer, rest for hands
+            tokenizer_action_dim = self.config.get("tokenizer_action_dim")
+            tokenizer_meta_actions = meta_actions[:, :tokenizer_action_dim]
+            hand_actions_raw = meta_actions[:, tokenizer_action_dim:]
+
+            # Override hand actions with motion data if configured
+            if self.config.get("use_motion_hand_actions", False):
+                motion_cmd = self.env.command_manager.get_term("motion")
+                left_action = motion_cmd.get_hand_action("left_hand")
+                right_action = motion_cmd.get_hand_action("right_hand")
+
+                if left_action is None or right_action is None:
+                    raise ValueError(
+                        "use_motion_hand_actions=True but hand_action_left/right not found in motion data. "
+                        "Ensure processed_robot_motions.pkl contains 'hand_action_left' and 'hand_action_right' arrays."  # noqa: E501
+                    )
+
+                # Use motion data directly: -1.0 = open, +1.0 = closed
+                # Threshold at 0 in _convert_primitive_to_finger_actions
+                hand_actions_raw = torch.stack([left_action, right_action], dim=-1)
+
+            # Convert primitive actions to finger joint targets if enabled
+            if self._use_finger_primitive and self._finger_primitive_map:
+                # Store raw primitive actions on env for reward computation (before clamping)
+                self.env._finger_primitive_actions_raw = hand_actions_raw  # noqa: SLF001
+                hand_actions = self._convert_primitive_to_finger_actions(hand_actions_raw)
+            else:
+                self.env._finger_primitive_actions_raw = None  # noqa: SLF001
+                hand_actions = hand_actions_raw
+
+            if action_mode == "direct_latent":
+                # Student direct latent mode: policy outputs FULL latent, not residual
+                # Skip ATM encoder entirely - go directly to decoder
+                body_actions = self._decode_direct_latent(tokenizer_meta_actions, atm_obs_dict)
+
+            elif action_mode == "residual":
+                # Teacher/residual mode: policy outputs residual that's added to ATM encoded tokens
+                # Apply scaling to residual before passing to ATM
+                scaled_residual = tokenizer_meta_actions * self._latent_residual_scale
+                # Add residual in latent/token space (after encoding, before decoding)
+                body_actions = self.action_transform_module(
+                    atm_obs_dict,
+                    latent_residual=scaled_residual,
+                    latent_residual_mode=self._latent_residual_mode,
+                )
+
+            elif action_mode == "mixed":
+                # Mixed rollout: some envs use teacher (residual), some use student (direct_latent)
+                # is_teacher_env is a boolean mask: True = teacher/residual, False = student/direct_latent
+                is_teacher_env = actions.get("is_teacher_env")
+                if is_teacher_env is None:
+                    raise ValueError(
+                        "action_mode='mixed' requires 'is_teacher_env' mask in actions dict"
+                    )
+
+                num_envs = tokenizer_meta_actions.shape[0]
+                atm = self.action_transform_module.actor_module
+
+                # Get proprioception for decoder (needed for student/direct_latent mode)
+                if "policy_atm" in atm_obs_dict:
+                    proprioception = atm_obs_dict["policy_atm"]
+                else:
+                    proprioception = atm_obs_dict["actor_obs"]
+                if proprioception.dim() == 2:
+                    proprioception = proprioception.unsqueeze(1)
+
+                teacher_mask = is_teacher_env
+                student_mask = ~is_teacher_env
+                num_teacher = teacher_mask.sum().item()
+                num_student = student_mask.sum().item()
+
+                # Initialize placeholders
+                teacher_indices = None
+                student_indices = None
+                teacher_body_actions = None
+                student_body_actions = None
+
+                # Process teacher envs (residual mode) if any
+                if num_teacher > 0:
+                    teacher_indices = teacher_mask.nonzero(as_tuple=True)[0]
+                    teacher_latent = tokenizer_meta_actions[teacher_indices]
+
+                    # Prepare teacher obs dict (subset of envs)
+                    # Only include keys that ATM actually needs: tokenizer and actor_obs
+                    teacher_atm_obs = {}
+                    atm_keys = ["tokenizer", "actor_obs"]
+                    for obs_key in atm_keys:
+                        if obs_key in atm_obs_dict:
+                            teacher_atm_obs[obs_key] = atm_obs_dict[obs_key][teacher_indices]
+
+                    scaled_residual = teacher_latent * self._latent_residual_scale
+                    teacher_body_actions = self.action_transform_module(
+                        teacher_atm_obs,
+                        latent_residual=scaled_residual,
+                        latent_residual_mode=self._latent_residual_mode,
+                    )
+
+                # Process student envs (direct_latent mode) if any
+                if num_student > 0:
+                    student_indices = student_mask.nonzero(as_tuple=True)[0]
+                    student_latent = tokenizer_meta_actions[student_indices]
+                    student_proprio = proprioception[student_indices]
+
+                    student_body_actions = self._decode_direct_latent_batch(
+                        student_latent, student_proprio, atm
+                    )
+
+                # Merge results - determine output shape from whichever mode ran
+                if teacher_body_actions is not None:
+                    out_seq_len = teacher_body_actions.shape[1]
+                    out_action_dim = teacher_body_actions.shape[2]
+                    dtype = teacher_body_actions.dtype
+                elif student_body_actions is not None:
+                    out_seq_len = student_body_actions.shape[1]
+                    out_action_dim = student_body_actions.shape[2]
+                    dtype = student_body_actions.dtype
+                else:
+                    raise RuntimeError("Mixed mode: no envs to process (both masks empty)")
+
+                body_actions = torch.zeros(
+                    num_envs, out_seq_len, out_action_dim, device=self.device, dtype=dtype
+                )
+
+                if teacher_body_actions is not None and teacher_indices is not None:
+                    body_actions[teacher_indices] = teacher_body_actions
+                if student_body_actions is not None and student_indices is not None:
+                    body_actions[student_indices] = student_body_actions
+
+            else:
+                raise ValueError(
+                    f"Unknown action_mode: {action_mode}. "
+                    f"Valid modes are 'direct_latent', 'residual', or 'mixed'."
+                )
+
+            # Store full latent (decoder input) for full_latent_rate_l2 reward
+            # Only needed for residual mode (teacher RL training); student modes
+            # (direct_latent, mixed) use L2 distillation loss, not RL rewards.
+            if action_mode == "residual":
+                self.env._prev_full_latent = self.env._full_latent.clone()  # noqa: SLF001
+                atm_module = self.action_transform_module.actor_module
+                if (
+                    hasattr(atm_module, "_last_full_latent_flat")
+                    and atm_module._last_full_latent_flat is not None  # noqa: SLF001
+                ):
+                    fl = atm_module._last_full_latent_flat  # noqa: SLF001
+                    if fl.dim() == 3:
+                        fl = fl[:, -1, :]  # (batch, latent_dim)
+                    self.env._full_latent = fl.to(self.env.device)  # noqa: SLF001
+
+            body_actions = body_actions[:, -1]  # Take last timestep
+
+            if (
+                self._body_joint_indices is not None
+                and self._hand_joint_indices is not None
+                and len(self._body_joint_indices) > 0
+                and len(self._hand_joint_indices) > 0
+            ):
+                num_envs = body_actions.shape[0]
+                env_actions = torch.zeros(
+                    num_envs, self._current_num_joints, device=self.device, dtype=body_actions.dtype
+                )
+                env_actions[:, self._body_joint_indices] = body_actions
+                env_actions[:, self._hand_joint_indices] = hand_actions
+            else:
+                env_actions = torch.cat([body_actions, hand_actions], dim=-1)
+
+        else:
+            env_actions = actions["actions"]
+
+        action_clip_value = self.config.get("action_clip_value", None)
+
+        if action_clip_value is not None and action_clip_value > 0:
+            env_actions = torch.clip(env_actions, -action_clip_value, action_clip_value)
+
+        # Lightweight action plot update (env 0, first N joints)
+        if self.turn_on_visualization:
+            try:
+                if self._action_fig is None:
+                    self._init_action_plot(env_actions.shape[-1])
+                if (self._step_counter % self._viz_every_n_steps) == 0:
+                    self._update_action_plot(env_actions)
+                self._step_counter += 1
+            except Exception:  # noqa: S110, BLE001
+                pass
+
+        obs_dict, rew, terminated, truncated, extras = self.env.step(env_actions)
+
+        # compute dones for compatibility with RSL-RL
+        dones = (terminated | truncated).to(dtype=torch.long)
+
+        # Zero out action/latent rate buffers for envs that just reset
+        # This prevents a false large rate penalty on the first step of a new episode
+        # Only applies when action_transform_module is used (buffers created in reset())
+        reset_mask = dones.bool()
+        if reset_mask.any() and hasattr(self.env, "_prev_meta_action"):
+            self.env._prev_meta_action[reset_mask] = 0.0  # noqa: SLF001
+            self.env._last_meta_action[reset_mask] = 0.0  # noqa: SLF001
+            self.env._prev_full_latent[reset_mask] = 0.0  # noqa: SLF001
+            self.env._full_latent[reset_mask] = 0.0  # noqa: SLF001
+
+        # Compute success_lift metric: check if object has no contact with table (lifted)
+        # Skip frames before first contact (from contact label data)
+        if (
+            hasattr(self.env, "scene")
+            and "object_to_table_contact_sensor" in self.env.scene.sensors
+        ):
+            from isaaclab.sensors import ContactSensor
+
+            sensor: ContactSensor = self.env.scene["object_to_table_contact_sensor"]
+            contact_force = sensor.data.force_matrix_w  # [num_envs, 1, 1, 3]
+            force_magnitude = torch.norm(contact_force, dim=-1).sum(dim=(-1, -2))  # [num_envs]
+
+            contact_force_threshold = self.config.get("lift_contact_force_threshold", 0.1)
+
+            if hasattr(self.env, "success_lift"):
+                self.env.success_lift = self.env.success_lift & (~dones.bool())
+
+            # Get first contact frame from motion command (contact label data)
+            is_before_contact = torch.ones(
+                self.env.num_envs, dtype=torch.bool, device=self.env.device
+            )
+            if self.motion_command is not None:
+                per_env_first_contact = getattr(self.motion_command, "_per_env_first_contact", None)
+                if per_env_first_contact is not None:
+                    current_time = (
+                        self.motion_command.motion_start_time_steps + self.motion_command.time_steps
+                    )
+                    is_before_contact = current_time < per_env_first_contact
+
+            # Object is currently lifted if there's no contact force from table
+            is_currently_lifted = force_magnitude <= contact_force_threshold
+
+            # Only update success_lift after first contact (cumulative OR)
+            # Before first contact, success_lift stays False regardless of contact
+            self.env.success_lift = torch.where(
+                is_before_contact,
+                self.env.success_lift,  # Keep current value (False after reset)
+                is_currently_lifted | self.env.success_lift,
+            )
+        else:
+            # If no contact sensor, set to False for all envs
+            self.env.success_lift = torch.zeros(
+                self.env.num_envs, dtype=torch.bool, device=self.env.device
+            )
+
+        extras["time_outs"] = truncated
+        extras["episode"] = {}
+        extras["to_log"] = {}
+        for k, v in extras["log"].items():
+            if isinstance(v, torch.Tensor):
+                extras["to_log"][k] = v
+            else:
+                extras["to_log"][k] = torch.tensor(v, dtype=torch.float)
+        if self._motion_lib is not None and self._motion_lib.use_adaptive_sampling:
+            extras["to_log"][
+                "adp_samp/num_episodes_min"
+            ] = self._motion_lib.adp_samp_num_episodes.min()
+            extras["to_log"][
+                "adp_samp/num_episodes_max"
+            ] = self._motion_lib.adp_samp_num_episodes.max()
+            extras["to_log"][
+                "adp_samp/num_episodes_mean"
+            ] = self._motion_lib.adp_samp_num_episodes.mean()
+            extras["to_log"][
+                "adp_samp/num_failures_min"
+            ] = self._motion_lib.adp_samp_num_failures.min()
+            extras["to_log"][
+                "adp_samp/num_failures_max"
+            ] = self._motion_lib.adp_samp_num_failures.max()
+            extras["to_log"][
+                "adp_samp/num_failures_mean"
+            ] = self._motion_lib.adp_samp_num_failures.mean()
+            extras["to_log"][
+                "adp_samp/failure_rate_min"
+            ] = self._motion_lib.adp_samp_failure_rate_raw.min()
+            extras["to_log"][
+                "adp_samp/failure_rate_max"
+            ] = self._motion_lib.adp_samp_failure_rate_raw.max()
+            extras["to_log"][
+                "adp_samp/failure_rate_mean"
+            ] = self._motion_lib.adp_samp_failure_rate_raw.mean()
+
+            if hasattr(self._motion_lib, "adp_sampling_active_prob"):
+                prob = self._motion_lib.adp_sampling_active_prob
+                uniform_prob = 1.0 / len(prob) if len(prob) > 0 else 1.0
+                extras["to_log"]["adp_samp/prob_max"] = prob.max()
+                extras["to_log"]["adp_samp/prob_min"] = prob.min()
+                extras["to_log"]["adp_samp/prob_mean"] = prob.mean()
+                extras["to_log"]["adp_samp/prob_max_over_uniform"] = prob.max() / uniform_prob
+                extras["to_log"]["adp_samp/effective_num_bins"] = 1.0 / (prob**2).sum()
+                # How many bins have prob > 10x uniform (significantly concentrated)
+                # Note: max allowed is 50x uniform, so 10x is 20% of the cap
+                extras["to_log"]["adp_samp/num_concentrated_bins"] = (
+                    (prob > 10 * uniform_prob).sum().float()
+                )
+
+            eps_mean = self._motion_lib.adp_samp_num_episodes.mean()
+            if eps_mean > 0:
+                extras["to_log"]["adp_samp/episodes_max_over_mean"] = (
+                    self._motion_lib.adp_samp_num_episodes.max() / eps_mean
+                )
+        new_obs = self.process_raw_obs(obs_dict, flatten_dict_obs=True)
+        # Store obs for action_transform_module when obs_dict is not provided in next step()
+        self._last_obs_dict = new_obs
+        self.extras = extras
+        # Store env_actions for callbacks (e.g., MultiLatentSaveCallback)
+        extras["env_actions"] = env_actions.detach().cpu()
+        return new_obs, rew, dones, extras
+
+    def get_env_data(self, key):
+        if key == "ref_body_pos_extend":
+            return self.motion_command.robot_body_pos_w
+        elif key == "rigid_body_pos_extend":
+            return self.motion_command.body_pos_w
+        else:
+            return self.env.get_env_data(key)
+
+    def render_results(self):
+        pass
+
+    def end_render_results(self):
+        if "render_envs" in self.env.recorder_manager._terms:  # noqa: SLF001
+            self.env.recorder_manager._terms["render_envs"].close_writers()  # noqa: SLF001
+        if "trajectory" in self.env.recorder_manager._terms:  # noqa: SLF001
+            self.env.recorder_manager._terms["trajectory"].close_writers()  # noqa: SLF001
+        if self._action_fig is not None:
+            try:
+                import matplotlib.pyplot as plt
+
+                plt.close(self._action_fig)
+            except Exception:  # noqa: S110, BLE001
+                pass
+            self._action_fig = None
+            self._action_lines = None
+            self._action_hist = None
+            self._hist_idx = 0
+            self._blit_background = None
+
+    def set_is_evaluating(self, is_evaluating: bool = True, global_rank=0, **_kwargs):
+        self.is_evaluating = is_evaluating
+        if self.motion_command is not None:
+            self.motion_command.set_is_evaluating(is_evaluating)
+        if self.force_command is not None:
+            self.force_command.set_is_evaluating(is_evaluating)
+        if is_evaluating:
+            self.begin_seq_motion_samples(global_rank)
+
+    def begin_seq_motion_samples(self, global_rank=0):
+        logger.info("Loading motions for evaluation")
+        self.start_idx = global_rank * self.num_envs
+        self._motion_lib.load_motions_for_evaluation(start_idx=self.start_idx)
+        self.reset_all(global_rank=global_rank)
+
+    def forward_motion_samples(self, global_rank=0, world_size=1):
+        old_start_idx = self.start_idx
+        self.start_idx += world_size * self.num_envs
+        logger.info(
+            f"Forward motions for evaluation from {old_start_idx} to {self.start_idx} - rank: {global_rank} - world size: {world_size}"  # noqa: E501
+        )
+        self._motion_lib.load_motions_for_evaluation(start_idx=self.start_idx)
+        self.reset_all(global_rank=global_rank)
+
+    def focusing_viewer(self):
+        if not self.viewer_focused:
+            # Focus on robot asset (tracking mode)
+            self.env.viewport_camera_controller.update_view_to_world()
+            self.env.viewport_camera_controller.update_view_to_asset_root("robot")
+            self.viewer_focused = True
+            logger.info("Focused on robot")
+        else:
+            # Switch to free camera mode centered on robot
+            self.env.viewport_camera_controller.viewer_origin = torch.zeros_like(
+                self.env.scene["robot"].data.root_pos_w[0]
+            )
+            self.env.viewport_camera_controller.cfg.origin_type = "world"
+            cam_eye = self.env.scene["robot"].data.root_pos_w[0].cpu().numpy() + np.array(
+                [2.0, 2.0, 1.5]
+            )
+            cam_target = self.env.scene["robot"].data.root_pos_w[0].cpu().numpy()
+            self.env.viewport_camera_controller._env.sim.set_camera_view(  # noqa: SLF001
+                eye=cam_eye, target=cam_target
+            )
+            self.viewer_focused = False
+            logger.info("Switched to free camera mode")  # noqa: RUF100, T201
+
+    def set_is_training(self, **_kwargs):
+        self.is_evaluating = False
+        self.resample_motion()
+
+    def resample_motion(self):
+        res = self._motion_lib.load_motions_for_training(
+            max_num_seqs=min(self.num_envs, self.motion_command.max_num_load_motions)
+        )
+        if res:
+            self.reset_all()
+        else:
+            logger.info("No new motions loaded, skipping reset")  # noqa: RUF100, T201
+
+    def sync_and_compute_adaptive_sampling(self, accelerator, sync_across_gpus=False):
+        if self._motion_lib is not None:
+            self._motion_lib.sync_and_compute_adaptive_sampling(
+                accelerator, sync_across_gpus=sync_across_gpus
+            )
+
+    def load_env_state_dict(self, state_dict):
+        if "motion_lib" in state_dict:
+            self._motion_lib.load_state_dict(state_dict["motion_lib"])
+            if self._motion_lib.use_adaptive_sampling:
+                self.resample_motion()
+
+    def get_env_state_dict(self):
+        state_dict = {
+            "motion_lib": self._motion_lib.get_state_dict(),
+        }
+        return state_dict
+
+    def reinit_dr(self, **_kwargs):
+        pass
+
+    def setup_keyboard(self):
+        try:
+            from isaaclab.devices.keyboard.se2_keyboard import Se2Keyboard, Se2KeyboardCfg
+
+            cfg = Se2KeyboardCfg()
+            self.keyboard_interface = Se2Keyboard(cfg)
+            self.keyboard_interface.add_callback("R", self.reset_all)
+            self.keyboard_interface.add_callback("T", self.forward_motion_samples)
+            self.keyboard_interface.add_callback("F", self.focusing_viewer)
+            self.keyboard_interface.add_callback("V", self.toggle_debug_vis)
+        except Exception as e:  # noqa: BLE001
+            logger.info(f"Error setting up keyboard: {e}")  # noqa: RUF100, T201
+
+    def toggle_debug_vis(self):
+        if self.motion_command is not None and hasattr(self.motion_command, "_set_debug_vis_impl"):
+            self._debug_vis_enabled = not getattr(self, "_debug_vis_enabled", True)
+            self.motion_command._set_debug_vis_impl(self._debug_vis_enabled)  # noqa: SLF001
+            logger.info(f"Debug visualization: {'ON' if self._debug_vis_enabled else 'OFF'}")  # noqa: RUF100, T201
+
+    # --- Action plotting helpers ---
+    def _init_action_plot(self, action_dim: int):
+        plot_dim = min(int(self._plot_action_dim), int(action_dim))
+        if plot_dim <= 0:
+            return
+        try:
+            import matplotlib.pyplot as plt
+        except Exception as e:  # noqa: BLE001, F841
+            return
+        plt.ion()
+        import math
+
+        cols = math.ceil(math.sqrt(plot_dim))
+        rows = math.ceil(plot_dim / cols)
+        fig, axes = plt.subplots(
+            rows, cols, sharex=True, sharey=True, figsize=(cols * 3.0, rows * 2.2)
+        )
+        axes = np.array(axes).reshape(-1)
+        self._plot_window = max(10, int(self._plot_window))
+        x_vals = np.arange(self._plot_window)
+        self._action_hist = np.zeros((self._plot_window, plot_dim), dtype=np.float32)
+        lines = []
+        y_min, y_max = -float(self._action_ylim), float(self._action_ylim)
+        for i in range(rows * cols):
+            ax_i = axes[i]
+            if i < plot_dim:
+                (ln,) = ax_i.plot(x_vals, self._action_hist[:, i], linewidth=1.0)
+                ln.set_animated(True)
+                lines.append(ln)
+                ax_i.set_ylim((y_min, y_max))
+                ax_i.set_xlim((0, self._plot_window - 1))
+                ax_i.set_title(f"J{i}", fontsize=8)
+                if i // cols == rows - 1:
+                    ax_i.set_xlabel("step", fontsize=8)
+                if i % cols == 0:
+                    ax_i.set_ylabel("act", fontsize=8)
+            else:
+                ax_i.axis("off")
+        fig.suptitle(f"Action time series (env 0) - first {plot_dim} joints", fontsize=10)
+        fig.tight_layout(rect=(0, 0.02, 1, 0.96))
+        self._action_fig = fig
+        self._action_lines = lines
+        self._hist_idx = -1
+        try:
+            fig.canvas.draw()
+            self._blit_background = fig.canvas.copy_from_bbox(fig.bbox)
+        except Exception:  # noqa: BLE001
+            self._blit_background = None
+            self._blit_supported = False
+        fig.canvas.flush_events()
+
+    def _update_action_plot(self, env_actions: torch.Tensor):
+        if self._action_lines is None or self._action_hist is None:
+            return
+        plot_dim = len(self._action_lines)
+        with torch.no_grad():
+            vals = env_actions[0, :plot_dim].detach().to("cpu").numpy().astype(np.float32)
+        vals = np.clip(vals, -self._action_ylim, self._action_ylim)
+        # advance circular buffer
+        self._hist_idx = (self._hist_idx + 1) % self._plot_window
+        self._action_hist[self._hist_idx, :plot_dim] = vals
+        # display in chronological order using a view with roll
+        disp = np.roll(self._action_hist, shift=-(self._hist_idx + 1), axis=0)
+        for i, ln in enumerate(self._action_lines):
+            ln.set_ydata(disp[:, i])
+        if self._action_fig is not None:
+            canvas = self._action_fig.canvas
+            # Periodically refresh background to handle resizes or overdraw
+            if self._blit_supported and (
+                self._blit_background is None
+                or (self._step_counter % max(1, self._blit_refresh_interval) == 0)
+            ):
+                try:
+                    self._action_fig.canvas.draw()
+                    self._blit_background = canvas.copy_from_bbox(self._action_fig.bbox)
+                except Exception:  # noqa: BLE001
+                    self._blit_background = None
+                    self._blit_supported = False
+            if self._blit_supported and self._blit_background is not None:
+                try:
+                    canvas.restore_region(self._blit_background)
+                    for ln in self._action_lines:
+                        ln.axes.draw_artist(ln)
+                    canvas.blit(self._action_fig.bbox)
+                    canvas.flush_events()
+                    return
+                except Exception:  # noqa: BLE001
+                    self._blit_supported = False
+            # Fallback full redraw
+            canvas.draw_idle()
+            canvas.flush_events()
+
+    @property
+    def motion_ids(self):
+        return self.motion_command.motion_ids
+
+    def setup_replay_grid(self, spacing=2.0, rows=None, cols=None):
+        """Setup a custom grid layout for environment origins during replay.
+
+        Args:
+            spacing: Distance between environments in meters (default: 2.0)
+            rows: Number of rows in the grid (default: auto-calculated)
+            cols: Number of columns in the grid (default: auto-calculated)
+
+        Returns:
+            Tensor of shape (num_envs, 3) with custom origins
+        """
+        import math
+
+        num_envs = self.num_envs
+
+        # Auto-calculate grid dimensions if not provided
+        if rows is None and cols is None:
+            # Try to make a square-ish grid
+            cols = int(math.ceil(math.sqrt(num_envs)))  # noqa: RUF046
+            rows = int(math.ceil(num_envs / cols))  # noqa: RUF046
+        elif rows is None:
+            rows = int(math.ceil(num_envs / cols))  # noqa: RUF046
+        elif cols is None:
+            cols = int(math.ceil(num_envs / rows))  # noqa: RUF046
+
+        logger.info(f"Setting up replay grid: {rows} rows x {cols} cols with {spacing}m spacing")
+
+        # Create grid origins
+        custom_origins = torch.zeros((num_envs, 3), device=self.device, dtype=torch.float32)
+
+        for i in range(num_envs):
+            row = i // cols
+            col = i % cols
+
+            # Center the grid around origin
+            x_offset = (col - (cols - 1) / 2.0) * spacing
+            y_offset = (row - (rows - 1) / 2.0) * spacing
+
+            custom_origins[i, 0] = x_offset
+            custom_origins[i, 1] = y_offset
+            custom_origins[i, 2] = 0.0  # Keep z at ground level
+
+        # Store custom origins
+        self._replay_custom_origins = custom_origins
+
+        logger.info(
+            f"Grid bounds: X=[{custom_origins[:, 0].min():.1f}, {custom_origins[:, 0].max():.1f}], "
+            f"Y=[{custom_origins[:, 1].min():.1f}, {custom_origins[:, 1].max():.1f}]"
+        )
+
+        return custom_origins
+
+    def run_replay(
+        self,
+        motion_id=None,
+        start_time_step=0,
+        speed=1.0,
+        loop=True,
+        enable_vis=False,
+        target_fps=50,
+        grid_spacing=2.0,
+        grid_rows=None,
+        grid_cols=None,
+        save_video_path=None,
+    ):
+        """Run a complete motion replay with automatic rendering loop for all environments.
+        This is the high-level convenience function that handles everything.
+
+        Args:
+            motion_id: The motion ID(s) to replay. Can be:
+                       - None: uses current motion_ids for all envs
+                       - int: single motion ID for all envs
+                       - list/tensor: motion ID per environment
+            start_time_step: Starting time step in the motion (default: 0)
+            speed: Playback speed multiplier (default: 1.0)
+            loop: Whether to loop the motion (default: True)
+            enable_vis: Whether to enable debug visualization (default: False)
+            target_fps: Target frames per second for rendering (default: 50)
+            grid_spacing: Distance between environments in meters (default: 2.0)
+            grid_rows: Number of rows in grid layout (default: auto)
+            grid_cols: Number of columns in grid layout (default: auto)
+            save_video_path: Path to save video file (default: None, no video saved)
+                             Requires render_results=True in config to enable eval_camera.
+
+        Returns:
+            List of dictionaries with motion metadata for each environment
+        """  # noqa: D205
+        import os
+        import time
+
+        # Setup custom grid origins
+        self.setup_replay_grid(spacing=grid_spacing, rows=grid_rows, cols=grid_cols)
+
+        # Initialize replay
+        info = self.setup_replay_motion(
+            motion_id=motion_id,
+            start_time_step=start_time_step,
+            speed=speed,
+            loop=loop,
+            enable_vis=enable_vis,
+        )
+
+        if info is None:
+            return None
+
+        # Initialize video writer if save_video_path is specified
+        video_writer = None
+        if save_video_path is not None:
+            if "eval_camera" not in self.env.scene.sensors:
+                logger.info("WARNING: save_video_path specified but eval_camera not available.")
+                logger.info("         Set ++manager_env.config.render_results=True to enable it.")
+            else:
+                import imageio
+
+                os.makedirs(os.path.dirname(os.path.abspath(save_video_path)), exist_ok=True)
+                video_writer = imageio.get_writer(
+                    save_video_path,
+                    fps=target_fps,
+                    codec="libx264",
+                    quality=5,
+                    pixelformat="yuv420p",
+                )
+                logger.info(f"[Video] Recording to: {save_video_path}")
+
+        # Run the rendering loop
+        frame_time = 1.0 / target_fps
+        last_time = time.time()
+        frame_count = 0
+
+        try:
+            while self.step_replay():
+                # Position camera BEFORE rendering (if recording)
+                if video_writer is not None and "eval_camera" in self.env.scene.sensors:
+                    root_pos = self.motion_command.robot.data.body_pos_w[:, 0]
+                    camera_offset = self.config.get("eval_camera_offset", [-2, -2, 1])
+                    eye = root_pos + torch.tensor(camera_offset, device=self.device)
+                    self.env.scene["eval_camera"].set_world_poses_from_view(eye, root_pos)
+
+                # Render the scene (triggers camera render in headless mode)
+                if hasattr(self.env, "sim"):
+                    try:
+                        self.env.sim.render()
+                    except Exception as e:  # noqa: BLE001
+                        logger.info(f"Render error: {e}")
+
+                # Capture video frame if recording
+                if video_writer is not None and "eval_camera" in self.env.scene.sensors:
+                    # Update camera to refresh its data after render
+                    self.env.scene["eval_camera"].update(dt=0.0)
+
+                    # Grab frame from camera
+                    rgb_frame = self.env.scene["eval_camera"].data.output["rgb"]
+                    frame = rgb_frame[0].cpu().numpy().astype(np.uint8)
+                    video_writer.append_data(frame)
+                    frame_count += 1
+
+                # Frame rate limiting (only when not headless or not saving video)
+                if not self.config.get("headless", False) or video_writer is None:
+                    current_time = time.time()
+                    elapsed = current_time - last_time
+                    sleep_time = max(0, (frame_time / speed) - elapsed)
+
+                    if sleep_time > 0:
+                        time.sleep(sleep_time)
+
+                    last_time = time.time()
+
+        except KeyboardInterrupt:
+            logger.info("\nReplay interrupted by user")
+
+        # Close video writer
+        if video_writer is not None:
+            video_writer.close()
+            logger.info(f"[Video] Saved {frame_count} frames to: {save_video_path}")
+
+        logger.info("\nReplay complete!")
+        return info
+
+    def setup_replay_motion(
+        self, motion_id=None, start_time_step=0, speed=1.0, loop=True, enable_vis=False
+    ):
+        """Initialize motion replay for all environments (low-level function for manual control).
+        Use run_replay() instead if you want automatic rendering loop.
+
+        After calling this, you need to manually call step_replay() in a loop
+        and handle rendering yourself. For automatic rendering, use run_replay().
+
+        Args:
+            motion_id: The motion ID(s) to replay. Can be:
+                       - None: uses current motion_ids for all envs
+                       - int: single motion ID for all envs
+                       - list/tensor: motion ID per environment
+            start_time_step: Starting time step in the motion (default: 0)
+            speed: Playback speed multiplier (default: 1.0)
+            loop: Whether to loop the motion (default: True)
+            enable_vis: Whether to enable debug visualization (default: False)
+
+        Returns:
+            List of dictionaries with motion metadata for each environment
+        """  # noqa: D205
+        if self._motion_lib is None:
+            logger.info("No motion library available for replay")
+            return None
+
+        # Replay all environments
+        num_replay_envs = self.num_envs
+        replay_env_ids = torch.arange(num_replay_envs, device=self.device, dtype=torch.long)
+
+        # Get motion IDs
+        if motion_id is None:
+            motion_ids = self.motion_command.motion_ids
+        elif isinstance(motion_id, int):
+            motion_ids = torch.full(
+                (num_replay_envs,), motion_id, device=self.device, dtype=torch.long
+            )
+        elif isinstance(motion_id, list | tuple):
+            if len(motion_id) != num_replay_envs:
+                logger.info(
+                    f"Error: motion_id list length ({len(motion_id)}) doesn't match number of environments ({num_replay_envs})"  # noqa: E501
+                )
+                return None
+            motion_ids = torch.tensor(motion_id, device=self.device, dtype=torch.long)
+        elif isinstance(motion_id, torch.Tensor):
+            if motion_id.shape[0] != num_replay_envs:
+                logger.info(
+                    f"Error: motion_id tensor size ({motion_id.shape[0]}) doesn't match number of environments ({num_replay_envs})"  # noqa: E501
+                )
+                return None
+            motion_ids = motion_id.to(self.device)
+        else:
+            logger.info(f"Error: Invalid motion_id type: {type(motion_id)}")
+            return None
+
+        # Get motion info for display
+        num_steps_per_env = self._motion_lib.get_motion_num_steps(motion_ids)
+        max_num_steps = num_steps_per_env.max().item()
+
+        unique_motions = torch.unique(motion_ids)
+        logger.info(f"\n{'='*60}")
+        logger.info(f"Batch Replaying {num_replay_envs} Environments")
+        logger.info(f"Unique motion IDs: {unique_motions.tolist()}")
+        logger.info(f"Max frames: {max_num_steps}")
+        logger.info(f"Max duration: {max_num_steps / 50.0:.2f}s (@ 50 FPS)")
+        logger.info(f"Speed: {speed}x")
+        logger.info(f"{'='*60}\n")
+
+        # Enable visualization if requested
+        if enable_vis and hasattr(self.motion_command, "_set_debug_vis_impl"):
+            self.motion_command._set_debug_vis_impl(True)  # noqa: SLF001
+
+        # Initialize contact point visualization if available and enabled
+        self._replay_contact_visualizer = None
+        self._replay_vis_enabled = enable_vis
+        if enable_vis and VISUALIZATION_AVAILABLE:
+            self._setup_contact_center_visualizer()
+
+        # Replay state
+        self._replay_active = True
+        self._replay_paused = False
+        self._replay_reverse = False
+        self._replay_time_steps = torch.full(
+            (num_replay_envs,), start_time_step, device=self.device, dtype=torch.long
+        )
+        self._replay_motion_ids = motion_ids
+        self._replay_env_ids = replay_env_ids
+        self._replay_speed = speed
+        self._replay_loop = loop
+        self._replay_num_steps_per_env = num_steps_per_env
+        self._replay_max_num_steps = max_num_steps
+
+        # Pre-load table metadata from pkl files for ALL motions (per-env)
+        # This supports multi-motion replay where each env can have different table positions
+        self._replay_table_pos = None  # Shape: (num_envs, 3)
+        self._replay_table_quat = None  # Shape: (num_envs, 4)
+        if hasattr(self.env, "scene") and "table" in self.env.scene.rigid_objects:
+            try:
+                import os
+
+                import joblib
+
+                # Derive meta directory from motion_file path in config
+                motion_lib_cfg = getattr(self.motion_command.cfg, "motion_lib_cfg", None)
+                motion_file = motion_lib_cfg.get("motion_file", "") if motion_lib_cfg else ""
+
+                if motion_file and "/robot" in motion_file:
+                    if os.path.isdir(motion_file):
+                        meta_dir = motion_file.replace("/robot", "/meta")
+                    else:
+                        meta_dir = os.path.dirname(motion_file).replace("/robot", "/meta")
+                else:
+                    meta_dir = "data/motion_lib_grab/meta"
+
+                # Pre-load motion file data if it's a single file (for table fallback)
+                motion_file_data = None
+                if motion_file and os.path.isfile(motion_file) and not os.path.isdir(meta_dir):
+                    motion_file_data = joblib.load(motion_file)
+
+                # Load table metadata for each environment's assigned motion
+                table_pos_list = []
+                table_quat_list = []
+                table_scale_list = []  # Also load scales for offset calculation
+                loaded_count = 0
+                loaded_from_motion_file = 0
+
+                for env_idx in range(num_replay_envs):
+                    motion_idx = motion_ids[env_idx].item()
+                    motion_key = self._motion_lib.curr_motion_keys[motion_idx]
+                    meta_file = os.path.join(meta_dir, f"{motion_key}.pkl")
+
+                    if os.path.exists(meta_file):
+                        meta = joblib.load(meta_file)
+                        pos = torch.tensor(
+                            meta.get("table_pos", [0.0, 0.0, 0.8]), device=self.device
+                        ).float()
+                        quat = torch.tensor(
+                            meta.get("table_quat", [1.0, 0.0, 0.0, 0.0]), device=self.device
+                        ).float()
+                        scale = torch.tensor(
+                            meta.get("table_scale", [1.0, 1.0, 1.0]), device=self.device
+                        ).float()
+                        table_pos_list.append(pos)
+                        table_quat_list.append(quat)
+                        table_scale_list.append(scale)
+                        loaded_count += 1
+                    elif motion_file_data is not None and motion_key in motion_file_data:
+                        # Fallback: try to get table data from motion file directly
+                        motion_data = motion_file_data[motion_key]
+                        if "table_pos" in motion_data and "table_quat" in motion_data:
+                            pos = torch.tensor(motion_data["table_pos"], device=self.device).float()
+                            quat = torch.tensor(
+                                motion_data["table_quat"], device=self.device
+                            ).float()
+                            scale = torch.tensor(
+                                motion_data.get("table_scale", [1.0, 1.0, 1.0]), device=self.device
+                            ).float()
+                            table_pos_list.append(pos)
+                            table_quat_list.append(quat)
+                            table_scale_list.append(scale)
+                            loaded_from_motion_file += 1
+                        else:
+                            # No table data in motion file, use default
+                            table_pos_list.append(
+                                torch.tensor([0.0, 0.0, 0.8], device=self.device).float()
+                            )
+                            table_quat_list.append(
+                                torch.tensor([1.0, 0.0, 0.0, 0.0], device=self.device).float()
+                            )
+                            table_scale_list.append(
+                                torch.tensor([1.0, 1.0, 1.0], device=self.device).float()
+                            )
+                    else:
+                        # Fallback: use default table position
+                        table_pos_list.append(
+                            torch.tensor([0.0, 0.0, 0.8], device=self.device).float()
+                        )
+                        table_quat_list.append(
+                            torch.tensor([1.0, 0.0, 0.0, 0.0], device=self.device).float()
+                        )
+                        table_scale_list.append(
+                            torch.tensor([1.0, 1.0, 1.0], device=self.device).float()
+                        )
+
+                self._replay_table_pos = torch.stack(table_pos_list, dim=0)  # (num_envs, 3)
+                self._replay_table_quat = torch.stack(table_quat_list, dim=0)  # (num_envs, 4)
+                self._replay_table_scale = torch.stack(table_scale_list, dim=0)  # (num_envs, 3)
+
+                # Apply fixed table_offset if configured
+                table_offset = self.config.get("table_offset", None)
+                if table_offset is not None:
+                    offset_tensor = torch.tensor(
+                        table_offset, device=self.device, dtype=self._replay_table_pos.dtype
+                    )
+                    self._replay_table_pos = self._replay_table_pos + offset_tensor
+
+                # Apply maximal X offset based on object starting position
+                # This computes the maximum valid offset range such that object stays on table
+                x_offset_maximal = self.config.get("replay_table_x_offset_maximal", False)
+                x_offset_margin = self.config.get(
+                    "replay_table_x_offset_margin", 0.05
+                )  # 5cm safety margin
+
+                if x_offset_maximal:
+                    # Get base table width from config (required for maximal offset mode)
+                    table_size_cfg = self.config.get("table_size")
+                    if table_size_cfg is None:
+                        raise ValueError(
+                            "replay_table_x_offset_maximal requires table_size to be set in config"
+                        )
+                    base_table_width = table_size_cfg[0]
+
+                    # Get object starting position (step 0) for each env
+                    object_start_pos = self._motion_lib.get_object_root_pos(
+                        motion_ids,
+                        torch.zeros(num_replay_envs, device=self.device, dtype=torch.long),
+                    )[
+                        :, 0
+                    ]  # Shape: (num_envs, 3) -> take first object if multiple
+                    object_start_x = object_start_pos[:, 0]  # (num_envs,)
+
+                    # Get table position and scaled width per env
+                    table_x = self._replay_table_pos[:, 0]  # (num_envs,)
+                    scale_x = self._replay_table_scale[:, 0]  # (num_envs,)
+                    actual_width = base_table_width * scale_x  # (num_envs,)
+                    half_width = actual_width / 2.0
+
+                    # Compute valid offset range per env:
+                    # Object must stay on table after offset dx
+                    # dx_min = object_x - table_x - half_width + margin
+                    # dx_max = object_x - table_x + half_width - margin
+                    relative_obj_x = object_start_x - table_x
+                    dx_min = relative_obj_x - half_width + x_offset_margin
+                    dx_max = relative_obj_x + half_width - x_offset_margin
+
+                    # Sample X offset uniformly within valid range per env
+                    rand_vals = torch.rand(num_replay_envs, device=self.device)
+                    x_offsets = dx_min + rand_vals * (dx_max - dx_min)
+
+                    self._replay_table_pos[:, 0] += x_offsets
+
+                    logger.info("Replay table X offset (maximal mode):")
+                    logger.info(
+                        f"  Base table width: {base_table_width:.3f}m, margin: {x_offset_margin:.3f}m"
+                    )
+                    logger.info(f"  dx_min range: [{dx_min.min():.3f}, {dx_min.max():.3f}]")
+                    logger.info(f"  dx_max range: [{dx_max.min():.3f}, {dx_max.max():.3f}]")
+                    logger.info(f"  Applied offsets: [{x_offsets.min():.3f}, {x_offsets.max():.3f}]")
+
+                # Apply fixed XY offset range (legacy mode, used if maximal mode is disabled)
+                # Format: [x_min, x_max, y_min, y_max] - per-env randomized
+                xy_offset_range = self.config.get("replay_table_xy_offset_range", None)
+                if xy_offset_range is not None and not x_offset_maximal:
+                    x_min, x_max, y_min, y_max = xy_offset_range
+
+                    # Sample X offset per environment
+                    x_offsets = torch.empty(num_replay_envs, device=self.device).uniform_(
+                        x_min, x_max
+                    )
+                    self._replay_table_pos[:, 0] += x_offsets
+
+                    # Sample Y offset per environment (typically y_min = y_max for fixed offset)
+                    y_offsets = torch.empty(num_replay_envs, device=self.device).uniform_(
+                        y_min, y_max
+                    )
+                    self._replay_table_pos[:, 1] += y_offsets
+
+                    logger.info(
+                        f"Replay table XY offset: X=[{x_min:.3f}, {x_max:.3f}], Y=[{y_min:.3f}, {y_max:.3f}]"
+                    )
+
+                if loaded_count > 0:
+                    logger.info(
+                        f"Loaded table metadata for {loaded_count}/{num_replay_envs} environments from {meta_dir}"
+                    )
+                elif loaded_from_motion_file > 0:
+                    logger.info(
+                        f"Loaded table metadata for {loaded_from_motion_file}/{num_replay_envs} environments from motion file"  # noqa: E501
+                    )
+
+            except Exception as e:  # noqa: BLE001
+                logger.info(f"Warning: Could not load table metadata: {e}")
+
+        # Setup keyboard controls for replay
+        if not self.config.headless:
+            self._setup_replay_keyboard()
+
+        logger.info("Replay Controls:")
+        logger.info("  G: Pause/Resume")
+        logger.info("  B: Toggle Reverse Play")
+        logger.info("  LEFT/RIGHT: Step backward/forward (when paused)")
+        logger.info("  R: Restart from beginning")
+        logger.info("  ESC: Exit replay")
+        logger.info("  +/-: Increase/Decrease speed")
+        logger.info("")
+
+        # Return metadata
+        results = []
+        for i, (mid, nsteps) in enumerate(
+            zip(motion_ids.tolist(), num_steps_per_env.tolist(), strict=False)
+        ):
+            # Get motion key from motion library (curr_motion_keys is the active list)
+            if hasattr(self._motion_lib, "curr_motion_keys") and mid < len(
+                self._motion_lib.curr_motion_keys
+            ):
+                motion_key = self._motion_lib.curr_motion_keys[mid]
+            elif hasattr(self._motion_lib, "motion_keys") and mid < len(
+                self._motion_lib.motion_keys
+            ):
+                motion_key = self._motion_lib.motion_keys[mid]
+            else:
+                motion_key = f"motion_{mid}"
+            results.append(
+                {
+                    "env_id": i,
+                    "motion_id": mid,
+                    "motion_key": motion_key,
+                    "num_steps": nsteps,
+                    "duration": nsteps / 50.0,
+                }
+            )
+        return results
+
+    def _setup_replay_keyboard(self):
+        """Setup keyboard controls for replay mode"""  # noqa: D415
+        try:
+            if not hasattr(self, "keyboard_interface"):
+                from isaaclab.devices.keyboard.se2_keyboard import Se2Keyboard, Se2KeyboardCfg
+
+                cfg = Se2KeyboardCfg()
+                self.keyboard_interface = Se2Keyboard(cfg)
+
+            # Add replay-specific callbacks
+            self.keyboard_interface.add_callback("G", self._toggle_replay_pause)
+            self.keyboard_interface.add_callback("LEFT", self._replay_step_backward)
+            self.keyboard_interface.add_callback("RIGHT", self._replay_step_forward)
+            self.keyboard_interface.add_callback("R", self._restart_replay)
+            self.keyboard_interface.add_callback("ESCAPE", self._exit_replay)
+            self.keyboard_interface.add_callback("EQUAL", self._increase_replay_speed)
+            self.keyboard_interface.add_callback("MINUS", self._decrease_replay_speed)
+            self.keyboard_interface.add_callback("B", self._toggle_reverse_play)
+        except Exception as e:  # noqa: BLE001
+            logger.info(f"Could not setup replay keyboard controls: {e}")
+
+    def _setup_contact_center_visualizer(self):
+        """Setup visualization markers for per-hand contact centers during replay."""
+        if not VISUALIZATION_AVAILABLE:
+            return
+
+        has_left = hasattr(self._motion_lib, "_motion_object_contact_center_left")
+        has_right = hasattr(self._motion_lib, "_motion_object_contact_center_right")
+        if not has_left and not has_right:
+            logger.info("No contact center available in motion library for visualization")
+            return
+
+        try:
+            self._replay_contact_visualizers = {}
+
+            if has_left:
+                left_cfg = VisualizationMarkersCfg(
+                    prim_path="/Visuals/Replay/contact_center_left",
+                    markers={
+                        "contact": sim_utils.SphereCfg(
+                            radius=0.03,
+                            visual_material=sim_utils.PreviewSurfaceCfg(
+                                diffuse_color=(0.0, 0.0, 1.0)  # Blue for left hand
+                            ),
+                        ),
+                    },
+                )
+                self._replay_contact_visualizers["left_hand"] = VisualizationMarkers(left_cfg)
+                self._replay_contact_visualizers["left_hand"].set_visibility(True)
+
+            if has_right:
+                right_cfg = VisualizationMarkersCfg(
+                    prim_path="/Visuals/Replay/contact_center_right",
+                    markers={
+                        "contact": sim_utils.SphereCfg(
+                            radius=0.03,
+                            visual_material=sim_utils.PreviewSurfaceCfg(
+                                diffuse_color=(0.0, 1.0, 1.0)  # Green for right hand
+                            ),
+                        ),
+                    },
+                )
+                self._replay_contact_visualizers["right_hand"] = VisualizationMarkers(right_cfg)
+                self._replay_contact_visualizers["right_hand"].set_visibility(True)
+
+            hands = list(self._replay_contact_visualizers.keys())
+            logger.info(f"Contact center visualizer initialized for: {', '.join(hands)}")
+
+            # Keep _replay_contact_visualizer as a truthy check for the update loop
+            self._replay_contact_visualizer = True
+
+        except Exception as e:  # noqa: BLE001
+            logger.info(f"Could not setup contact center visualizer: {e}")
+            self._replay_contact_visualizer = None
+
+    def _update_contact_center_visualization(self):
+        """Update per-hand contact center visualization during replay."""
+        if not getattr(self, "_replay_contact_visualizers", None):
+            return
+
+        from gear_sonic.isaac_utils.rotations import quat_rotate
+
+        # Get object pose (shared by both hands)
+        object_root_pos = self._motion_lib.get_object_root_pos(
+            self._replay_motion_ids, self._replay_time_steps
+        )[
+            :, 0, :
+        ]  # (num_envs, 3)
+        object_root_quat = self._motion_lib.get_object_root_quat(
+            self._replay_motion_ids, self._replay_time_steps
+        )[
+            :, 0, :
+        ]  # (num_envs, 4)
+
+        # Environment origin offsets
+        if hasattr(self, "_replay_custom_origins"):
+            env_origins = self._replay_custom_origins[self._replay_env_ids]
+        else:
+            env_origins = self.env.scene.env_origins[self._replay_env_ids]
+
+        # Hidden position for markers when no contact
+        hidden_pos = torch.tensor([[0.0, 0.0, -1000.0]], device=self.device)
+
+        for hand, visualizer in self._replay_contact_visualizers.items():
+            contact_center = self._motion_lib.get_object_contact_center(
+                self._replay_motion_ids, self._replay_time_steps, hand=hand
+            )
+            if contact_center is None:
+                visualizer.visualize(translations=hidden_pos)
+                continue
+
+            # Move markers far away for envs with no contact (zero contact center)
+            valid_mask = torch.norm(contact_center, dim=-1) > 1e-6
+            rotated = quat_rotate(object_root_quat, contact_center, w_last=False)
+            world_center = rotated + object_root_pos + env_origins
+            world_center[~valid_mask] = hidden_pos
+            visualizer.visualize(translations=world_center)
+
+    def _toggle_replay_pause(self):
+        """Toggle pause state for replay"""  # noqa: D415
+        if hasattr(self, "_replay_active") and self._replay_active:
+            self._replay_paused = not self._replay_paused
+            status = "PAUSED" if self._replay_paused else "PLAYING"
+            max_frame = self._replay_time_steps.max().item()
+            logger.info(f"Replay {status} (max frame {max_frame}/{self._replay_max_num_steps})")
+
+    def _replay_step_backward(self):
+        """Step backward one frame when paused"""  # noqa: D415
+        if hasattr(self, "_replay_active") and self._replay_active and self._replay_paused:
+            self._replay_time_steps = torch.clamp(self._replay_time_steps - 1, min=0)
+            self._update_replay_frame()
+            max_frame = self._replay_time_steps.max().item()
+            logger.info(f"Frame {max_frame}/{self._replay_max_num_steps}")
+            self.env.sim.render()
+
+    def _replay_step_forward(self):
+        """Step forward one frame when paused"""  # noqa: D415
+        if hasattr(self, "_replay_active") and self._replay_active and self._replay_paused:
+            self._replay_time_steps = torch.minimum(
+                self._replay_time_steps + 1, self._replay_num_steps_per_env - 1
+            )
+            self._update_replay_frame()
+            max_frame = self._replay_time_steps.max().item()
+            logger.info(f"Frame {max_frame}/{self._replay_max_num_steps}")
+            self.env.sim.render()
+
+    def _restart_replay(self):
+        """Restart replay from beginning"""  # noqa: D415
+        if hasattr(self, "_replay_active") and self._replay_active:
+            self._replay_time_steps.fill_(0)
+            self._replay_paused = False
+            self._update_replay_frame()
+            logger.info("Replay restarted from beginning")
+            self.env.sim.render()
+
+    def _exit_replay(self):
+        """Exit replay mode"""  # noqa: D415
+        if hasattr(self, "_replay_active"):
+            self._replay_active = False
+            # Hide contact center visualizers
+            for vis in getattr(self, "_replay_contact_visualizers", {}).values():
+                vis.set_visibility(False)
+            logger.info("Exiting replay mode")
+            self.env.sim.render()
+
+    def _increase_replay_speed(self):
+        """Increase replay speed"""  # noqa: D415
+        if hasattr(self, "_replay_active") and self._replay_active:
+            self._replay_speed = min(5.0, self._replay_speed * 1.25)
+            logger.info(f"Replay speed: {self._replay_speed:.2f}x")
+
+    def _decrease_replay_speed(self):
+        """Decrease replay speed"""  # noqa: D415
+        if hasattr(self, "_replay_active") and self._replay_active:
+            self._replay_speed = max(0.1, self._replay_speed * 0.8)
+            logger.info(f"Replay speed: {self._replay_speed:.2f}x")
+
+    def _toggle_reverse_play(self):
+        """Toggle reverse playback mode"""  # noqa: D415
+        if hasattr(self, "_replay_active") and self._replay_active:
+            self._replay_reverse = not self._replay_reverse
+            direction = "REVERSE" if self._replay_reverse else "FORWARD"
+            logger.info(f"Replay direction: {direction}")
+
+    def _update_replay_frame(self):
+        """Update all robot states to match the current replay frames"""  # noqa: D415
+        if not hasattr(self, "_replay_active") or not self._replay_active:
+            return
+
+        # Get motion data at current time steps for all environments
+        root_pos = self._motion_lib.get_root_pos_w(self._replay_motion_ids, self._replay_time_steps)
+        root_quat = self._motion_lib.get_root_quat_w(
+            self._replay_motion_ids, self._replay_time_steps
+        )
+        root_lin_vel = self._motion_lib.get_root_lin_vel_w(
+            self._replay_motion_ids, self._replay_time_steps
+        )
+        root_ang_vel = self._motion_lib.get_root_ang_vel_w(
+            self._replay_motion_ids, self._replay_time_steps
+        )
+        motion_lib_joint_pos = self._motion_lib.get_dof_pos(
+            self._replay_motion_ids, self._replay_time_steps
+        )
+        motion_lib_joint_vel = self._motion_lib.get_dof_vel(
+            self._replay_motion_ids, self._replay_time_steps
+        )
+
+        # Handle DOF mismatch between motion library (e.g., 29 DOF) and robot (e.g., 43 DOF)
+        robot_num_joints = self.motion_command.robot.num_joints
+        motion_lib_num_dof = motion_lib_joint_pos.shape[-1]
+
+        if robot_num_joints > motion_lib_num_dof and self._body_joint_indices is not None:
+            # Robot has more DOFs than motion lib (e.g., 43 DOF robot with 29 DOF motion data)
+            # Use body joint indices for proper mapping
+            num_envs = motion_lib_joint_pos.shape[0]
+
+            # Create full joint tensors with zeros for all DOFs
+            joint_pos = torch.zeros(
+                num_envs, robot_num_joints, device=self.device, dtype=motion_lib_joint_pos.dtype
+            )
+            joint_vel = torch.zeros(
+                num_envs, robot_num_joints, device=self.device, dtype=motion_lib_joint_vel.dtype
+            )
+
+            # Map motion lib data to body joint indices (using G1_ISAACLab_ORDER mapping)
+            joint_pos[:, self._body_joint_indices] = motion_lib_joint_pos
+            joint_vel[:, self._body_joint_indices] = motion_lib_joint_vel
+
+            # Use hand DOFs from motion lib if available, otherwise default to zero
+            hand_dof_pos = self._motion_lib.get_hand_dof_pos(
+                self._replay_motion_ids, self._replay_time_steps
+            )
+            if hand_dof_pos is not None:
+                # Hand DOFs are the last N joints (in Isaac order, not G1_HAND_JOINTS order)
+                num_hand_dof = hand_dof_pos.shape[-1]
+                joint_pos[:, -num_hand_dof:] = hand_dof_pos
+        else:
+            joint_pos = motion_lib_joint_pos
+            joint_vel = motion_lib_joint_vel
+
+        # Add environment origin offsets (use custom grid origins if set)
+        if hasattr(self, "_replay_custom_origins"):
+            root_pos = root_pos + self._replay_custom_origins[self._replay_env_ids]
+        else:
+            root_pos = root_pos + self.env.scene.env_origins[self._replay_env_ids]
+
+        # Write state to simulation for all environments
+        self.motion_command.robot.write_joint_state_to_sim(
+            joint_pos, joint_vel, env_ids=self._replay_env_ids
+        )
+        self.motion_command.robot.write_root_state_to_sim(
+            torch.cat([root_pos, root_quat, root_lin_vel, root_ang_vel], dim=-1),
+            env_ids=self._replay_env_ids,
+        )
+
+        # Get object motion data from motion library
+        if hasattr(self._motion_lib, "_motion_object_root_pos") and hasattr(
+            self._motion_lib, "_motion_object_root_quat"
+        ):
+            object_root_pos = self._motion_lib.get_object_root_pos(
+                self._replay_motion_ids, self._replay_time_steps
+            )
+            object_root_quat = self._motion_lib.get_object_root_quat(
+                self._replay_motion_ids, self._replay_time_steps
+            )
+
+            # Add environment origin offsets to object position
+            if hasattr(self, "_replay_custom_origins"):
+                object_root_pos = object_root_pos + self._replay_custom_origins[
+                    self._replay_env_ids
+                ].unsqueeze(1)
+            else:
+                object_root_pos = object_root_pos + self.env.scene.env_origins[
+                    self._replay_env_ids
+                ].unsqueeze(1)
+
+            # Write object state to simulation (handle multiple objects)
+            # Shape: object_root_pos is (num_envs, max_num_objects, 3)
+            # Shape: object_root_quat is (num_envs, max_num_objects, 4)
+            for obj_idx in range(object_root_pos.shape[1]):
+                obj_pos = object_root_pos[:, obj_idx, :]
+                obj_quat = object_root_quat[:, obj_idx, :]
+                object_root_pose = torch.cat([obj_pos, obj_quat], dim=-1)
+
+                self.env.scene["object"].write_root_pose_to_sim(
+                    object_root_pose, env_ids=self._replay_env_ids
+                )
+
+        if hasattr(self.env, "scene") and "table" in self.env.scene.rigid_objects:
+            # Use pre-loaded per-motion table metadata (loaded during setup_replay_motion)
+            # _replay_table_pos and _replay_table_quat have shape (num_envs, 3) and (num_envs, 4)
+            if self._replay_table_pos is not None:
+                # Use per-env table positions directly (already the right shape)
+                table_pos = self._replay_table_pos.clone()
+                table_quat = self._replay_table_quat.clone()
+            else:
+                # Fallback: derive from object position with hardcoded offset
+                table_pos = self._motion_lib.get_object_root_pos(
+                    self._replay_motion_ids, torch.zeros_like(self._replay_time_steps)
+                )[:, 0].clone()
+                table_pos[:, 2] = 0.76  # Table height
+                table_pos[:, 1] -= 0.15
+                table_quat = torch.tensor([[1.0, 0.0, 0.0, 0.0]], device=self.device).repeat(
+                    len(self._replay_env_ids), 1
+                )
+
+                # Apply table_offset if configured (for fallback path)
+                table_offset = self.config.get("table_offset", None)
+                if table_offset is not None:
+                    offset_tensor = torch.tensor(
+                        table_offset, device=self.device, dtype=table_pos.dtype
+                    )
+                    table_pos = table_pos + offset_tensor
+
+            # Add environment origin offsets (same as robot and object)
+            if hasattr(self, "_replay_custom_origins"):
+                table_pos = table_pos + self._replay_custom_origins[self._replay_env_ids]
+            else:
+                table_pos = table_pos + self.env.scene.env_origins[self._replay_env_ids]
+
+            table_root_pose = torch.cat([table_pos, table_quat], dim=-1)
+            self.env.scene["table"].write_root_pose_to_sim(
+                table_root_pose, env_ids=self._replay_env_ids
+            )
+
+        # Visualize contact points if enabled
+        if (
+            hasattr(self, "_replay_contact_visualizer")
+            and self._replay_contact_visualizer is not None
+            and hasattr(self, "_replay_vis_enabled")
+            and self._replay_vis_enabled
+        ):
+            self._update_contact_center_visualization()
+
+        self.env.sim.forward()
+
+    def step_replay(self):
+        """Step the replay forward for all environments. Call this in a loop to animate the motions.
+        Returns False when replay is complete.
+        """  # noqa: D205
+        if not hasattr(self, "_replay_active") or not self._replay_active:
+            return False
+
+        # Update frames based on speed and direction
+        if not self._replay_paused:
+            if self._replay_reverse:
+                # Playing in reverse
+                self._replay_time_steps -= int(self._replay_speed)
+
+                # Handle beginning of motions per environment
+                at_start = self._replay_time_steps < 0
+                if at_start.any():
+                    if self._replay_loop:
+                        # Loop to end when reaching start
+                        self._replay_time_steps[at_start] = (
+                            self._replay_num_steps_per_env[at_start] - 1
+                        ).to(torch.long)
+                    else:
+                        # If any motion reached start and not looping, end replay
+                        logger.info("Replay complete (reversed to start)!")
+                        self._replay_active = False
+                        return False
+            else:
+                # Playing forward (normal)
+                self._replay_time_steps += int(self._replay_speed)
+
+                # Handle end of motions per environment
+                completed = self._replay_time_steps >= self._replay_num_steps_per_env
+                if completed.any():
+                    if self._replay_loop:
+                        # Loop completed motions
+                        self._replay_time_steps[completed] = 0
+                        # Decrement loop counter if it's an integer (countdown mode)
+                        if not isinstance(self._replay_loop, bool):
+                            self._replay_loop -= 1
+                            logger.info(f"Loop completed, {self._replay_loop} loops remaining")
+                            if self._replay_loop <= 0:
+                                logger.info("All loops complete!")
+                                self._replay_active = False
+                                return False
+                    else:
+                        # If any motion is complete and not looping, end replay
+                        logger.info("Replay complete!")
+                        self._replay_active = False
+                        return False
+
+        # Update the robot states
+        self._update_replay_frame()
+
+        # Print progress every 50 frames (based on max time step)
+        max_time_step = self._replay_time_steps.max().item()
+        if max_time_step % 50 == 0 and max_time_step > 0:
+            progress = (max_time_step / self._replay_max_num_steps) * 100
+            logger.info(
+                f"Progress: {progress:.1f}% (max frame {max_time_step}/{self._replay_max_num_steps})"
+            )
+
+        return True
diff --git a/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__pycache__/__init__.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0ef6118214e14f7caef2e5db5de8b6f453ba9c45
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__pycache__/__init__.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__pycache__/__init__.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..70ec1301309a4291d49558ab74be30d7211dc14f
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__pycache__/__init__.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__pycache__/maths.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__pycache__/maths.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1a8c8ba82979ee84ff0010b14c084c0cd4dbee3e
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__pycache__/maths.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__pycache__/maths.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__pycache__/maths.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0d099691a6c95a692b4d577cb4e5157881fc1710
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__pycache__/maths.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__pycache__/rotations.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__pycache__/rotations.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..45c3be9bccc50c8e5fb451c02b045533ae09f815
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__pycache__/rotations.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__pycache__/rotations.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__pycache__/rotations.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f3bbac6e6e10ae08da658c08d8f7dcda48bb72f5
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/isaac_utils/__pycache__/rotations.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/__pycache__/__init__.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cff2c4aae78f99789e7ec51383df92c05e3d6774
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/__pycache__/__init__.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/__pycache__/__init__.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..860b30c9f54651d93b66afb9f61d2cc16e2b5976
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/__pycache__/__init__.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/__init__.py b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/__pycache__/__init__.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..abf8c2a50e7267ca4b5f0113eff0f2540d7bc661
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/__pycache__/__init__.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/__pycache__/__init__.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9e632368a2227f9d95120cd42ae98ae4772eda5c
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/__pycache__/__init__.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/__pycache__/hv_callback_handler.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/__pycache__/hv_callback_handler.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4f4cb9cb76b3e9aa6000397f5c815d69d5685b0f
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/__pycache__/hv_callback_handler.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/__pycache__/hv_callback_handler.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/__pycache__/hv_callback_handler.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..aeea37861ac64e27c1a6d32ebd0f71d1cd5e5417
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/__pycache__/hv_callback_handler.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/hv_callback_handler.py b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/hv_callback_handler.py
new file mode 100644
index 0000000000000000000000000000000000000000..5fc88df5af14f9da113a63a9ede10a5bc3251138
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/hv_callback_handler.py
@@ -0,0 +1,148 @@
+from transformers.trainer_callback import *
+from trl.trainer.ppo_trainer import *
+
+
+class HVCallbackHandler(TrainerCallback):
+    """Internal class that just calls the list of callbacks in order."""
+
+    def __init__(
+        self, callbacks, model, processing_class, optimizer, lr_scheduler, env, accelerator
+    ):
+        self.callbacks = []
+        for cb in callbacks:
+            self.add_callback(cb)
+        self.model = model
+        self.processing_class = processing_class
+        self.optimizer = optimizer
+        self.lr_scheduler = lr_scheduler
+        self.train_dataloader = None
+        self.eval_dataloader = None
+        self.env = env
+        self.accelerator = accelerator
+        if not any(isinstance(cb, DefaultFlowCallback) for cb in self.callbacks):
+            logger.warning(
+                "The Trainer will not work properly if you don't have a `DefaultFlowCallback` in its callbacks. You\n"
+                + "should add one before training with `trainer.add_callback(DefaultFlowCallback). The current list of"
+                + "callbacks is\n:"
+                + self.callback_list
+            )
+
+    def add_callback(self, callback):
+        cb = callback() if isinstance(callback, type) else callback
+        cb_class = callback if isinstance(callback, type) else callback.__class__
+        if cb_class in [c.__class__ for c in self.callbacks]:
+            logger.warning(
+                f"You are adding a {cb_class} to the callbacks of this Trainer, but there is already one. The current"
+                + "list of callbacks is\n:"
+                + self.callback_list
+            )
+        self.callbacks.append(cb)
+
+    def pop_callback(self, callback):
+        if isinstance(callback, type):
+            for cb in self.callbacks:
+                if isinstance(cb, callback):
+                    self.callbacks.remove(cb)
+                    return cb
+        else:
+            for cb in self.callbacks:
+                if cb == callback:
+                    self.callbacks.remove(cb)
+                    return cb
+
+    def remove_callback(self, callback):
+        if isinstance(callback, type):
+            for cb in self.callbacks:
+                if isinstance(cb, callback):
+                    self.callbacks.remove(cb)
+                    return
+        else:
+            self.callbacks.remove(callback)
+
+    @property
+    def callback_list(self):
+        return "\n".join(cb.__class__.__name__ for cb in self.callbacks)
+
+    def on_init_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
+        return self.call_event("on_init_end", args, state, control)
+
+    def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
+        control.should_training_stop = False
+        return self.call_event("on_train_begin", args, state, control)
+
+    def on_train_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
+        return self.call_event("on_train_end", args, state, control)
+
+    def on_epoch_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
+        control.should_epoch_stop = False
+        return self.call_event("on_epoch_begin", args, state, control)
+
+    def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
+        return self.call_event("on_epoch_end", args, state, control)
+
+    def on_step_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
+        control.should_log = False
+        control.should_evaluate = False
+        control.should_save = False
+        return self.call_event("on_step_begin", args, state, control)
+
+    def on_pre_optimizer_step(
+        self, args: TrainingArguments, state: TrainerState, control: TrainerControl
+    ):
+        return self.call_event("on_pre_optimizer_step", args, state, control)
+
+    def on_optimizer_step(
+        self, args: TrainingArguments, state: TrainerState, control: TrainerControl
+    ):
+        return self.call_event("on_optimizer_step", args, state, control)
+
+    def on_substep_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
+        return self.call_event("on_substep_end", args, state, control)
+
+    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
+        return self.call_event("on_step_end", args, state, control)
+
+    def on_evaluate(
+        self, args: TrainingArguments, state: TrainerState, control: TrainerControl, metrics
+    ):
+        control.should_evaluate = False
+        return self.call_event("on_evaluate", args, state, control, metrics=metrics)
+
+    def on_predict(
+        self, args: TrainingArguments, state: TrainerState, control: TrainerControl, metrics
+    ):
+        return self.call_event("on_predict", args, state, control, metrics=metrics)
+
+    def on_save(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
+        control.should_save = False
+        return self.call_event("on_save", args, state, control)
+
+    def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, logs):
+        control.should_log = False
+        return self.call_event("on_log", args, state, control, logs=logs)
+
+    def on_prediction_step(
+        self, args: TrainingArguments, state: TrainerState, control: TrainerControl
+    ):
+        return self.call_event("on_prediction_step", args, state, control)
+
+    def call_event(self, event, args, state, control, **kwargs):
+        for callback in self.callbacks:
+            result = getattr(callback, event)(
+                args,
+                state,
+                control,
+                model=self.model,
+                processing_class=self.processing_class,
+                optimizer=self.optimizer,
+                lr_scheduler=self.lr_scheduler,
+                train_dataloader=self.train_dataloader,
+                eval_dataloader=self.eval_dataloader,
+                env=self.env,
+                accelerator=self.accelerator,
+                **kwargs,
+            )
+            # A Callback can skip the return of `control` if it doesn't change it.
+            if result is not None:
+                control = result
+        return control
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/im_eval_callback.py b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/im_eval_callback.py
new file mode 100644
index 0000000000000000000000000000000000000000..83a3e5a749ddfcc6c7b5f5cd64e445c957f3b665
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/im_eval_callback.py
@@ -0,0 +1,880 @@
+from datetime import datetime
+import gc
+import json
+import os
+import time
+
+import numpy as np
+import torch
+from tqdm import tqdm
+from transformers import TrainerCallback
+import wandb
+
+
+def create_html_table(metrics_dict):
+    """
+    Create a sortable HTML table for metrics logging using DataTables.
+
+    Args:
+        metrics_dict: Dictionary containing metrics data with keys like 'mpjpe_g', 'mpjpe_l', 'mpjpe_pa',
+                     'terminated', 'motion_keys', etc.
+
+    Returns:
+        str: HTML string containing the sortable table
+    """
+    if not metrics_dict or len(metrics_dict) == 0:
+        return wandb.Html("<p>No metrics data available</p>")
+
+    # Get the motion keys and number of motions
+    motion_keys = metrics_dict.get("motion_keys", [])
+    if len(motion_keys) == 0:
+        return wandb.Html("<p>No motion data available</p>")
+
+    num_motions = len(motion_keys)
+
+    # Get metric names (excluding special keys)
+    special_keys = {"terminated", "motion_keys"}
+    metric_names = [key for key in metrics_dict.keys() if key not in special_keys]
+
+    # Create table header
+    html = """
+<!-- HTML Table -->
+<table id="my-table" class="display">
+  <thead>
+    <tr>
+      <th>Motion Key</th>
+      <th>Terminated</th>
+"""
+
+    # Add metric column headers
+    for metric_name in metric_names:
+        html += f"      <th>{metric_name}</th>\n"
+
+    html += """    </tr>
+  </thead>
+  <tbody>
+"""
+
+    # Create table rows
+    for i in range(num_motions):
+        motion_key = motion_keys[i] if i < len(motion_keys) else f"Motion_{i}"
+        terminated = "Yes" if metrics_dict.get("terminated", [True] * num_motions)[i] else "No"
+
+        html += f"    <tr><td>{motion_key}</td><td>{terminated}</td>"
+
+        # Add metric values
+        for metric_name in metric_names:
+            metric_values = metrics_dict[metric_name]
+            if i < len(metric_values):
+                value = metric_values[i]
+                # Format the value appropriately
+                if isinstance(value, int | float):
+                    if abs(value) < 0.001:
+                        formatted_value = f"{value:.6f}"
+                    elif abs(value) < 1:
+                        formatted_value = f"{value:.4f}"
+                    else:
+                        formatted_value = f"{value:.3f}"
+                else:
+                    formatted_value = str(value)
+                html += f"<td>{formatted_value}</td>"
+            else:
+                html += "<td>N/A</td>"
+
+        html += "</tr>\n"
+
+    html += """  </tbody>
+</table>
+
+<!-- DataTables CSS and JS dependencies -->
+<link rel="stylesheet" href="https://cdn.datatables.net/1.11.5/css/jquery.dataTables.min.css">
+<script src="https://code.jquery.com/jquery-3.3.1.min.js"></script>
+<script src="https://cdn.datatables.net/1.11.5/js/jquery.dataTables.min.js"></script>
+
+<!-- DataTables Initialization -->
+<script>
+  $(document).ready(function () {
+    $('#my-table').DataTable({
+      pageLength: 100,
+      order: [[0, 'asc']],
+    });
+  });
+</script>
+"""
+    return wandb.Html(html)
+
+
+class ImEvalCallback(TrainerCallback):
+    """Callback to evaluate motion imtiation during training. Supports multigpu ."""
+
+    def __init__(
+        self, eval_frequency, empty_cache_freq=20, eval_only=False, output_dir=None, log_keys=None
+    ):
+        super().__init__()
+        self.eval_frequency = eval_frequency
+        self.empty_cache_freq = empty_cache_freq
+        self.output_dir = output_dir
+        self.eval_only = eval_only
+        self.in_eval_mode = False
+        self.render_only = False
+        self.log_keys = log_keys
+        self._has_object = False
+
+    def on_step_end(self, args, state, control, **kwargs):
+
+        self.env = kwargs.get("env")
+        self.model = kwargs.get("model")
+        self.accelerator = kwargs.get("accelerator")
+        self.device = self.accelerator.device
+        self.args = args
+        self.model.eval()
+
+        if (state.global_step + 1) % self.eval_frequency == 0:
+            metrics_eval = self.evaluate_policy()
+
+    def save_metrics_eval(self, metrics_eval):
+        metrics_json = {}
+        for k, v in metrics_eval.items():
+            if k in ["eval/all_metrics_dict", "eval/failed_metrics_dict"]:
+                metrics_json[k] = {}
+                for kk, vv in v.items():
+                    if isinstance(vv, np.ndarray):
+                        metrics_json[k][kk] = vv.tolist()
+                    else:
+                        metrics_json[k][kk] = vv
+            elif isinstance(v, np.ndarray):
+                metrics_json[k] = v.tolist()
+            else:
+                metrics_json[k] = v
+
+        os.makedirs(self.output_dir, exist_ok=True)
+        with open(os.path.join(self.output_dir, "metrics_eval.json"), "w") as f:
+            print(f"Saving metrics_eval to {os.path.join(self.output_dir, 'metrics_eval.json')}")
+            if self.log_keys is not None:
+                metrics_json["log_keys"] = self.log_keys
+            json.dump(metrics_json, f, indent=4)
+
+    @torch.no_grad()
+    def evaluate_policy(self):
+
+        self.accelerator.wait_for_everyone()
+        with torch.no_grad():
+            self._eval_mode()
+            print(
+                "============================================Evaluating policy============================================"
+            )
+
+            self._pre_evaluate_policy()
+            actor_state = {"done_indices": [], "stop": False}
+            step = 0
+            self.eval_policy = self._get_inference_policy()
+            obs_dict = self.env.reset_all(global_rank=self.args.global_rank)
+            self.model.policy.init_rollout()
+
+            init_actions = torch.zeros(
+                self.env.num_envs, self.env.config.robot.actions_dim, device=self.env.device
+            )
+            actor_state.update({"obs": obs_dict, "actions": init_actions})
+            actor_state = self._pre_eval_env_step(actor_state)
+
+            while not actor_state.get("end_eval", False):
+                self.env.render_results()
+                actor_state["step"] = step
+                actor_state = self._pre_eval_env_step(actor_state)
+                actor_state = self.env_step(actor_state)
+                actor_state = self._post_eval_env_step(actor_state)
+                step += 1
+
+                if step % self.empty_cache_freq == 0:
+                    gc.collect()
+                    torch.cuda.empty_cache()
+
+            if self.render_only:
+                self.env.end_render_results()
+                return {}
+
+            metrics_eval = self._post_evaluate_policy(actor_state)
+
+            if self.eval_only:
+                if self.output_dir is not None:
+                    self.save_metrics_eval(metrics_eval)
+            else:
+                metrics_eval["eval/all_metrics_dict"] = create_html_table(
+                    metrics_eval["eval/all_metrics_dict"]
+                )
+                metrics_eval["eval/failed_metrics_dict"] = create_html_table(
+                    metrics_eval["eval/failed_metrics_dict"]
+                )
+
+        self._train_mode()
+        self.model.policy.clear_rollout()
+        if not self.eval_only:
+            gc.collect()
+            torch.cuda.empty_cache()
+        if self.eval_frequency == 1:  # Exit if eval frequency is 1.
+            os._exit(0)
+        return metrics_eval
+
+    def _post_evaluate_policy(self, eval_res):
+        metrics_success = eval_res["metrics_success"]
+        metrics_all = eval_res["metrics_all"]
+        metrics_eval = {}
+        for k, v in metrics_success.items():
+            metrics_eval[f"eval/success/{k}"] = v
+        for k, v in metrics_all.items():
+            metrics_eval[f"eval/all/{k}"] = v
+
+        # Add failed_keys to metrics_eval for wandb logging
+        metrics_eval["eval/all_metrics_dict"] = eval_res["all_metrics_dict"]
+        metrics_eval["eval/failed_metrics_dict"] = eval_res["failed_metrics_dict"]
+        if self.eval_only:
+            metrics_eval["failed_keys"] = eval_res["failed_keys"]
+            metrics_eval["failed_idxes"] = eval_res["failed_idxes"]
+
+        return metrics_eval
+
+    def _get_inference_policy(self, device=None):
+        self.model.policy.eval()  # switch to evaluation mode (dropout for example)
+        if device is not None:
+            self.model.policy.to(device)
+        return self.model.policy.act_inference
+
+    def _eval_mode(self):
+        if self.eval_only and self.in_eval_mode:
+            return
+        self.in_eval_mode = True
+        self.model.eval()
+        if hasattr(self.model.policy, "eval_mode"):
+            self.model.policy.eval_mode()  # For VAE, eval mode means that we are no longer sampling from the VAE but using the mean latent value.
+        self.env.set_is_evaluating(True, global_rank=self.args.global_rank)
+
+    def _train_mode(self):
+        if self.eval_only and self.in_eval_mode:
+            return
+        self.in_eval_mode = False
+        self.model.train()
+        if hasattr(self.model.policy, "train_mode"):
+            self.model.policy.train_mode()
+        self.env.set_is_evaluating(False)
+        self.env.set_is_training()
+
+    def _pre_evaluate_policy(self, reset_env=True):
+        if reset_env:
+            _ = self.env.reset_all()
+
+        self.num_total_env_eval_loops = int(
+            np.ceil(
+                self.env._motion_lib._num_unique_motions
+                / (self.env.num_envs * self.args.world_size)
+            )
+        )
+        if "max_render_envs" in self.env.config:
+            self.num_total_env_eval_loops = 1
+            self.render_only = True
+
+        self.env_eval_loop_idx = 0
+        self.pbar = tqdm(range(self.num_total_env_eval_loops), desc="Total evaluation progress")
+        self.steps_pbar = None
+        self.success_rate = 0
+        self.curr_steps = 0
+        # self.env.start_compute_metrics(global_rank=self.args.global_rank)
+        self.terminate_state = torch.zeros(self.env.num_envs, device=self.env.device)
+        self.progress_state = torch.zeros(self.env.num_envs, device=self.env.device)
+        self.terminate_memory = []
+        self.progress_memory = []
+        self.mpjpe, self.mpjpe_all = [], []
+        self.gt_pos, self.gt_pos_all = [], []
+        self.gt_rot, self.gt_rot_all = [], []
+        self.pred_pos, self.pred_pos_all = [], []
+        self.pred_rot, self.pred_rot_all = [], []
+        self.sampled_motion_idx = []
+        self.time_eval_start = time.time()
+
+        # Object tracking metrics
+        self._has_object = (
+            hasattr(self.env, "env")
+            and hasattr(self.env.env, "scene")
+            and "object" in self.env.env.scene.rigid_objects
+            and hasattr(self.env, "motion_command")
+            and self.env.motion_command is not None
+        )
+        self.obj_pos_error, self.obj_pos_error_all = [], []
+        self.obj_ori_error, self.obj_ori_error_all = [], []
+
+    def _collect_object_tracking_errors(self):
+        """Collect per-step object position and orientation errors (ref vs simulated)."""
+        try:
+            obj = self.env.env.scene["object"]
+            motion_cmd = self.env.motion_command
+            current_obj_pos = obj.data.root_pos_w[:, :3]  # (num_envs, 3)
+            current_obj_quat = obj.data.root_quat_w  # (num_envs, 4)
+            target_obj_pos = motion_cmd.object_root_pos[:, 0, :3]  # (num_envs, 3)
+            target_obj_quat = motion_cmd.object_root_quat[:, 0]  # (num_envs, 4)
+
+            pos_error = torch.norm(target_obj_pos - current_obj_pos, dim=-1)  # (num_envs,)
+
+            # Quaternion error: angle between two quaternions
+            from isaaclab.utils.math import quat_error_magnitude
+
+            ori_error = quat_error_magnitude(target_obj_quat, current_obj_quat)  # (num_envs,)
+
+            self.obj_pos_error.append(pos_error.cpu())
+            self.obj_ori_error.append(ori_error.cpu())
+        except (KeyError, AttributeError, IndexError):
+            # Gracefully handle missing object, motion data, or shape mismatches
+            pass
+
+    def env_step(self, actor_state):
+        obs_dict, rewards, dones, extras = self.env.step(actor_state)
+        actor_state.update({"obs": obs_dict, "rewards": rewards, "dones": dones, "extras": extras})
+        return actor_state
+
+    def _pre_eval_env_step(self, actor_state: dict):
+        dones = actor_state.get("dones", torch.zeros(self.env.num_envs, device=self.env.device))
+        actions = self.eval_policy(
+            obs_dict=actor_state["obs"], cur_dones=dones, skip_episode_attnmask=True
+        )
+        actor_state.update({"actions": actions})
+        return actor_state
+
+    def _post_eval_env_step(self, actor_state):
+        step = actor_state["step"]
+        actor_state["end_eval"] = False
+
+        if "ref_body_pos_extend" in self.env.extras:
+            self.gt_pos.append(self.env.extras["ref_body_pos_extend"].cpu().numpy())
+            self.pred_pos.append(self.env.extras["rigid_body_pos_extend"].cpu().numpy())
+            self.mpjpe.append(self.env.dif_global_body_pos.norm(dim=-1).cpu() * 1000)
+        else:
+            gt_pos = self.env.get_env_data("ref_body_pos_extend")
+            pred_pos = self.env.get_env_data("rigid_body_pos_extend")
+            mpjpe = (gt_pos - pred_pos).norm(dim=-1) * 1000
+            self.gt_pos.append(gt_pos.cpu().numpy())
+            self.pred_pos.append(pred_pos.cpu().numpy())
+            self.mpjpe.append(mpjpe.cpu())
+
+        # Collect object tracking errors if object exists in scene
+        if self._has_object:
+            self._collect_object_tracking_errors()
+
+        # self.gt_rot.append(self.env.extras['ref_body_rot_extend'].cpu().numpy())
+        # self.pred_rot.append(self.env._rigid_body_rot_extend.cpu().numpy())
+
+        died = actor_state["dones"]
+        died[actor_state["extras"]["time_outs"]] = False
+
+        termination_state = torch.logical_and(
+            self.curr_steps <= self.env._motion_lib.get_motion_num_steps(self.env.motion_ids) - 1,
+            died,
+        )  # if terminate after the last frame, then it is not a termination. curr_step is one step behind simulation.
+        self.terminate_state = torch.logical_or(termination_state, self.terminate_state)
+
+        self.progress_state[~self.terminate_state] += 1
+
+        if (~self.terminate_state).sum() > 0:
+            max_possible_id = self.env._motion_lib._num_unique_motions - 1
+            curr_ids = self.env._motion_lib._curr_motion_ids
+            if (max_possible_id == curr_ids).sum() > 0:  # When you are running out of motions.
+                bound = (max_possible_id == curr_ids).nonzero()[0] + 1
+                if (~self.terminate_state[:bound]).sum() > 0:
+                    curr_max = (
+                        self.env._motion_lib.get_motion_num_steps(self.env.motion_ids)[:bound][
+                            ~self.terminate_state[:bound]
+                        ]
+                        .max()
+                        .item()
+                    )
+                else:
+                    curr_max = self.curr_steps - 1  # the ones that should be counted have teimrated
+            else:
+                curr_max = (
+                    self.env._motion_lib.get_motion_num_steps(self.env.motion_ids)[
+                        ~self.terminate_state
+                    ]
+                    .max()
+                    .item()
+                )
+
+            if self.curr_steps >= curr_max:
+                curr_max = self.curr_steps + 1  # For matching up the current steps and max steps.
+        else:
+            curr_max = self.env._motion_lib.get_motion_num_steps(self.env.motion_ids).max().item()
+
+        if self.steps_pbar is None and (~self.terminate_state).sum() > 0:
+            self.steps_pbar = tqdm(total=int(curr_max), desc="Sequence progress", leave=False)
+
+        if self.steps_pbar is not None:
+            self.steps_pbar.update(1)
+            if self.steps_pbar.total != int(curr_max):
+                self.steps_pbar.total = int(curr_max)
+                self.steps_pbar.refresh()
+
+        self.curr_steps += 1
+        if self.curr_steps >= curr_max or self.terminate_state.sum() == self.env.num_envs:
+            if self.steps_pbar is not None:
+                self.steps_pbar.close()
+                self.steps_pbar = None
+
+            self.terminate_memory.append(self.terminate_state.cpu().numpy())
+            self.progress_memory.append(
+                (
+                    self.progress_state
+                    / self.env._motion_lib.get_motion_num_steps(self.env.motion_ids)
+                )
+                .cpu()
+                .numpy()
+            )
+
+            self.success_rate = (
+                1
+                - np.concatenate(self.terminate_memory)[
+                    : self.env._motion_lib._num_unique_motions
+                ].mean()
+            )
+            self.progress_rate = np.concatenate(self.progress_memory)[
+                : self.env._motion_lib._num_unique_motions
+            ].mean()
+
+            # MPJPE
+            all_mpjpe = torch.stack(self.mpjpe)
+            try:
+                assert (
+                    all_mpjpe.shape[0] == curr_max
+                    or self.terminate_state.sum() == self.env.num_envs
+                )  # Max should be the same as the number of frames in the motion.
+            except AssertionError:
+                print(
+                    f"Warning: MPJPE shape mismatch: {all_mpjpe.shape[0]} vs curr_max={curr_max}, terminated={self.terminate_state.sum()}/{self.env.num_envs}"
+                )
+
+            all_body_pos_pred = np.stack(self.pred_pos)
+            all_body_pos_gt = np.stack(self.gt_pos)
+            # all_body_rot_pred = np.stack(self.pred_rot)
+            # all_body_rot_gt = np.stack(self.gt_rot)
+
+            all_mpjpe = [
+                all_mpjpe[: (i - 1), idx].mean()
+                for idx, i in enumerate(
+                    self.env._motion_lib.get_motion_num_steps(self.env.motion_ids)
+                )
+            ]  # -1 since we do not count the first frame.
+            all_body_pos_pred = [
+                all_body_pos_pred[: (i - 1), idx]
+                for idx, i in enumerate(
+                    self.env._motion_lib.get_motion_num_steps(self.env.motion_ids)
+                )
+            ]
+            all_body_pos_gt = [
+                all_body_pos_gt[: (i - 1), idx]
+                for idx, i in enumerate(
+                    self.env._motion_lib.get_motion_num_steps(self.env.motion_ids)
+                )
+            ]
+            # all_body_rot_pred = [all_body_rot_pred[: (i - 1), idx] for idx, i in enumerate(self.env._motion_lib.get_motion_num_steps())]
+            # all_body_rot_gt = [all_body_rot_gt[: (i - 1), idx] for idx, i in enumerate(self.env._motion_lib.get_motion_num_steps())]
+
+            self.mpjpe_all.append(all_mpjpe)
+            self.pred_pos_all += all_body_pos_pred
+            self.gt_pos_all += all_body_pos_gt
+            # self.pred_rot_all += all_body_rot_pred
+            # self.gt_rot_all += all_body_rot_gt
+
+            # Aggregate object tracking errors for this batch
+            if self._has_object and len(self.obj_pos_error) > 0:
+                all_obj_pos_err = torch.stack(self.obj_pos_error)  # (T, num_envs)
+                all_obj_ori_err = torch.stack(self.obj_ori_error)  # (T, num_envs)
+                motion_num_steps = self.env._motion_lib.get_motion_num_steps(self.env.motion_ids)
+                per_env_obj_pos_err = [
+                    all_obj_pos_err[: (i - 1), idx].mean().item()
+                    for idx, i in enumerate(motion_num_steps)
+                ]
+                per_env_obj_ori_err = [
+                    all_obj_ori_err[: (i - 1), idx].mean().item()
+                    for idx, i in enumerate(motion_num_steps)
+                ]
+                self.obj_pos_error_all.append(per_env_obj_pos_err)
+                self.obj_ori_error_all.append(per_env_obj_ori_err)
+
+            env_motion_ids = self.env.start_idx + self.env.motion_ids
+            self.sampled_motion_idx.append(env_motion_ids)
+            self.env_eval_loop_idx += 1
+
+            if self.env_eval_loop_idx >= self.num_total_env_eval_loops:
+                if self.render_only:
+                    print("Rendering only. Reached the end of the evaluation loop.")
+                    self.env.end_render_results()
+                    actor_state["end_eval"] = True
+                    return actor_state
+
+                terminate_hist = np.concatenate(self.terminate_memory)
+                progress_hist = np.concatenate(self.progress_memory)
+                succ_idxes = np.nonzero(
+                    ~terminate_hist[: self.env._motion_lib._num_unique_motions]
+                )[0].tolist()
+                self.accelerator.wait_for_everyone()
+                # metrics_all = compute_metrics_lite(self.pred_pos_all, self.gt_pos_all, self.pred_rot_all, self.gt_rot_all, concatenate = False) # OOM
+
+                print(
+                    f"!!!!!!! {len(self.pred_pos_all)} {len(self.gt_pos_all)} {self.env.start_idx} {self.args.global_rank} Time: {datetime.now().strftime('%H:%M:%S')}"
+                )
+
+                if hasattr(self.env, "motion_command"):
+                    body_names = self.env.motion_command.cmd_body_names
+                else:
+                    print("No self.env.motion_command.cmd_body_names found!!!!")
+                    exit()
+
+                """
+                # gear_sonic/config/manager_env/commands/terms/motion.yaml
+                body_names: [
+                    "pelvis",
+                    "left_hip_roll_link",
+                    "left_knee_link",
+                    "left_ankle_roll_link",
+                    "right_hip_roll_link",
+                    "right_knee_link",
+                    "right_ankle_roll_link",
+                    "torso_link",
+                    "left_shoulder_roll_link",
+                    "left_elbow_link",
+                    "left_wrist_yaw_link",
+                    "right_shoulder_roll_link",
+                    "right_elbow_link",
+                    "right_wrist_yaw_link",
+                ]
+                """
+
+                # Define subsets
+                # 6 + 3 + 5 = 14
+                legs_subset_names = [
+                    "left_hip_roll_link",
+                    "left_knee_link",
+                    "left_ankle_roll_link",
+                    "right_hip_roll_link",
+                    "right_knee_link",
+                    "right_ankle_roll_link",
+                ]
+                # NOTE use torso_link instead of head for vr_3points_subset_names
+                vr_3points_subset_names = [
+                    "torso_link",
+                    "left_wrist_yaw_link",
+                    "right_wrist_yaw_link",
+                ]
+                other_upper_bodies_subset_names = [
+                    "pelvis",
+                    "left_shoulder_roll_link",
+                    "left_elbow_link",
+                    "right_shoulder_roll_link",
+                    "right_elbow_link",
+                ]
+
+                foot_subset_names = ["left_ankle_roll_link", "right_ankle_roll_link"]
+
+                # Get indices for subsets
+                legs_indices = [body_names.index(name) for name in legs_subset_names]
+                vr_3points_indices = [body_names.index(name) for name in vr_3points_subset_names]
+                other_upper_bodies_indices = [
+                    body_names.index(name) for name in other_upper_bodies_subset_names
+                ]
+                foot_indices = [body_names.index(name) for name in foot_subset_names]
+                # Extract subset data
+                pred_pos_legs = [p[:, legs_indices, :] for p in self.pred_pos_all]
+                gt_pos_legs = [g[:, legs_indices, :] for g in self.gt_pos_all]
+
+                pred_pos_foot = [p[:, foot_indices, :] for p in self.pred_pos_all]
+                gt_pos_foot = [g[:, foot_indices, :] for g in self.gt_pos_all]
+
+                pred_pos_vr_3points = [p[:, vr_3points_indices, :] for p in self.pred_pos_all]
+                gt_pos_vr_3points = [g[:, vr_3points_indices, :] for g in self.gt_pos_all]
+
+                pred_pos_other_upper_bodies = [
+                    p[:, other_upper_bodies_indices, :] for p in self.pred_pos_all
+                ]
+                gt_pos_other_upper_bodies = [
+                    g[:, other_upper_bodies_indices, :] for g in self.gt_pos_all
+                ]
+
+                # Lazy import to avoid cffi version conflict with IsaacSim
+                from smpl_sim.smpllib.smpl_eval import compute_metrics_lite
+
+                metrics_all = compute_metrics_lite(
+                    self.pred_pos_all, self.gt_pos_all, concatenate=False
+                )  # list of length N_env
+                metrics_legs = compute_metrics_lite(pred_pos_legs, gt_pos_legs, concatenate=False)
+                metrics_vr_3points = compute_metrics_lite(
+                    pred_pos_vr_3points, gt_pos_vr_3points, concatenate=False
+                )
+                metrics_other_upper_bodies = compute_metrics_lite(
+                    pred_pos_other_upper_bodies, gt_pos_other_upper_bodies, concatenate=False
+                )
+                metrics_foot = compute_metrics_lite(pred_pos_foot, gt_pos_foot, concatenate=False)
+
+                # Rename keys for subset metrics
+                metrics_legs = {f"{k}_legs": v for k, v in metrics_legs.items()}
+                metrics_vr_3points = {f"{k}_vr_3points": v for k, v in metrics_vr_3points.items()}
+                metrics_other_upper_bodies = {
+                    f"{k}_other_upper_bodies": v for k, v in metrics_other_upper_bodies.items()
+                }
+                metrics_foot = {f"{k}_foot": v for k, v in metrics_foot.items()}
+
+                metrics_all.update(metrics_legs)
+                metrics_all.update(metrics_vr_3points)
+                metrics_all.update(metrics_other_upper_bodies)
+                metrics_all.update(metrics_foot)
+
+                metrics_all_sum = {
+                    k: torch.tensor(
+                        [np.sum(i) / i.shape[1] if "mpjpe" in k else np.sum(i) for i in v]
+                    ).to(self.env.device)
+                    for k, v in metrics_all.items()
+                }  # of length N_env -- mean over joint but sum over length
+                length_all = torch.tensor([len(i) for i in self.pred_pos_all]).to(self.env.device)
+
+                metrics_all_contactnate = torch.stack(
+                    [v for k, v in metrics_all_sum.items()] + [length_all], dim=-1
+                )
+                terminate_hist_concatenate = torch.tensor(terminate_hist).to(self.env.device)
+                progress_hist_concatenate = torch.tensor(progress_hist).to(self.env.device)
+                all_motion_idxes = torch.cat(self.sampled_motion_idx).to(self.env.device)
+
+                # Prepare object tracking metrics for gathering
+                has_obj_metrics = self._has_object and len(self.obj_pos_error_all) > 0
+                if has_obj_metrics:
+                    obj_pos_err_flat = torch.tensor(
+                        [v for batch in self.obj_pos_error_all for v in batch]
+                    ).to(self.env.device)
+                    obj_ori_err_flat = torch.tensor(
+                        [v for batch in self.obj_ori_error_all for v in batch]
+                    ).to(self.env.device)
+
+                # Tensor layout: [metrics_all_sum..., length, terminate, progress, (obj_pos_err, obj_ori_err,) motion_idx]
+                tail_tensors = [
+                    terminate_hist_concatenate[:, None],
+                    progress_hist_concatenate[:, None],
+                ]
+                if has_obj_metrics:
+                    tail_tensors.append(obj_pos_err_flat[:, None])
+                    tail_tensors.append(obj_ori_err_flat[:, None])
+                tail_tensors.append(all_motion_idxes[:, None])
+
+                all_tensors = torch.cat(
+                    [metrics_all_contactnate] + tail_tensors,
+                    dim=-1,
+                )
+                print("Gathering eval tensors", all_tensors.shape, self.accelerator.process_index)
+
+                chunk_size = 1024  # Chunk gathering since it's 4096 is too large.
+                chunks = all_tensors.split(chunk_size)
+                gathered_chunks = [
+                    self.accelerator.gather(chunk).reshape(-1, *chunk.shape) for chunk in chunks
+                ]  # each with shape (num_processes x 1024 (chunked_num_env), D_metrics)
+                all_metrics = torch.cat(gathered_chunks, dim=1)
+
+                metric_size = all_metrics.shape[-1]
+                gathered_metrics_stack = (
+                    all_metrics.reshape(
+                        self.accelerator.num_processes, -1, self.env.num_envs, metric_size
+                    )
+                    .transpose(0, 1)
+                    .reshape(-1, metric_size)[: self.env._motion_lib._num_unique_motions]
+                )  # make sure that we are selecting the correct ones.
+
+                # Extract tail columns: terminate, progress, (obj_pos_err, obj_ori_err,) motion_idx
+                num_tail = 3 + (
+                    2 if has_obj_metrics else 0
+                )  # terminate + progress + (obj*2) + motion_idx
+                num_body_metrics = metric_size - num_tail  # metrics_all_sum columns + length
+
+                gathered_terminate_hist_stack = gathered_metrics_stack[:, num_body_metrics].bool()
+                gathered_progress_hist_stack = gathered_metrics_stack[:, num_body_metrics + 1]
+                if has_obj_metrics:
+                    gathered_obj_pos_err = gathered_metrics_stack[:, num_body_metrics + 2]
+                    gathered_obj_ori_err = gathered_metrics_stack[:, num_body_metrics + 3]
+                    gathered_motion_idxes = gathered_metrics_stack[:, num_body_metrics + 4].long()
+                else:
+                    gathered_motion_idxes = gathered_metrics_stack[:, num_body_metrics + 2].long()
+                gathered_progress_hist_stack[~gathered_terminate_hist_stack] = 1
+
+                assert (gathered_motion_idxes.diff(dim=0) == 1).all()
+
+                # Micro-average: sum all frame-level sums, divide by total frames
+                # (each timestep weighted equally, longer motions contribute more)
+                metric_sums = gathered_metrics_stack[:, : num_body_metrics - 1]
+                frame_counts = gathered_metrics_stack[:, num_body_metrics - 1 : num_body_metrics]
+
+                success_mask = ~gathered_terminate_hist_stack
+                success_metrics_mean = metric_sums[success_mask].sum(dim=0) / frame_counts[
+                    success_mask
+                ].sum(dim=0)
+                all_metrics_mean = metric_sums.sum(dim=0) / frame_counts.sum(dim=0)
+
+                # Also keep per-motion metrics for downstream use
+                all_metrics = metric_sums / frame_counts
+                metrics_all_print = {
+                    k: all_metrics_mean[idx].cpu().numpy()
+                    for idx, (k, v) in enumerate(metrics_all_sum.items())
+                }
+                metrics_succ_print = {
+                    k: success_metrics_mean[idx].cpu().numpy()
+                    for idx, (k, v) in enumerate(metrics_all_sum.items())
+                }
+
+                # Add object tracking metrics to printed summaries
+                if has_obj_metrics:
+                    obj_pos_err_mean = gathered_obj_pos_err.mean().cpu().numpy()
+                    obj_ori_err_mean = gathered_obj_ori_err.mean().cpu().numpy()
+                    obj_pos_err_succ = (
+                        gathered_obj_pos_err[~gathered_terminate_hist_stack].mean().cpu().numpy()
+                        if (~gathered_terminate_hist_stack).any()
+                        else 0.0
+                    )
+                    obj_ori_err_succ = (
+                        gathered_obj_ori_err[~gathered_terminate_hist_stack].mean().cpu().numpy()
+                        if (~gathered_terminate_hist_stack).any()
+                        else 0.0
+                    )
+                    metrics_all_print["obj_pos_error"] = obj_pos_err_mean
+                    metrics_all_print["obj_ori_error"] = obj_ori_err_mean
+                    metrics_succ_print["obj_pos_error"] = obj_pos_err_succ
+                    metrics_succ_print["obj_ori_error"] = obj_ori_err_succ
+
+                failed_keys = self.env._motion_lib._motion_data_keys[
+                    gathered_terminate_hist_stack.cpu().numpy()
+                ]
+                success_keys = self.env._motion_lib._motion_data_keys[
+                    ~gathered_terminate_hist_stack.cpu().numpy()
+                ]
+                success_rate = 1 - gathered_terminate_hist_stack.cpu().numpy().mean()
+                progress_rate = gathered_progress_hist_stack.cpu().numpy().mean()
+
+                all_metrics_dict = {
+                    k: all_metrics[:, idx].cpu().numpy()
+                    for idx, (k, v) in enumerate(metrics_all_sum.items())
+                }
+                all_metrics_dict["terminated"] = gathered_terminate_hist_stack.cpu().numpy()
+                all_metrics_dict["progress"] = gathered_progress_hist_stack.cpu().numpy()
+                all_metrics_dict["motion_keys"] = self.env._motion_lib._motion_data_keys[
+                    gathered_motion_idxes.cpu().numpy()
+                ]
+                all_metrics_dict["sampling_prob"] = (
+                    self.env._motion_lib._sampling_prob[gathered_motion_idxes.cpu().numpy()]
+                    .cpu()
+                    .numpy()
+                )
+                # Add per-motion object tracking metrics
+                if has_obj_metrics:
+                    all_metrics_dict["obj_pos_error"] = gathered_obj_pos_err.cpu().numpy()
+                    all_metrics_dict["obj_ori_error"] = gathered_obj_ori_err.cpu().numpy()
+                    # Save per-env obj_pos_error for threshold-based success rate analysis
+                    if self.eval_only and len(self.obj_pos_error_all) > 0:
+                        all_metrics_dict["per_env_obj_pos_error"] = [
+                            v for batch in self.obj_pos_error_all for v in batch
+                        ]
+                        all_metrics_dict["per_env_obj_ori_error"] = [
+                            v for batch in self.obj_ori_error_all for v in batch
+                        ]
+
+                failed_metrics_dict = {
+                    k: all_metrics[gathered_terminate_hist_stack, idx].cpu().numpy()
+                    for idx, (k, v) in enumerate(metrics_all_sum.items())
+                }
+                failed_metrics_dict["motion_keys"] = failed_keys
+                failed_metrics_dict["sampling_prob"] = (
+                    self.env._motion_lib._sampling_prob[gathered_terminate_hist_stack.cpu().numpy()]
+                    .cpu()
+                    .numpy()
+                )
+                if has_obj_metrics:
+                    failed_metrics_dict["obj_pos_error"] = (
+                        gathered_obj_pos_err[gathered_terminate_hist_stack].cpu().numpy()
+                    )
+                    failed_metrics_dict["obj_ori_error"] = (
+                        gathered_obj_ori_err[gathered_terminate_hist_stack].cpu().numpy()
+                    )
+
+                if self.accelerator.is_main_process:
+                    print(f"Success Rate: {success_rate:.10f}")
+                    print(f"Progress Rate: {progress_rate:.10f}")
+                    if has_obj_metrics:
+                        print(
+                            f"Object Pos Error (all): {obj_pos_err_mean:.4f}m | "
+                            f"Object Ori Error (all): {obj_ori_err_mean:.4f}rad"
+                        )
+                    print(
+                        "All: ", " \t".join([f"{k}: {v:.3f}" for k, v in metrics_all_print.items()])
+                    )
+                    print(
+                        "Succ: ",
+                        " \t".join([f"{k}: {v:.3f}" for k, v in metrics_succ_print.items()]),
+                    )
+
+                metrics_succ_print["success_rate"] = success_rate
+                metrics_succ_print["progress_rate"] = progress_rate
+                actor_state["metrics_all"] = metrics_all_print
+                actor_state["metrics_success"] = metrics_succ_print
+                actor_state["failed_keys"] = failed_keys
+                actor_state["success_keys"] = success_keys
+                actor_state["all_metrics_dict"] = all_metrics_dict
+                actor_state["failed_metrics_dict"] = failed_metrics_dict
+                actor_state["failed_idxes"] = (
+                    gathered_terminate_hist_stack.cpu().numpy().nonzero()[0]
+                )
+
+                if not self.eval_only:
+                    del (
+                        self.mpjpe,
+                        self.mpjpe_all,
+                        self.gt_pos,
+                        self.gt_pos_all,
+                        self.gt_rot,
+                        self.gt_rot_all,
+                        self.pred_pos,
+                        self.pred_pos_all,
+                        self.pred_rot,
+                        self.pred_rot_all,
+                        self.sampled_motion_idx,
+                        self.obj_pos_error,
+                        self.obj_pos_error_all,
+                        self.obj_ori_error,
+                        self.obj_ori_error_all,
+                    )
+                    gc.collect()
+                    torch.cuda.empty_cache()
+
+                actor_state["end_eval"] = True
+                self.pbar.update(1)
+                self.pbar.refresh()
+                return actor_state
+
+            self.env.forward_motion_samples(self.args.global_rank, self.args.world_size)
+            self.terminate_state = torch.zeros(self.env.num_envs, device=self.device)
+            self.progress_state = torch.zeros(self.env.num_envs, device=self.env.device)
+
+            self.success_rate = 0
+            self.curr_steps = 0
+
+            self.pbar.update(1)
+            self.pbar.refresh()
+            (
+                self.mpjpe,
+                self.gt_pos,
+                self.pred_pos,
+                self.obj_pos_error,
+                self.obj_ori_error,
+            ) = (
+                [],
+                [],
+                [],
+                [],
+                [],
+            )
+
+        eval_time = (time.time() - self.time_eval_start) / 60  # in minutes
+        obj_str = ""
+        if self._has_object and len(self.obj_pos_error_all) > 0:
+            mean_obj_pos_err = np.mean([v for batch in self.obj_pos_error_all for v in batch])
+            obj_str = f" | ObjPosErr: {mean_obj_pos_err:.4f}m"
+        update_str = f"Terminated: {self.terminate_state.sum().item()} | max frames: {curr_max} | steps {self.curr_steps} | env_loop: {self.env_eval_loop_idx} | eval_time: {eval_time:.1f}m | Start: {self.env.start_idx} | Succ rate: {self.success_rate:.3f} | Mpjpe: {np.mean(self.mpjpe_all) * 1000:.3f}{obj_str}"
+        self.pbar.set_description(update_str)
+
+        return actor_state
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/im_resample_callback.py b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/im_resample_callback.py
new file mode 100644
index 0000000000000000000000000000000000000000..27bd6f1e3a79584e5d41ea92537ae3966c16d717
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/im_resample_callback.py
@@ -0,0 +1,25 @@
+from transformers import TrainerCallback
+
+
+class ImResampleCallback(TrainerCallback):
+    """Callback to resample motion during training. Supports multigpu ."""
+
+    def __init__(self, motion_resample_frequency, skip_resample_frequency=None):
+        super().__init__()
+        self.motion_resample_frequency = motion_resample_frequency
+        self.skip_resample_frequency = skip_resample_frequency
+
+    def on_step_end(self, args, state, control, **kwargs):
+
+        self.env = kwargs.get("env")
+        self.accelerator = kwargs.get("accelerator")
+        self.device = self.accelerator.device
+
+        should_resample = (state.global_step + 1) % self.motion_resample_frequency == 0
+        should_skip = (
+            self.skip_resample_frequency is not None
+            and (state.global_step + 1) % self.skip_resample_frequency == 0
+        )
+
+        if should_resample and not should_skip:
+            self.env.resample_motion()
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/model_save_callback.py b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/model_save_callback.py
new file mode 100644
index 0000000000000000000000000000000000000000..75f47e1d1ff0fccc4c510a358b31cfdd0fe02006
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/model_save_callback.py
@@ -0,0 +1,176 @@
+import copy
+import os
+from pathlib import Path
+import shutil
+import subprocess
+import sys
+
+from loguru import logger
+import torch
+from transformers import TrainerCallback
+import wandb
+
+from gear_sonic.trl.utils.common import wandb_run_exists
+
+
+class ModelSaveCallback(TrainerCallback):
+    """Callback to save model state_dict during training."""
+
+    def __init__(self, save_dir, save_frequency=1000, save_last_frequency=50, max_disk_usage=None):
+        """
+        Args:
+            save_dir (str): Directory to save model checkpoints
+            save_frequency (int): Save model every N steps
+            max_disk_usage (float): Maximum disk usage in TB
+        """
+        self.save_dir = Path(save_dir)
+        self.save_frequency = save_frequency
+        self.save_last_frequency = save_last_frequency
+        self.save_dir.mkdir(parents=True, exist_ok=True)
+        self.save_last_only = False
+        self.max_disk_usage = max_disk_usage
+
+    def check_disk_usage(self):
+        """Check if current working directory has more than 0.9TB of usage."""
+        if self.save_last_only or self.max_disk_usage is None:
+            return
+
+        try:
+            result = subprocess.run(["du", "-sb", "."], capture_output=True, text=True, check=True)
+            size_bytes = int(result.stdout.split()[0])
+            size_tb = size_bytes / (1024**4)
+            if size_tb > self.max_disk_usage:
+                self.save_last_only = True
+                logger.info(
+                    f"Directory size {size_tb:.2f}TB > {self.max_disk_usage}TB, setting save_last_only=True"
+                )
+                if wandb_run_exists():
+                    wandb.alert(
+                        title="Disk usage warning",
+                        text=f"Directory size {size_tb:.2f}TB > {self.max_disk_usage}TB, setting save_last_only=True",
+                        level="WARN",
+                    )
+        except Exception as e:
+            logger.error(f"Error checking disk usage: {e}. Skip!")
+
+    def on_step_end(self, args, state, control, **kwargs):
+        """Save model state_dict at the end of each step if frequency matches."""
+        model = kwargs.get("model")
+        optimizer = kwargs.get("optimizer")
+        lr_scheduler = kwargs.get("lr_scheduler")
+        env = kwargs.get("env")
+
+        if state.is_world_process_zero and not env.is_evaluating:
+            # Only save regular checkpoints if save_last_only is False
+            if not self.save_last_only and state.global_step % self.save_frequency == 0:
+                env_state_dict = env.get_env_state_dict()
+                ModelSaveCallback.save_checkpoint(
+                    model,
+                    optimizer,
+                    lr_scheduler,
+                    state,
+                    env_state_dict,
+                    args,
+                    f"{self.save_dir}/model_step_{state.global_step:06d}.pt",
+                )
+
+            # Always save last checkpoint every 50 steps
+            if state.global_step % self.save_last_frequency == 0:
+                env_state_dict = env.get_env_state_dict()
+                ModelSaveCallback.save_checkpoint(
+                    model,
+                    optimizer,
+                    lr_scheduler,
+                    state,
+                    env_state_dict,
+                    args,
+                    f"{self.save_dir}/last.pt",
+                )
+                # self.export_policy_to_onnx(env, state, model)
+
+    def get_example_obs(self, env):
+        obs_dict = copy.deepcopy(env.obs_buf_dict)
+        for obs_key in obs_dict.keys():
+            print(obs_key, sorted(env.config.obs.obs_dict[obs_key]))
+        # move to cpu
+        for k in obs_dict:
+            obs_dict[k] = obs_dict[k].cpu()[0:1]
+        return obs_dict
+
+    def export_policy_to_onnx(self, env, state, model):
+        checkpoint_path = os.path.join(env.config.experiment_dir, "last.pt")
+        cmd = [
+            sys.executable,
+            "gear_sonic/eval_agent_trl.py",
+            f"+checkpoint={checkpoint_path}",
+            "+num_envs=1",
+            "+headless=true",
+            "+export_onnx_only=true",
+        ]
+        result = subprocess.run(cmd, capture_output=False, text=True, cwd=os.getcwd())
+        onnx_last_path = os.path.join(env.config.experiment_dir, "exported", "last.onnx")
+        onnx_step_path = os.path.join(
+            env.config.experiment_dir, "exported", f"model_step_{state.global_step:06d}.onnx"
+        )
+        shutil.copy(onnx_last_path, onnx_step_path)
+
+    @classmethod
+    def save_checkpoint(
+        cls, model, optimizer, lr_scheduler, state, env_state_dict, args, save_path
+    ):
+        if model is not None:
+            # Save model, optimizer, scheduler and training state
+            state_without_log_history = copy.copy(state)
+            state_without_log_history.__dict__.pop("log_history")
+            _state = copy.deepcopy(state_without_log_history)
+            checkpoint = {
+                "policy_state_dict": model.policy.state_dict(),
+                "value_state_dict": (
+                    model.value_model.state_dict() if model.value_model is not None else None
+                ),  # Value model is not always preset (e.g. for distillation)
+                "optimizer_state_dict": optimizer.state_dict() if optimizer is not None else None,
+                "lr_scheduler_state_dict": (
+                    lr_scheduler.state_dict() if lr_scheduler is not None else None
+                ),
+                "state": _state,
+                "args": args,
+                "env_state_dict": env_state_dict,
+            }
+
+            if hasattr(model, "disc_model") and model.disc_model is not None:
+                checkpoint["disc_state_dict"] = model.disc_model.state_dict()
+
+            import tempfile
+            import time
+
+            save_dir = os.path.dirname(save_path)
+
+            for attempt in range(5):
+                try:
+                    # Save to a temp file first, then atomically rename
+                    # This prevents corrupted partial checkpoints on filesystem failures
+                    with tempfile.NamedTemporaryFile(
+                        dir=save_dir, delete=False, suffix=".pt.tmp"
+                    ) as tmp_file:
+                        tmp_path = tmp_file.name
+
+                    torch.save(checkpoint, tmp_path)
+
+                    # Atomic rename (os.replace is atomic on POSIX systems)
+                    os.replace(tmp_path, save_path)
+                    print(f"Saved model checkpoint to {save_path}")
+                    break
+                except Exception as e:
+                    # Clean up temp file if it exists
+                    if "tmp_path" in locals() and os.path.exists(tmp_path):
+                        try:
+                            os.remove(tmp_path)
+                        except:
+                            pass
+
+                    if attempt == 4:  # Last attempt
+                        print(f"Failed to save checkpoint after 5 attempts. Error: {e}")
+                    print(f"Attempt {attempt + 1} failed to save checkpoint. Retrying...")
+                    time.sleep(
+                        5
+                    )  # Wait a bit before retrying (helps with transient filesystem issues)
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/read_eval_callback.py b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/read_eval_callback.py
new file mode 100644
index 0000000000000000000000000000000000000000..aa0a081f65a165ec1fe85320b7173d2b60b48c34
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/read_eval_callback.py
@@ -0,0 +1,163 @@
+import json
+from pathlib import Path
+
+from transformers import TrainerCallback
+import wandb
+
+from gear_sonic.trl.callbacks.im_eval_callback import create_html_table
+from gear_sonic.trl.utils.common import wandb_run_exists
+
+
+class ReadEvalCallback(TrainerCallback):
+    """Callback to read evaluation metrics and log them to wandb."""
+
+    def __init__(self, eval_dir: str, check_interval: int = 1):
+        """
+        Initialize the callback.
+
+        Args:
+            experiment_dir: Path to the experiment directory where eval results are stored
+            check_interval: How often to check for new evaluation results (in log steps)
+        """
+        super().__init__()
+        self.eval_dir = Path(eval_dir)
+        self.check_interval = check_interval
+        self.last_check_step = 0
+        self.state = None
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        del args, control, logs
+        return  # Disable this callback for now. We will let eval script to handle the wandb logging.
+
+    def find_next_unread_checkpoint(self, current_eval_step: int = 0, mode: str = "metrics"):
+        """
+        Find the next unread checkpoint folder that has finished evaluation.
+
+        Args:
+            current_eval_step: Current evaluation step to start searching from
+
+        Returns:
+            tuple: (eval_step, eval_step_dir) or (None, None) if not found
+        """
+        if not self.eval_dir.exists():
+            return []
+
+        # Find all evaluation directories with 6-digit zero-padded names
+        eval_step_dirs = []
+        for d in self.eval_dir.iterdir():
+            if d.is_dir() and d.name.isdigit():
+                eval_step = int(d.name)
+                if eval_step > current_eval_step:
+                    eval_step_dirs.append((eval_step, d))
+
+        if not eval_step_dirs:
+            return []
+
+        # Sort by step number to get the next one in sequence
+        eval_step_dirs.sort(key=lambda x: x[0])
+
+        # Find the first one that has finished evaluation
+        for eval_step, eval_step_dir in eval_step_dirs:
+            metrics_finish_file = eval_step_dir / "all_eval_finish.txt"
+            if metrics_finish_file.exists():
+                eval_dir_res = []
+                for f in eval_step_dir.iterdir():
+                    if f.is_dir():
+                        mode_finish_file = f / f"{mode}_finish.txt"
+                        if mode_finish_file.exists():
+                            eval_dir_res.append((eval_step, f))
+                return eval_dir_res
+
+        return []
+
+    def _check_and_log_eval_results(self, eval_step, eval_step_dir):
+        """Check for new evaluation results and log them to wandb."""
+
+        metrics_file = eval_step_dir / "metrics_eval.json"
+        # Read and log the metrics
+        with open(metrics_file) as f:
+            try:
+                metrics_eval = json.load(f)
+            except json.JSONDecodeError:
+                print(f"Error loading metrics_eval.json for step {eval_step}")
+                return
+
+        if "log_keys" in metrics_eval:
+            self.log_keys = metrics_eval["log_keys"]
+        else:
+            self.log_keys = None
+
+        # Add eval_step if not already present
+        if "eval_step" not in metrics_eval:
+            metrics_eval["eval_step"] = eval_step
+
+        metrics_eval["eval/all_metrics_dict"][
+            "sampling_prob"
+        ] = self.env._motion_lib._sampling_prob.cpu().numpy()
+        metrics_eval["eval/failed_metrics_dict"][
+            "sampling_prob"
+        ] = self.env._motion_lib._sampling_prob.cpu().numpy()[metrics_eval["failed_idxes"]]
+        metrics_eval["eval/all_metrics_dict"] = create_html_table(
+            metrics_eval["eval/all_metrics_dict"]
+        )
+        metrics_eval["eval/failed_metrics_dict"] = create_html_table(
+            metrics_eval["eval/failed_metrics_dict"]
+        )
+
+        # Log to wandb
+        if self.accelerator.is_main_process and wandb_run_exists():
+            if self.log_keys is not None:
+                metrics_eval = {f"{self.log_keys}/{k}": v for k, v in metrics_eval.items()}
+            wandb.log(metrics_eval)
+
+        for key in ["failed_keys", "failed_idxes"]:
+            # ZL: why do we need to do this?
+            if key in metrics_eval:
+                del metrics_eval[key]
+
+        print(f"Logged evaluation metrics for step {eval_step} {eval_step_dir}")
+
+    def _check_and_log_eval_render_results(self, eval_step, eval_step_dir):
+        """Check for new evaluation results and log them to wandb."""
+
+        video_dir = eval_step_dir / "render_results"
+        metrics_file = eval_step_dir / "metrics_eval.json"
+
+        if not video_dir.exists():
+            print(f"No render_results directory for step {eval_step}, skipping render logging")
+            return
+
+        # Read and log the metrics
+        with open(metrics_file) as f:
+            try:
+                metrics_eval = json.load(f)
+            except json.JSONDecodeError:
+                print(f"Error loading metrics_eval.json for step {eval_step}")
+                return
+
+        if "log_keys" in metrics_eval:
+            self.log_keys = metrics_eval["log_keys"]
+        else:
+            self.log_keys = None
+
+        video_files = []
+        for i, video_file in enumerate(sorted(video_dir.iterdir())):
+            if video_file.is_file() and video_file.name.endswith(".mp4"):
+                video_files.append((i, video_file))
+        if self.log_keys is not None:
+            wandb_videos = {
+                f"videos_hard_{self.log_keys}/{i: 04d}": wandb.Video(str(video_file), format="mp4")
+                for i, video_file in reversed(video_files)
+            }
+        else:
+            wandb_videos = {
+                f"videos_hard/{i: 04d}": wandb.Video(str(video_file), format="mp4")
+                for i, video_file in reversed(video_files)
+            }
+        wandb_videos["eval_step"] = eval_step
+
+        # Log to wandb
+        if self.accelerator.is_main_process and wandb_run_exists():
+            wandb.log(wandb_videos)
+
+        print(f"Logged rendered video for step {eval_step}")
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/wandb_callback.py b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/wandb_callback.py
new file mode 100644
index 0000000000000000000000000000000000000000..28e74948c8d7ff71c7e5bd098ee767d9140b6516
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/callbacks/wandb_callback.py
@@ -0,0 +1,19 @@
+from transformers import TrainerCallback
+import wandb
+
+from gear_sonic.trl.utils.common import wandb_run_exists
+
+
+class WandbCallback(TrainerCallback):
+    """Callback to save model state_dict during training."""
+
+    def __init__(
+        self,
+    ):
+        super().__init__()
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+
+        if state.is_world_process_zero and wandb_run_exists():
+            logs["global_step"] = state.global_step
+            wandb.log(logs, step=state.global_step)
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/losses/__init__.py b/GR00T-WholeBodyControl/gear_sonic/trl/losses/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/losses/__pycache__/__init__.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/losses/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c14c4d5fcbd5520f04b2e610b667694b445bcb1b
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/losses/__pycache__/__init__.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/losses/__pycache__/__init__.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/losses/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d4a9f3bc3df793935eaa1e40c909ffeed4f1874b
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/losses/__pycache__/__init__.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/losses/__pycache__/token_losses.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/losses/__pycache__/token_losses.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f7ce6292c13e7fd14e40d158574f346ac794b580
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/losses/__pycache__/token_losses.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/losses/__pycache__/token_losses.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/losses/__pycache__/token_losses.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..973fa7b329e8629851f573a348c40b05d1a3bea2
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/losses/__pycache__/token_losses.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/losses/token_losses.py b/GR00T-WholeBodyControl/gear_sonic/trl/losses/token_losses.py
new file mode 100644
index 0000000000000000000000000000000000000000..2c63664ab9055383ab67bba12e46933dae4a5bec
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/losses/token_losses.py
@@ -0,0 +1,1646 @@
+"""Loss functions for token-based models."""
+
+from pathlib import Path
+
+import omegaconf
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from gear_sonic.isaac_utils import rotations
+from gear_sonic.trl.utils import order_converter, torch_transform
+from gear_sonic.utils import batch_normalizer
+from gear_sonic.utils.motion_lib import torch_humanoid_batch
+
+
+def create_humanoid(
+    skeleton_name: str, device: torch.device = None
+) -> torch_humanoid_batch.Humanoid_Batch:
+    """Create a Humanoid_Batch for FK computations.
+
+    Args:
+        skeleton_name: Name of the skeleton config file (without .yaml extension)
+                       e.g., "motion_g1_extended_toe" for the extended G1 skeleton
+        device: Device to place the humanoid on
+
+    Returns:
+        Humanoid_Batch instance
+    """
+    if device is None:
+        device = torch.device("cpu")
+
+    groot_root = Path(__file__).parent.parent.parent.parent
+    motion_yaml = (
+        groot_root
+        / "rl"
+        / "config"
+        / "manager_env"
+        / "commands"
+        / "terms"
+        / f"{skeleton_name}.yaml"
+    )
+
+    cfg = omegaconf.OmegaConf.load(motion_yaml).motion.motion_lib_cfg
+    return torch_humanoid_batch.Humanoid_Batch(cfg, device=device)
+
+
+def decoder_output_to_egocentric_transforms(
+    decoder_output: dict,
+    decoder_cfg: dict,
+    humanoid: torch_humanoid_batch.Humanoid_Batch,
+    dof_converter: order_converter.G1Converter = None,
+    include_extended: bool = False,
+):
+    """Convert decoder output to joint positions and 6D rotations (ortho6d).
+
+    Always returns 6D rotations. For geodesic loss, convert to matrices in the loss function.
+
+    Args:
+        decoder_output: Dict with decoder output tensors
+        decoder_cfg: Decoder config with 'outputs' key
+        humanoid: Humanoid_Batch instance (already on correct device)
+        dof_converter: DOF order converter (if None, assumes qpos is in MuJoCo order)
+        include_extended: If True, compute and append extended body transforms (e.g., head, toes)
+
+    Returns:
+        egocentric_pos: Joint positions [..., num_bodies, 3]
+        egocentric_rot_6d: Joint rotations in 6D format [..., num_bodies, 6]
+    """
+    output_keys = list(decoder_cfg["outputs"])
+
+    if set(output_keys) == {"command_multi_future_nonflat", "motion_anchor_ori_b_mf_nonflat"}:
+        """NOTE: there's a bug in command_multi_future_nonflat, where the temporal axis is incorrectly flattened.
+        command_multi_future_nonflat: [..., num_future, 58] = [dof_pos(29), dof_vel(29)] in IsaacLab order
+        motion_anchor_ori_b_mf_nonflat: [..., num_future, 6] = relative 6D orientation
+        6D format: rot_mat[..., :2].reshape(6) = [R00, R01, R10, R11, R20, R21] (row-major of first 2 cols)
+
+        Since we don't use the raltive motion anchor, we are considering the egocentric transforms.
+        """
+        orig_shape = decoder_output["command_multi_future_nonflat"].shape[:-1]  # [..., num_future]
+        num_timesteps = orig_shape[-1]
+
+        # The below processing to obtain dof_qpos is needed because of the observation bug
+        dof_pos = decoder_output["command_multi_future_nonflat"][
+            ..., : num_timesteps // 2, : humanoid.num_dof * 2
+        ]
+        dof_pos = dof_pos.reshape(-1, humanoid.num_dof)
+
+        root_quat = torch.tensor([1.0, 0.0, 0.0, 0.0]).repeat(dof_pos.shape[0], 1).to(dof_pos)
+        root_trans = torch.zeros(dof_pos.shape[0], 3).to(dof_pos)
+        qpos = torch.cat([root_trans, root_quat, dof_pos], dim=-1)
+
+        # Convert from IsaacLab to MuJoCo DOF order if converter provided
+        if dof_converter is not None:
+            qpos = dof_converter.to_mujoco(qpos)
+
+        egocentric_pos, egocentric_rot = humanoid.qpos_to_global_transforms(
+            qpos, include_extended=include_extended
+        )
+
+        egocentric_pos = egocentric_pos.view(*orig_shape, egocentric_pos.shape[-2], 3)
+        egocentric_rot = egocentric_rot.view(*orig_shape, egocentric_rot.shape[-3], 3, 3)
+
+        # Convert 3x3 rotation matrix to 6D (first two columns flattened)
+        egocentric_rot_6d = rotations.mat_to_rot6d_first_two_cols(egocentric_rot)
+
+        return egocentric_pos, egocentric_rot_6d
+
+    elif set(output_keys) == {
+        "command_multi_future_egocentric_joint_transforms_nonflat",
+        "command_multi_future_root_transforms_nonflat",
+    }:
+        """
+        command_multi_future_egocentric_joint_transforms_nonflat: [..., num_future, num_bodies, 9]
+            - 9 = 3 (position) + 6 (6D rotation)
+            - Joint positions and rotations relative to each frame's projected root
+        command_multi_future_root_transforms_nonflat: [..., num_future, 9]
+            - 9 = 3 (position) + 6 (6D rotation)
+            - Root position and rotation relative to first reference frame
+        """
+        joint_transforms = decoder_output[
+            "command_multi_future_egocentric_joint_transforms_nonflat"
+        ]
+        joint_transforms = joint_transforms.reshape(*joint_transforms.shape[:-1], -1, 9)
+
+        egocentric_pos = joint_transforms[..., :3]  # [..., num_future, num_bodies, 3]
+        egocentric_rot_6d = joint_transforms[..., 3:]  # [..., num_future, num_bodies, 6]
+
+        # Convert from IsaacLab to MuJoCo DOF order if converter provided
+        if dof_converter is not None:
+            egocentric_pos = dof_converter.to_mujoco(egocentric_pos)
+            egocentric_rot_6d = dof_converter.to_mujoco(egocentric_rot_6d)
+
+        # Compute extended body transforms if requested
+        if include_extended and humanoid.num_bodies_augment > humanoid.num_bodies:
+            # Convert to matrices for FK computation (only for extended joints)
+            egocentric_rot_mat = rotations.rot6d_to_mat_first_two_cols(egocentric_rot_6d)
+            full_pos, full_rot_mat = humanoid.append_extended_transforms(
+                egocentric_pos, egocentric_rot_mat
+            )
+
+            # Extract only extended joints and convert to 6D
+            extended_pos = full_pos[..., humanoid.num_bodies :, :]
+            extended_rot_mat = full_rot_mat[..., humanoid.num_bodies :, :, :]
+            extended_rot_6d = rotations.mat_to_rot6d_first_two_cols(extended_rot_mat)
+
+            # Concatenate: keep original 6D (no normalization) + extended 6D
+            egocentric_pos = torch.cat([egocentric_pos, extended_pos], dim=-2)
+            egocentric_rot_6d = torch.cat([egocentric_rot_6d, extended_rot_6d], dim=-2)
+
+        return egocentric_pos, egocentric_rot_6d
+
+    raise NotImplementedError(f"Unsupported decoder output format: {output_keys}")
+
+
+def decoder_output_to_world_transforms(
+    decoder_output: dict,
+    decoder_cfg: dict,
+    humanoid: torch_humanoid_batch.Humanoid_Batch,
+    dof_converter: order_converter.G1Converter = None,
+    include_extended: bool = False,
+):
+    """Convert decoder output to joint positions and rotation matrices in a consistent
+    world coordinate frame (first reference frame's projected root, heading-aligned).
+
+    For formats with per-frame egocentric transforms and root transforms, this
+    reconstructs world transforms by applying each frame's root rotation/translation.
+
+    Coordinate system: first frame's heading-aligned, ground-projected root.
+
+    Args:
+        decoder_output: Dict with decoder output tensors
+        decoder_cfg: Decoder config with 'outputs' key
+        humanoid: Humanoid_Batch instance (already on correct device)
+        dof_converter: DOF order converter (if None, assumes qpos is in MuJoCo order)
+        include_extended: If True, compute and append extended body transforms
+
+    Returns:
+        world_pos: Joint positions [..., num_future, num_bodies, 3]
+        world_rot: Joint rotation matrices [..., num_future, num_bodies, 3, 3]
+    """  # noqa: D205
+    output_keys = list(decoder_cfg["outputs"])
+
+    if set(output_keys) == {"command_multi_future_nonflat", "motion_anchor_ori_b_mf_nonflat"}:
+        # qpos-based format with identity root — egocentric IS the world frame
+        egocentric_pos, egocentric_rot_6d = decoder_output_to_egocentric_transforms(
+            decoder_output,
+            decoder_cfg,
+            humanoid,
+            dof_converter,
+            include_extended=include_extended,
+        )
+        egocentric_rot_mat = rotations.rot6d_to_mat_first_two_cols(egocentric_rot_6d)
+        return egocentric_pos, egocentric_rot_mat
+
+    elif set(output_keys) == {
+        "command_multi_future_egocentric_joint_transforms_nonflat",
+        "command_multi_future_root_transforms_nonflat",
+    }:
+        # Get egocentric transforms and root transforms
+        egocentric_pos, egocentric_rot_6d = decoder_output_to_egocentric_transforms(
+            decoder_output,
+            decoder_cfg,
+            humanoid,
+            dof_converter,
+            include_extended=include_extended,
+        )
+        root_pos, root_rot_6d = decoder_output_to_root_transforms(
+            decoder_output,
+            decoder_cfg,
+        )
+
+        # root_rot = heading_0_inv * q_current_t  (full rotation relative to first frame heading)
+        # egocentric transforms are in each frame's heading-aligned frame
+        # To go from egocentric to world (first frame heading-aligned):
+        #   - Extract heading-only (yaw) from root_rot for position/rotation transform
+        #   - ego_pos is already in heading frame, so only yaw rotation is needed
+        root_rot_mat = rotations.rot6d_to_mat_first_two_cols(root_rot_6d)  # [..., num_future, 3, 3]
+        ego_rot_mat = rotations.rot6d_to_mat_first_two_cols(
+            egocentric_rot_6d
+        )  # [..., num_future, num_bodies, 3, 3]
+
+        # Variable-frame masking can zero out padded root rotations.
+        # matrix_to_quaternion assumes valid SO(3); sanitize before heading extraction.
+        root_rot_mat = _sanitize_rotation_matrices(root_rot_mat)
+
+        # Extract heading (yaw-only) rotation from root_rot_mat
+        root_quat = rotations.matrix_to_quaternion(root_rot_mat)  # [..., num_future, 4] (wxyz)
+        heading_quat = torch_transform.get_heading_q(
+            root_quat
+        )  # [..., num_future, 4] (wxyz, yaw only)
+        heading_rot_mat = rotations.quaternion_to_matrix(heading_quat)  # [..., num_future, 3, 3]
+
+        # Positions: [..., num_future, num_bodies, 3]
+        # Zero out height (z) from root_pos — egocentric positions already encode
+        # height relative to ground; adding root height would double-count it.
+        root_pos_xy = root_pos.clone()
+        root_pos_xy[..., 2] = 0.0
+        world_pos = torch.matmul(
+            heading_rot_mat.unsqueeze(-3), egocentric_pos.unsqueeze(-1)
+        ).squeeze(-1) + root_pos_xy.unsqueeze(-2)
+
+        # Rotations: [..., num_future, num_bodies, 3, 3]
+        world_rot = torch.matmul(heading_rot_mat.unsqueeze(-3), ego_rot_mat)
+
+        return world_pos, world_rot
+
+    raise NotImplementedError(f"Unsupported decoder output format: {output_keys}")
+
+
+def decoder_output_to_root_transforms(
+    decoder_output: dict,
+    decoder_cfg: dict,
+):
+    """Extract root position and rotation (as 6D) from decoder output.
+
+    Always returns 6D rotations. For geodesic loss, convert to matrices in the loss function.
+
+    Args:
+        decoder_output: Dict with decoder output tensors
+        decoder_cfg: Decoder config with 'outputs' key
+
+    Returns:
+        root_pos: Root position [..., num_future, 3] (zeros if not available in format)
+        root_rot_6d: Root rotation in 6D format [..., num_future, 6]
+    """
+    output_keys = list(decoder_cfg["outputs"])
+
+    if set(output_keys) == {"command_multi_future_nonflat", "motion_anchor_ori_b_mf_nonflat"}:
+        # Root rotation from motion_anchor_ori_b_mf_nonflat (already 6D)
+        root_rot_6d = decoder_output["motion_anchor_ori_b_mf_nonflat"]  # [..., num_future, 6]
+
+        # Root position not directly available - return zeros with matching shape
+        root_pos = torch.zeros(
+            *root_rot_6d.shape[:-1], 3, device=root_rot_6d.device, dtype=root_rot_6d.dtype
+        )
+
+        return root_pos, root_rot_6d
+
+    elif set(output_keys) == {
+        "command_multi_future_egocentric_joint_transforms_nonflat",
+        "command_multi_future_root_transforms_nonflat",
+    }:
+        # Root transforms are in command_multi_future_root_transforms_nonflat
+        root_transforms = decoder_output[
+            "command_multi_future_root_transforms_nonflat"
+        ]  # [..., num_future, 9]
+
+        root_pos = root_transforms[..., :3]  # [..., num_future, 3]
+        root_rot_6d = root_transforms[..., 3:]  # [..., num_future, 6]
+
+        return root_pos, root_rot_6d
+
+    raise NotImplementedError(f"Unsupported decoder output format: {output_keys}")
+
+
+def _extract_dof_pos(
+    decoder_output: dict,
+    decoder_cfg: dict,
+    humanoid: torch_humanoid_batch.Humanoid_Batch,
+    dof_converter: order_converter.G1Converter = None,
+    egocentric_pos: torch.Tensor = None,
+    egocentric_rot_6d: torch.Tensor = None,
+):
+    """Extract DOF angles from decoder output via inverse kinematics.
+
+    Runs ``humanoid.global_transforms_to_qpos`` on egocentric transforms to
+    recover DOF angles.  If ``egocentric_pos`` and ``egocentric_rot_6d`` are
+    provided, uses them directly (avoids re-parsing decoder output).  Otherwise,
+    calls ``decoder_output_to_egocentric_transforms`` to obtain them.
+
+    Only supports the ``command_multi_future_egocentric_joint_transforms_nonflat``
+    format; returns ``None`` for other formats.
+
+    Returns:
+        dof_pos: [..., num_future, num_dof] or None if format unsupported.
+    """
+    if egocentric_pos is None or egocentric_rot_6d is None:
+        output_keys = list(decoder_cfg["outputs"])
+        if set(output_keys) != {
+            "command_multi_future_egocentric_joint_transforms_nonflat",
+            "command_multi_future_root_transforms_nonflat",
+        }:
+            return None
+
+        # Parse decoder output (no extended bodies — they have no qpos)
+        egocentric_pos, egocentric_rot_6d = decoder_output_to_egocentric_transforms(
+            decoder_output,
+            decoder_cfg,
+            humanoid,
+            dof_converter,
+            include_extended=False,
+        )
+
+    # Convert 6D -> 3x3 rotation matrices for IK
+    rot_mat = rotations.rot6d_to_mat_first_two_cols(egocentric_rot_6d)  # [..., F, J, 3, 3]
+
+    # Variable-frame masking can zero out padded frames — sanitize for IK.
+    rot_mat = _sanitize_rotation_matrices(rot_mat)
+
+    # Flatten leading dims for global_transforms_to_qpos which expects [B, T, J, 3, 3]
+    num_frames = rot_mat.shape[-4]
+    num_joints = rot_mat.shape[-3]
+    lead_shape = rot_mat.shape[:-4]
+
+    rot_mat_5d = rot_mat.reshape(-1, num_frames, num_joints, 3, 3)
+    pos_4d = egocentric_pos.reshape(-1, num_frames, num_joints, 3)
+
+    # IK: egocentric transforms -> qpos (heading cancels in local rotation computation)
+    qpos = humanoid.global_transforms_to_qpos(rot_mat_5d, pos_4d)  # [flat_B, F, D]
+    dof_pos = qpos[..., 7:]  # [flat_B, F, num_dof]
+
+    return dof_pos.reshape(*lead_shape, num_frames, humanoid.num_dof)
+
+
+def _sanitize_rotation_matrices(rot_mat: torch.Tensor) -> torch.Tensor:
+    """Replace degenerate (near-zero) rotation matrices with identity.
+
+    Variable-frame masking can zero out padded frames, producing zero 3x3
+    matrices.  Downstream operations (matrix_to_quaternion, IK) assume valid
+    SO(3) and produce NaN on such inputs.
+    """
+    norm_sq = (rot_mat * rot_mat).sum(dim=(-2, -1))
+    eye = torch.eye(3, device=rot_mat.device, dtype=rot_mat.dtype)
+    return torch.where((norm_sq < 1e-6)[..., None, None], eye, rot_mat)
+
+
+def _apply_normalizer(normalizer, gt, pred, mask=None):
+    """Update normalizer from gt, normalize both gt and pred.
+
+    Flattens trailing dims to match normalizer's expected feature dim,
+    updates running stats from valid gt samples, normalizes both tensors.
+
+    Returns:
+        (gt_normed, pred_normed) with same shapes as inputs.
+    """
+    orig_shape = gt.shape
+    flat_dim = normalizer.num_features
+    gt_flat = gt.reshape(-1, flat_dim)
+    pred_flat = pred.reshape(-1, flat_dim)
+    if mask is not None:
+        # Broadcast mask to match gt shape (mask may have fewer trailing dims)
+        m = mask
+        while m.ndim < gt.ndim:
+            m = m.unsqueeze(-1)
+        valid_mask = m.expand_as(gt).reshape(-1, flat_dim)[:, 0].bool()
+        if valid_mask.any():
+            normalizer.update(gt_flat[valid_mask])
+    else:
+        normalizer.update(gt_flat)
+    return (
+        normalizer.normalize(gt_flat).reshape(orig_shape),
+        normalizer.normalize(pred_flat).reshape(orig_shape),
+    )
+
+
+# =============================================================================
+# Helper Functions
+# =============================================================================
+
+
+def compute_loss(pred: torch.Tensor, target: torch.Tensor, loss_type: str) -> torch.Tensor:
+    """Compute loss between prediction and target using specified loss type.
+
+    Args:
+        pred: Predicted tensor
+        target: Target tensor
+        loss_type: One of "mse", "l1", "huber", "cosine"
+
+    Returns:
+        Scalar loss tensor
+    """
+    if loss_type == "mse":
+        return F.mse_loss(pred, target)
+    elif loss_type == "l1":
+        return F.l1_loss(pred, target)
+    elif loss_type == "huber":
+        return F.huber_loss(pred, target)
+    elif loss_type == "cosine":
+        cosine_sim = F.cosine_similarity(pred, target, dim=-1)
+        return (1 - cosine_sim).mean()
+    else:
+        raise ValueError(f"Unknown loss_type: {loss_type}")
+
+
+def zero_loss(device: torch.device) -> torch.Tensor:
+    """Return a zero loss tensor with requires_grad=True."""
+    return torch.tensor(0.0, device=device, requires_grad=True)
+
+
+def _get_device_from_loss_inputs(loss_inputs: dict) -> torch.device:
+    """Get device from loss_inputs, handling kinematic-only mode where action_mean is None."""
+    if loss_inputs.get("action_mean") is not None:
+        return loss_inputs["action_mean"].device
+    for val in loss_inputs.get("tokenizer_obs", {}).values():
+        if isinstance(val, torch.Tensor):
+            return val.device
+    return torch.device("cpu")
+
+
+def _build_frame_mask_for_loss(frame_mask, num_frames_in_tensor):
+    """Build a frame mask for a loss tensor.
+
+    Args:
+        frame_mask: [..., max_frames] bool (from loss_inputs), or None
+        num_frames_in_tensor: actual frame count in the loss tensor (may be < max_frames)
+
+    Returns:
+        float mask [..., num_frames], or None.  Consumers (_compute_masked_loss,
+        _masked_geodesic_angle) auto-broadcast to the target shape.
+    """
+    if frame_mask is None:
+        return None
+    return frame_mask[..., :num_frames_in_tensor].float()
+
+
+def _masked_geodesic_angle(pred_rot, gt_rot, mask=None, dt=1.0, eps=1e-6):  # noqa: D417
+    """Compute geodesic angle loss on rotation matrices with optional masking.
+
+    Args:
+        pred_rot, gt_rot: [..., 3, 3] rotation matrices
+        mask: float mask (1=valid) with frame dim, auto-broadcast to angles shape, or None
+        dt: time step divisor (for velocity normalization)
+        eps: clamp epsilon for numerical stability
+
+    Returns:
+        scalar loss
+    """
+    R_diff = torch.matmul(gt_rot.transpose(-1, -2), pred_rot)
+    trace = R_diff[..., 0, 0] + R_diff[..., 1, 1] + R_diff[..., 2, 2]
+    cos_angle = torch.clamp((trace - 1) / 2, -1.0 + eps, 1.0 - eps)
+    angles = torch.acos(cos_angle) / dt  # [..., (B)]
+    if mask is not None:
+        # Auto-expand mask [..., F] to match angles [..., F, (B)]
+        while mask.dim() < angles.dim():
+            mask = mask.unsqueeze(-1)
+        return (angles * mask).sum() / mask.expand_as(angles).sum().clamp(min=1)
+    return angles.mean()
+
+
+def _build_vel_frame_mask(frame_mask, num_frames_in_tensor):
+    """Build a frame mask for velocity (finite-difference) tensors.
+
+    Velocity at frame i uses frames i and i+1, so frame i is valid only if
+    both frame i and frame i+1 are valid.
+
+    Args:
+        frame_mask: [..., max_frames] bool, or None
+        num_frames_in_tensor: number of frames in the position tensor (vel has num-1)
+
+    Returns:
+        float mask [..., num_frames-1], or None.  Consumers auto-broadcast.
+    """
+    if frame_mask is None:
+        return None
+    mask = frame_mask[..., :num_frames_in_tensor]
+    return (mask[..., :-1] & mask[..., 1:]).float()
+
+
+def _compute_masked_loss(pred, gt, mask, loss_type="mse"):
+    """Compute loss with optional mask, supporting all loss types.
+
+    Uses torch loss functions with reduction='none' so masking works uniformly.
+    mask is auto-broadcast to match the element-wise loss shape.
+    """
+    if loss_type == "mse":
+        elem = F.mse_loss(pred, gt, reduction="none")
+    elif loss_type == "l1":
+        elem = F.l1_loss(pred, gt, reduction="none")
+    elif loss_type == "huber":
+        elem = F.huber_loss(pred, gt, reduction="none")
+    else:
+        raise ValueError(f"Unknown loss_type: {loss_type}")
+    if mask is None:
+        return elem.mean()
+    # Auto-expand mask [..., F] to match elem [..., F, ...]
+    while mask.dim() < elem.dim():
+        mask = mask.unsqueeze(-1)
+    masked = elem * mask
+    num_valid = mask.expand_as(elem).sum().clamp(min=1)
+    return masked.sum() / num_valid
+
+
+# =============================================================================
+# Reconstruction Loss
+# =============================================================================
+
+
+class G1ReconLoss(nn.Module):
+
+    def __init__(self, loss_type="mse", **kwargs):  # noqa: ARG002
+        super().__init__()
+        self.loss_type = loss_type
+
+    def forward(self, loss_inputs):
+        tokenizer_obs = loss_inputs["tokenizer_obs"]
+        decoders_cfg = loss_inputs["decoders_cfg"]
+        decoded_outputs = loss_inputs["decoded_outputs"]
+        frame_mask = loss_inputs.get("frame_mask", None)
+
+        g1_motion_output = torch.cat(
+            [tokenizer_obs[key] for key in decoders_cfg["g1_kin"]["outputs"]], dim=-1
+        )
+        g1_motion_output_pred = torch.cat(
+            [decoded_outputs["g1_kin"][key] for key in decoded_outputs["g1_kin"]], dim=-1
+        )
+
+        # Align temporal dim: truncate the longer to the shorter
+        if g1_motion_output_pred.shape[-2] < g1_motion_output.shape[-2]:
+            g1_motion_output = g1_motion_output[..., : g1_motion_output_pred.shape[-2], :]
+        elif g1_motion_output_pred.shape[-2] > g1_motion_output.shape[-2]:
+            g1_motion_output_pred = g1_motion_output_pred[..., : g1_motion_output.shape[-2], :]
+
+        # Build frame mask: shape [..., num_frames, 1] for broadcasting
+        num_frames = g1_motion_output.shape[-2]
+        mask = _build_frame_mask_for_loss(frame_mask, num_frames)
+
+        return _compute_masked_loss(g1_motion_output_pred, g1_motion_output, mask, self.loss_type)
+
+
+class G1ReconLossAligned(G1ReconLoss):
+    """Same as G1ReconLoss but uses a single key list for target and pred so that
+    encoder input = g1_kin output = loss. Expects loss_inputs["recon_target_keys"].
+    """  # noqa: D205
+
+    def forward(self, loss_inputs):
+        tokenizer_obs = loss_inputs["tokenizer_obs"]
+        decoded_outputs = loss_inputs["decoded_outputs"]
+        frame_mask = loss_inputs.get("frame_mask", None)
+        keys = list(loss_inputs["recon_target_keys"])
+        g1_motion_output = torch.cat([tokenizer_obs[k] for k in keys], dim=-1)
+        g1_motion_output_pred = torch.cat([decoded_outputs["g1_kin"][k] for k in keys], dim=-1)
+
+        if g1_motion_output_pred.shape[-2] < g1_motion_output.shape[-2]:
+            g1_motion_output = g1_motion_output[..., : g1_motion_output_pred.shape[-2], :]
+        elif g1_motion_output_pred.shape[-2] > g1_motion_output.shape[-2]:
+            g1_motion_output_pred = g1_motion_output_pred[..., : g1_motion_output.shape[-2], :]
+
+        num_frames = g1_motion_output.shape[-2]
+        mask = _build_frame_mask_for_loss(frame_mask, num_frames)
+
+        return _compute_masked_loss(g1_motion_output_pred, g1_motion_output, mask, self.loss_type)
+
+
+# =============================================================================
+# G1-SMPL Latent Alignment Loss
+# =============================================================================
+
+
+class G1SmplLatentLoss(nn.Module):
+    """Loss that compares the encoded latents between g1 and smpl encoders.
+    This encourages the latent representations to be similar across different
+    motion representation formats.
+    """  # noqa: D205
+
+    def __init__(self, loss_type="mse", **kwargs):  # noqa: ARG002
+        super().__init__()
+        self.loss_type = loss_type
+
+    def forward(self, loss_inputs):
+        encoded_latents = loss_inputs["encoded_latents"]
+        encoder_masks = loss_inputs["encoder_masks"]
+
+        # Get g1 latents that have corresponding smpl data
+        g1_latents = encoded_latents["g1"]
+        smpl_latents = encoded_latents["smpl"]
+
+        # Extract only the g1 samples that have corresponding smpl
+        g1_latents_matched = g1_latents[encoder_masks["g1_has_smpl"]]
+        if g1_latents_matched.shape[0] == 0:
+            return torch.tensor(0.0, device=g1_latents.device)
+
+        # Compute loss based on loss_type
+        if self.loss_type == "mse":
+            loss = F.mse_loss(g1_latents_matched, smpl_latents)
+        elif self.loss_type == "l1":
+            loss = F.l1_loss(g1_latents_matched, smpl_latents)
+        elif self.loss_type == "huber":
+            loss = F.huber_loss(g1_latents_matched, smpl_latents)
+        elif self.loss_type == "cosine":
+            # Cosine distance: 1 - cosine_similarity
+            cosine_sim = F.cosine_similarity(g1_latents_matched, smpl_latents, dim=-1)
+            loss = (1 - cosine_sim).mean()
+        else:
+            raise ValueError(f"Unknown loss_type: {self.loss_type}")
+
+        return loss
+
+
+class TeleopSmplLatentLoss(nn.Module):
+    """Loss that compares the encoded latents between teleop and smpl encoders.
+    This encourages the latent representations to be similar across different
+    motion representation formats.
+    """  # noqa: D205
+
+    def __init__(self, loss_type="mse", **kwargs):  # noqa: ARG002
+        super().__init__()
+        self.loss_type = loss_type
+
+    def forward(self, loss_inputs):
+        encoded_latents = loss_inputs["encoded_latents"]
+        encoder_masks = loss_inputs["encoder_masks"]
+
+        # Get teleop and smpl latents
+        tlp_latents = encoded_latents["teleop"]
+        smpl_latents = encoded_latents["smpl"]
+
+        # teleop_has_smpl selects teleop samples that also have smpl active
+        tlp_latents_matched = tlp_latents[encoder_masks["teleop_has_smpl"]]
+        # smpl_has_teleop selects smpl samples that also have teleop active
+        smpl_latents_matched = smpl_latents[encoder_masks["smpl_has_teleop"]]
+
+        # Return 0 loss if no matching samples
+        if tlp_latents_matched.shape[0] == 0:
+            return torch.tensor(0.0, device=tlp_latents.device)
+
+        # Compute loss based on loss_type
+        if self.loss_type == "mse":
+            loss = F.mse_loss(tlp_latents_matched, smpl_latents_matched)
+        elif self.loss_type == "l1":
+            loss = F.l1_loss(tlp_latents_matched, smpl_latents_matched)
+        elif self.loss_type == "huber":
+            loss = F.huber_loss(tlp_latents_matched, smpl_latents_matched)
+        elif self.loss_type == "cosine":
+            # Cosine distance: 1 - cosine_similarity
+            cosine_sim = F.cosine_similarity(tlp_latents_matched, smpl_latents_matched, dim=-1)
+            loss = (1 - cosine_sim).mean()
+        else:
+            raise ValueError(f"Unknown loss_type: {self.loss_type}")
+
+        return loss
+
+
+# =============================================================================
+# G1-Teleop Latent Alignment Loss
+# =============================================================================
+
+
+class G1TeleopLatentLoss(nn.Module):
+    """Loss that compares the encoded latents between g1 and teleop encoders.
+    This encourages the latent representations to be similar across different
+    motion representation formats.
+    """  # noqa: D205
+
+    def __init__(self, loss_type="mse", **kwargs):  # noqa: ARG002
+        super().__init__()
+        self.loss_type = loss_type
+
+    def forward(self, loss_inputs):
+        encoded_latents = loss_inputs["encoded_latents"]
+        encoder_masks = loss_inputs["encoder_masks"]
+
+        # Get g1 and teleop latents
+        g1_latents = encoded_latents["g1"]
+        tlp_latents = encoded_latents["teleop"]
+
+        # g1_has_teleop selects g1 samples that also have teleop active
+        g1_latents_matched = g1_latents[encoder_masks["g1_has_teleop"]]
+        # teleop_has_g1 selects teleop samples that also have g1 active
+        tlp_latents_matched = tlp_latents[encoder_masks["teleop_has_g1"]]
+
+        # Return 0 loss if no matching samples
+        if g1_latents_matched.shape[0] == 0:
+            return torch.tensor(0.0, device=g1_latents.device)
+
+        # Compute loss based on loss_type
+        if self.loss_type == "mse":
+            loss = F.mse_loss(g1_latents_matched, tlp_latents_matched)
+        elif self.loss_type == "l1":
+            loss = F.l1_loss(g1_latents_matched, tlp_latents_matched)
+        elif self.loss_type == "huber":
+            loss = F.huber_loss(g1_latents_matched, tlp_latents_matched)
+        elif self.loss_type == "cosine":
+            # Cosine distance: 1 - cosine_similarity
+            cosine_sim = F.cosine_similarity(g1_latents_matched, tlp_latents_matched, dim=-1)
+            loss = (1 - cosine_sim).mean()
+        else:
+            raise ValueError(f"Unknown loss_type: {self.loss_type}")
+
+        return loss
+
+
+# =============================================================================
+# SOMA Latent Consistency Loss
+# =============================================================================
+
+
+class G1SomaLatentLoss(nn.Module):
+    """Loss that compares the encoded latents between g1 and soma encoders.
+    This encourages the latent representations to be similar across different
+    motion representation formats.
+    """  # noqa: D205
+
+    def __init__(self, loss_type="mse", **kwargs):  # noqa: ARG002
+        super().__init__()
+        self.loss_type = loss_type
+
+    def forward(self, loss_inputs):
+        encoded_latents = loss_inputs["encoded_latents"]
+        encoder_masks = loss_inputs["encoder_masks"]
+
+        # Get g1 latents that have corresponding soma data
+        g1_latents = encoded_latents["g1"]
+        soma_latents = encoded_latents["soma"]
+
+        # Extract only the g1 samples that have corresponding soma
+        g1_latents_matched = g1_latents[encoder_masks["g1_has_soma"]]
+        if g1_latents_matched.shape[0] == 0:
+            return torch.tensor(0.0, device=g1_latents.device)
+
+        assert g1_latents_matched.shape == soma_latents.shape, (
+            f"Shape mismatch: g1_matched={g1_latents_matched.shape}, soma={soma_latents.shape}. "
+            "Ensure g1 encoder is co-activated when soma is sampled."
+        )
+
+        # Compute loss based on loss_type
+        if self.loss_type == "mse":
+            loss = F.mse_loss(g1_latents_matched, soma_latents)
+        elif self.loss_type == "l1":
+            loss = F.l1_loss(g1_latents_matched, soma_latents)
+        elif self.loss_type == "huber":
+            loss = F.huber_loss(g1_latents_matched, soma_latents)
+        elif self.loss_type == "cosine":
+            cosine_sim = F.cosine_similarity(g1_latents_matched, soma_latents, dim=-1)
+            loss = (1 - cosine_sim).mean()
+        else:
+            raise ValueError(f"Unknown loss_type: {self.loss_type}")
+
+        return loss
+
+
+# =============================================================================
+# Cycle Consistency Loss
+# =============================================================================
+
+
+class ReencodedSmplG1LatentLoss(nn.Module):
+    """Loss that compares the reencoded g1 latents (from smpl-to-g1 reconstruction)
+    with the original g1 latents. This encourages the reconstructed g1 motion
+    to be encodable back to the same latent space as the original g1 motion.
+    """  # noqa: D205
+
+    def __init__(self, loss_type="mse", **kwargs):  # noqa: ARG002
+        super().__init__()
+        self.loss_type = loss_type
+
+    def forward(self, loss_inputs):
+        reencoded_smpl_g1_latents = loss_inputs["reencoded_smpl_g1_latents"]
+        encoded_latents = loss_inputs["encoded_latents"]
+        encoder_masks = loss_inputs["encoder_masks"]
+
+        # Get g1 latents that have corresponding smpl data
+        g1_latents = encoded_latents["g1"]
+        g1_latents_matched = g1_latents[encoder_masks["g1_has_smpl"]]
+        if g1_latents_matched.shape[0] == 0:
+            return torch.tensor(0.0, device=g1_latents.device)
+
+        # Compute loss based on loss_type
+        if self.loss_type == "mse":
+            loss = F.mse_loss(reencoded_smpl_g1_latents, g1_latents_matched)
+        elif self.loss_type == "l1":
+            loss = F.l1_loss(reencoded_smpl_g1_latents, g1_latents_matched)
+        elif self.loss_type == "huber":
+            loss = F.huber_loss(reencoded_smpl_g1_latents, g1_latents_matched)
+        elif self.loss_type == "cosine":
+            # Cosine distance: 1 - cosine_similarity
+            cosine_sim = F.cosine_similarity(reencoded_smpl_g1_latents, g1_latents_matched, dim=-1)
+            loss = (1 - cosine_sim).mean()
+        else:
+            raise ValueError(f"Unknown loss_type: {self.loss_type}")
+
+        return loss
+
+
+# =============================================================================
+# Compliance-Aware Latent Alignment Losses
+# =============================================================================
+
+
+class G1SmplComplianceLatentLoss(nn.Module):
+    """G1→SMPL latent alignment loss with compliance-aware filtering.
+
+    This loss ONLY applies when compliance ≈ 0 (stiff mode).
+
+    Rationale:
+    - G1 encoder encodes pure kinematics (no compliance input)
+    - SMPL encoder receives compliance as input
+    - Only in stiff mode (compliance=0) should both produce similar latents
+    - In compliant mode, latents may legitimately diverge
+
+    Uses paired_g1_smpl_latents which ensures both latents are for the SAME environments.
+
+    Note: G1 latents are pre-detached in UniversalTokenModule for memory optimization.
+    The detach_g1_target parameter is kept for API compatibility but has no effect
+    when using the default module configuration.
+    """
+
+    def __init__(
+        self,
+        loss_type: str = "mse",
+        compliance_threshold: float = 0.01,
+        detach_g1_target: bool = True,  # Note: pre-detached in module
+        debug: bool = False,
+        debug_print_every: int = 500,
+        **kwargs,  # noqa: ARG002
+    ):
+        super().__init__()
+        self.loss_type = loss_type
+        self.compliance_threshold = compliance_threshold
+        self.detach_g1_target = detach_g1_target
+        self.debug = debug
+        self.debug_print_every = debug_print_every
+        self._call_count = 0
+
+    def forward(self, loss_inputs: dict) -> torch.Tensor:
+        paired_g1_smpl_latents = loss_inputs.get("paired_g1_smpl_latents")
+
+        if paired_g1_smpl_latents is None:
+            return zero_loss(_get_device_from_loss_inputs(loss_inputs))
+
+        # Note: g1 latents are pre-detached in UniversalTokenModule
+        g1_latents = paired_g1_smpl_latents["g1"]
+        smpl_latents = paired_g1_smpl_latents["smpl"]
+
+        if g1_latents.shape[0] == 0 or smpl_latents.shape[0] == 0:
+            return zero_loss(g1_latents.device)
+
+        # Filter to stiff samples only (compliance ≈ 0)
+        compliance_values = paired_g1_smpl_latents.get("compliance")
+        if self.debug and compliance_values.max() > 0.001:
+            print(  # noqa: T201
+                f"Compliance values: {compliance_values.max()}, {compliance_values.min()}"
+            )  # noqa: T201
+
+        if compliance_values is not None:
+            # All 3 compliance dimensions must be near zero
+            is_stiff = (compliance_values.abs() <= self.compliance_threshold).all(dim=-1)
+
+            self._call_count += 1
+            if self.debug and self._call_count % self.debug_print_every == 1:
+                num_stiff = is_stiff.sum().item()
+                total = is_stiff.shape[0]
+                print(  # noqa: T201
+                    f"[G1SmplComplianceLatentLoss] Stiff samples: {num_stiff}/{total} "
+                    f"({100*num_stiff/max(1,total):.1f}%)"
+                )
+
+            g1_latents = g1_latents[is_stiff]
+            smpl_latents = smpl_latents[is_stiff]
+
+            if g1_latents.shape[0] == 0:
+                return zero_loss(_get_device_from_loss_inputs(loss_inputs))
+
+        return compute_loss(smpl_latents, g1_latents, self.loss_type)
+
+
+class TeleopSmplComplianceLatentLoss(nn.Module):
+    """Teleop→SMPL latent alignment for compliance-aware training.
+
+    Compares latents between teleop and smpl encoders
+    for the SAME motion under the SAME compliance level.
+
+    Key Design Points:
+    - Applied to ALL compliance values (enables compliant mode learning!)
+    - teleop is the BRIDGE (uses G1-space joints, closer to G1)
+    - smpl is further from G1, needs more learning
+    - detach_teleop_target=True: teleop is teacher, smpl learns
+
+    This loss enables smpl encoder to learn compliance-aware
+    behavior from the teleop bridge encoder.
+
+    Note: Teleop latents are pre-detached in UniversalTokenModule for memory optimization.
+    The detach_teleop_target parameter is kept for API compatibility.
+    """
+
+    def __init__(
+        self,
+        loss_type: str = "mse",
+        detach_teleop_target: bool = True,  # Note: pre-detached in module
+        debug: bool = False,
+        debug_print_every: int = 1000,
+        **kwargs,  # noqa: ARG002
+    ):
+        super().__init__()
+        self.loss_type = loss_type
+        self.detach_teleop_target = detach_teleop_target
+        self.debug = debug
+        self.debug_print_every = debug_print_every
+        self._call_count = 0
+
+    def forward(self, loss_inputs: dict) -> torch.Tensor:
+        # Use paired_compliance_latents computed for same motion + same compliance
+        paired_compliance_latents = loss_inputs.get("paired_compliance_latents")
+
+        if paired_compliance_latents is None:
+            return zero_loss(_get_device_from_loss_inputs(loss_inputs))
+
+        # Note: teleop latents are pre-detached in UniversalTokenModule
+        teleop_latents = paired_compliance_latents["teleop"]
+        smpl_latents = paired_compliance_latents["smpl"]
+
+        # Debug logging
+        self._call_count += 1
+        if self.debug and self._call_count % self.debug_print_every == 1:
+            self._print_debug_info(loss_inputs, teleop_latents, smpl_latents)
+
+        if teleop_latents.shape[0] == 0 or smpl_latents.shape[0] == 0:
+            return zero_loss(teleop_latents.device)
+
+        return compute_loss(smpl_latents, teleop_latents, self.loss_type)
+
+    def _print_debug_info(
+        self, loss_inputs: dict, teleop_latents: torch.Tensor, smpl_latents: torch.Tensor
+    ):
+        """Print debug information about paired latents."""
+        debug_info = loss_inputs.get("paired_compliance_debug_info")
+        if debug_info is None:
+            return
+
+        print("\n" + "=" * 80)  # noqa: T201
+        print(f"[TeleopSmplComplianceLatentLoss DEBUG] Call #{self._call_count}")  # noqa: T201
+        print("=" * 80)  # noqa: T201
+
+        total_envs = debug_info.get("total_envs", "unknown")
+        num_paired = debug_info["num_paired_envs"]
+        ratio = f"{num_paired}/{total_envs}" if total_envs != "unknown" else str(num_paired)
+        pct = (
+            f" ({100*num_paired/total_envs:.1f}%)"
+            if isinstance(total_envs, int) and total_envs > 0
+            else ""
+        )
+
+        print(f"  Paired envs: {ratio}{pct}")  # noqa: T201
+        print(f"  teleop_latents shape: {teleop_latents.shape}")  # noqa: T201
+        print(f"  smpl_latents shape: {smpl_latents.shape}")  # noqa: T201
+
+        if num_paired > 0:
+            print("\n  --- LATENT COMPARISON ---")  # noqa: T201
+            print(f"  teleop[0,:5]: {teleop_latents[0,:5].detach()}")  # noqa: T201
+            print(f"  smpl[0,:5]:   {smpl_latents[0,:5].detach()}")  # noqa: T201
+            l2_diff = ((teleop_latents - smpl_latents) ** 2).mean(dim=-1)
+            print(f"  L2 diff: {l2_diff[:min(4, l2_diff.shape[0])].detach()}")  # noqa: T201
+        else:
+            print("\n  WARNING: No paired envs!")  # noqa: T201
+            print("  Possible causes: no SMPL data or too few envs")  # noqa: T201
+
+        print("=" * 80 + "\n")  # noqa: T201
+
+
+class ReencodedSmplG1ComplianceLatentLoss(nn.Module):
+    """Cycle consistency loss for decoder regularization.
+
+    Re-encodes the decoded G1 motion back to latent space and compares
+    with the original G1 latent.
+
+    Compliance-Aware Design:
+    - compliance_aware=True (default): Only apply in stiff mode (compliance ≈ 0)
+      In stiff mode, decoder should produce pure kinematic motion that
+      can be re-encoded to the same latent space.
+    - compliance_aware=False: Apply to all samples
+
+    Note: Target G1 latents (original_g1_latents_for_reencode) are pre-detached
+    in UniversalTokenModule. The detach_target parameter is kept for API compatibility.
+    """
+
+    def __init__(
+        self,
+        loss_type: str = "mse",
+        detach_target: bool = True,  # Note: pre-detached in module
+        compliance_aware: bool = True,
+        compliance_threshold: float = 0.01,
+        debug: bool = False,
+        debug_print_every: int = 500,
+        **kwargs,  # noqa: ARG002
+    ):
+        super().__init__()
+        self.loss_type = loss_type
+        self.detach_target = detach_target
+        self.compliance_aware = compliance_aware
+        self.compliance_threshold = compliance_threshold
+        self.debug = debug
+        self.debug_print_every = debug_print_every
+        self._call_count = 0
+
+    def forward(self, loss_inputs: dict) -> torch.Tensor:
+        reencoded_latents = loss_inputs.get("reencoded_smpl_g1_latents")
+        # Note: g1 latents are pre-detached in UniversalTokenModule
+        g1_latents = loss_inputs.get("original_g1_latents_for_reencode")
+
+        if reencoded_latents is None or g1_latents is None:
+            return zero_loss(_get_device_from_loss_inputs(loss_inputs))
+
+        if reencoded_latents.shape[0] == 0 or g1_latents.shape[0] == 0:
+            return zero_loss(g1_latents.device)
+
+        # Compliance-aware filtering
+        if self.compliance_aware:
+            paired_g1_smpl_latents = loss_inputs.get("paired_g1_smpl_latents")
+            compliance_values = None
+            if paired_g1_smpl_latents is not None:
+                compliance_values = paired_g1_smpl_latents.get("compliance")
+
+            if compliance_values is not None:
+                is_stiff = (compliance_values.abs() < self.compliance_threshold).all(dim=-1)
+
+                self._call_count += 1
+                if self.debug and self._call_count % self.debug_print_every == 1:
+                    num_stiff = is_stiff.sum().item()
+                    total = is_stiff.shape[0]
+                    print(  # noqa: T201
+                        f"[ReencodedSmplG1LatentLoss] Stiff samples: {num_stiff}/{total} "
+                        f"({100*num_stiff/max(1,total):.1f}%)"
+                    )
+
+                reencoded_latents = reencoded_latents[is_stiff]
+                g1_latents = g1_latents[is_stiff]
+
+                if reencoded_latents.shape[0] == 0:
+                    return zero_loss(_get_device_from_loss_inputs(loss_inputs))
+
+        return compute_loss(reencoded_latents, g1_latents, self.loss_type)
+
+
+class LatentL2Loss(nn.Module):
+    """L2 regularization on the latent residual (first latent_dim dimensions of action_mean).
+
+    Encourages the policy to make small, focused adjustments to the pretrained ATM
+    rather than completely overriding its behavior.
+
+    loss = mean(latent_residual^2)
+    """
+
+    def __init__(self, latent_dim=64, **kwargs):  # noqa: ARG002
+        super().__init__()
+        self.latent_dim = latent_dim
+
+    def forward(self, loss_inputs):
+        action_mean = loss_inputs["action_mean"]
+        latent_residual = action_mean[..., : self.latent_dim]
+        return (latent_residual**2).mean()
+
+
+class LatentL1Loss(nn.Module):
+    """L1 regularization on the latent residual (sparsity-inducing).
+
+    Encourages sparse modifications to the pretrained ATM.
+    """
+
+    def __init__(self, latent_dim=64, **kwargs):  # noqa: ARG002
+        super().__init__()
+        self.latent_dim = latent_dim
+
+    def forward(self, loss_inputs):
+        action_mean = loss_inputs["action_mean"]
+        latent_residual = action_mean[..., : self.latent_dim]
+        return torch.abs(latent_residual).mean()
+
+
+# =============================================================================
+# FK-based Kinematic Losses
+# Use decoder_output_to_egocentric_transforms to get egocentric positions/rotations
+# =============================================================================
+
+
+class G1JointPositionLoss(nn.Module):
+    """Loss on joint body positions via FK with optional velocity loss."""
+
+    def __init__(
+        self,
+        loss_type="mse",
+        vel_weight=0.0,
+        dt=0.1,
+        include_extended=False,
+        coordinate_frame="egocentric",
+        normalize=False,
+        **kwargs,
+    ):
+        super().__init__()
+        self._skeleton_name = kwargs.get("skeleton_name", "motion_g1_extended_toe")
+        self.loss_type = loss_type
+        self.vel_weight = vel_weight
+        self.dt = (
+            dt  # Time step between future frames (from env.commands.motion.dt_future_ref_frames)
+        )
+        self._include_extended = include_extended
+        assert coordinate_frame in (
+            "egocentric",
+            "world",
+        ), f"coordinate_frame must be 'egocentric' or 'world', got '{coordinate_frame}'"
+        self.coordinate_frame = coordinate_frame
+        self._humanoid = create_humanoid(self._skeleton_name)
+        self._dof_converter = order_converter.G1Converter()
+        if normalize:
+            num_b = (
+                self._humanoid.num_bodies_augment
+                if self._include_extended
+                else self._humanoid.num_bodies
+            )
+            self._normalizer = batch_normalizer.BatchNormNormalizer((num_b * 3,))
+            self._vel_normalizer = batch_normalizer.BatchNormNormalizer((num_b * 3,))
+        else:
+            self._normalizer = None
+            self._vel_normalizer = None
+
+    def _get_positions(self, decoder_output, decoder_cfg):
+        if self.coordinate_frame == "world":
+            pos, _ = decoder_output_to_world_transforms(
+                decoder_output,
+                decoder_cfg,
+                self._humanoid,
+                self._dof_converter,
+                include_extended=self._include_extended,
+            )
+        else:
+            pos, _ = decoder_output_to_egocentric_transforms(
+                decoder_output,
+                decoder_cfg,
+                self._humanoid,
+                self._dof_converter,
+                include_extended=self._include_extended,
+            )
+        return pos
+
+    def forward(self, loss_inputs):
+        tokenizer_obs = loss_inputs["tokenizer_obs"]
+        decoders_cfg = loss_inputs["decoders_cfg"]
+        decoded_outputs = loss_inputs["decoded_outputs"]
+        frame_mask = loss_inputs.get("frame_mask", None)
+
+        device = _get_device_from_loss_inputs(loss_inputs)
+        if self._humanoid.device != device:
+            self._humanoid = self._humanoid.to(device=device)
+
+        pos_gt = self._get_positions(tokenizer_obs, decoders_cfg["g1_kin"])
+        pos_pred = self._get_positions(decoded_outputs["g1_kin"], decoders_cfg["g1_kin"])
+
+        # Position loss
+        # TCN/conv decoders may produce fewer frames; truncate gt to match pred (time dim -3)
+        if pos_pred.shape[-3] < pos_gt.shape[-3]:
+            pos_gt = pos_gt[..., : pos_pred.shape[-3], :, :]
+
+        num_frames = pos_gt.shape[-3]
+        pos_mask = _build_frame_mask_for_loss(frame_mask, num_frames)
+
+        # Compute velocity from raw positions BEFORE normalization
+        vel_gt = vel_pred = None
+        if self.vel_weight > 0 and num_frames > 1:
+            vel_gt = (pos_gt[..., 1:, :, :] - pos_gt[..., :-1, :, :]) / self.dt
+            vel_pred = (pos_pred[..., 1:, :, :] - pos_pred[..., :-1, :, :]) / self.dt
+
+        # Apply running normalization if enabled
+        if self._normalizer is not None:
+            pos_gt, pos_pred = _apply_normalizer(self._normalizer, pos_gt, pos_pred, pos_mask)
+
+        pos_loss = _compute_masked_loss(pos_pred, pos_gt, pos_mask, self.loss_type)
+
+        # Velocity loss (finite difference along future frames, normalized by dt)
+        if self.vel_weight > 0 and vel_gt is not None:
+            vel_mask = _build_vel_frame_mask(frame_mask, num_frames)
+            if self._vel_normalizer is not None:
+                vel_gt, vel_pred = _apply_normalizer(
+                    self._vel_normalizer, vel_gt, vel_pred, vel_mask
+                )
+            vel_loss = _compute_masked_loss(vel_pred, vel_gt, vel_mask, self.loss_type)
+            return pos_loss + self.vel_weight * vel_loss
+
+        return pos_loss
+
+
+class G1JointRotationLoss(nn.Module):
+    """Loss on joint rotations via FK with optional velocity loss. Uses geodesic or Frobenius norm."""
+
+    def __init__(
+        self,
+        loss_type="frobenius",
+        vel_weight=0.0,
+        dt=0.1,
+        include_extended=False,
+        normalize=False,
+        **kwargs,
+    ):
+        super().__init__()
+        self._skeleton_name = kwargs.get("skeleton_name", "motion_g1_extended_toe")
+        self.loss_type = loss_type
+        self.vel_weight = vel_weight
+        self.dt = (
+            dt  # Time step between future frames (from env.commands.motion.dt_future_ref_frames)
+        )
+        self._include_extended = include_extended
+        self._humanoid = create_humanoid(self._skeleton_name)
+        self._dof_converter = order_converter.G1Converter()
+        # Normalization only makes sense for frobenius (element-wise); geodesic operates on SO(3)
+        if normalize and loss_type == "frobenius":
+            num_b = (
+                self._humanoid.num_bodies_augment
+                if self._include_extended
+                else self._humanoid.num_bodies
+            )
+            self._normalizer = batch_normalizer.BatchNormNormalizer((num_b * 6,))
+            # Velocity normalizer operates on DOF angles (extended joints have no qpos)
+            self._vel_normalizer = batch_normalizer.BatchNormNormalizer((self._humanoid.num_dof,))
+        else:
+            self._normalizer = None
+            self._vel_normalizer = None
+
+    def forward(self, loss_inputs):
+        tokenizer_obs = loss_inputs["tokenizer_obs"]
+        decoders_cfg = loss_inputs["decoders_cfg"]
+        decoded_outputs = loss_inputs["decoded_outputs"]
+        frame_mask = loss_inputs.get("frame_mask", None)
+
+        device = _get_device_from_loss_inputs(loss_inputs)
+        if self._humanoid.device != device:
+            self._humanoid = self._humanoid.to(device=device)
+
+        # Get egocentric rotations (with extended bodies if configured)
+        _, egocentric_rot_6d_gt = decoder_output_to_egocentric_transforms(
+            tokenizer_obs,
+            decoders_cfg["g1_kin"],
+            self._humanoid,
+            self._dof_converter,
+            include_extended=self._include_extended,
+        )
+        _, egocentric_rot_6d_pred = decoder_output_to_egocentric_transforms(
+            decoded_outputs["g1_kin"],
+            decoders_cfg["g1_kin"],
+            self._humanoid,
+            self._dof_converter,
+            include_extended=self._include_extended,
+        )
+
+        # TCN/conv decoders may produce fewer frames; truncate gt to match pred (time dim -3)
+        if egocentric_rot_6d_pred.shape[-3] < egocentric_rot_6d_gt.shape[-3]:
+            egocentric_rot_6d_gt = egocentric_rot_6d_gt[
+                ..., : egocentric_rot_6d_pred.shape[-3], :, :
+            ]
+
+        num_frames = egocentric_rot_6d_gt.shape[-3]
+
+        if self.loss_type == "frobenius":
+            rot_mask = _build_frame_mask_for_loss(frame_mask, num_frames)
+
+            # Compute velocity from DOF positions (qpos) BEFORE normalization
+            vel_gt = vel_pred = None
+            if self.vel_weight > 0 and num_frames > 1:
+                dof_pos_gt = _extract_dof_pos(
+                    tokenizer_obs,
+                    decoders_cfg["g1_kin"],
+                    self._humanoid,
+                    self._dof_converter,
+                )
+                dof_pos_pred = _extract_dof_pos(
+                    decoded_outputs["g1_kin"],
+                    decoders_cfg["g1_kin"],
+                    self._humanoid,
+                    self._dof_converter,
+                )
+                if dof_pos_gt is not None and dof_pos_pred is not None:
+                    # Truncate gt to match pred frames (same as egocentric transforms above)
+                    if dof_pos_pred.shape[-2] < dof_pos_gt.shape[-2]:
+                        dof_pos_gt = dof_pos_gt[..., : dof_pos_pred.shape[-2], :]
+                    vel_gt = (dof_pos_gt[..., 1:, :] - dof_pos_gt[..., :-1, :]) / self.dt
+                    vel_pred = (dof_pos_pred[..., 1:, :] - dof_pos_pred[..., :-1, :]) / self.dt
+
+            # Apply running normalization if enabled
+            if self._normalizer is not None:
+                egocentric_rot_6d_gt, egocentric_rot_6d_pred = _apply_normalizer(
+                    self._normalizer, egocentric_rot_6d_gt, egocentric_rot_6d_pred, rot_mask
+                )
+
+            rot_loss = _compute_masked_loss(
+                egocentric_rot_6d_pred, egocentric_rot_6d_gt, rot_mask, "mse"
+            )
+
+            # Velocity loss on DOF positions
+            if self.vel_weight > 0 and vel_gt is not None:
+                vel_mask = _build_vel_frame_mask(frame_mask, num_frames)
+                if self._vel_normalizer is not None:
+                    vel_gt, vel_pred = _apply_normalizer(
+                        self._vel_normalizer, vel_gt, vel_pred, vel_mask
+                    )
+                vel_loss = _compute_masked_loss(vel_pred, vel_gt, vel_mask, "mse")
+                return rot_loss + self.vel_weight * vel_loss
+
+            return rot_loss
+
+        elif self.loss_type == "geodesic":
+            # Convert 6D to 3x3 rotation matrices for geodesic loss
+            egocentric_rot_gt = rotations.rot6d_to_mat_first_two_cols(egocentric_rot_6d_gt)
+            egocentric_rot_pred = rotations.rot6d_to_mat_first_two_cols(egocentric_rot_6d_pred)
+
+            geo_mask = _build_frame_mask_for_loss(frame_mask, num_frames)
+            rot_loss = _masked_geodesic_angle(egocentric_rot_pred, egocentric_rot_gt, mask=geo_mask)
+
+            # Angular velocity loss (relative rotation between consecutive frames)
+            if (
+                self.vel_weight > 0 and egocentric_rot_gt.shape[-4] > 1
+            ):  # [..., num_future, num_bodies, 3, 3]
+                rot_vel_gt = torch.matmul(
+                    egocentric_rot_gt[..., 1:, :, :, :],
+                    egocentric_rot_gt[..., :-1, :, :, :].transpose(-1, -2),
+                )
+                rot_vel_pred = torch.matmul(
+                    egocentric_rot_pred[..., 1:, :, :, :],
+                    egocentric_rot_pred[..., :-1, :, :, :].transpose(-1, -2),
+                )
+                vel_geo_mask = _build_vel_frame_mask(frame_mask, num_frames)
+                vel_loss = _masked_geodesic_angle(
+                    rot_vel_pred, rot_vel_gt, mask=vel_geo_mask, dt=self.dt
+                )
+                return rot_loss + self.vel_weight * vel_loss
+
+            return rot_loss
+
+        else:
+            raise ValueError(f"Unknown loss_type: {self.loss_type}")
+
+
+class G1RootPositionLoss(nn.Module):
+    """Loss on root (pelvis) position relative to first reference frame.
+
+    Only available for decoder output formats that include root transforms
+    (e.g., command_multi_future_root_transforms_nonflat).
+    """
+
+    def __init__(
+        self, loss_type="mse", vel_weight=0.0, dt=0.1, normalize=False, **kwargs  # noqa: ARG002
+    ):  # noqa: ARG002
+        super().__init__()
+        self.loss_type = loss_type
+        self.vel_weight = vel_weight
+        self.dt = dt
+        self._normalizer = batch_normalizer.BatchNormNormalizer((3,)) if normalize else None
+        self._vel_normalizer = batch_normalizer.BatchNormNormalizer((3,)) if normalize else None
+
+    def forward(self, loss_inputs):
+        tokenizer_obs = loss_inputs["tokenizer_obs"]
+        decoders_cfg = loss_inputs["decoders_cfg"]
+        decoded_outputs = loss_inputs["decoded_outputs"]
+        frame_mask = loss_inputs.get("frame_mask", None)
+
+        root_pos_gt, _ = decoder_output_to_root_transforms(tokenizer_obs, decoders_cfg["g1_kin"])
+        root_pos_pred, _ = decoder_output_to_root_transforms(
+            decoded_outputs["g1_kin"], decoders_cfg["g1_kin"]
+        )
+
+        # TCN/conv decoders may produce fewer frames; truncate gt to match pred (time dim -2)
+        if root_pos_pred.shape[-2] < root_pos_gt.shape[-2]:
+            root_pos_gt = root_pos_gt[..., : root_pos_pred.shape[-2], :]
+
+        num_frames = root_pos_gt.shape[-2]
+        pos_mask = _build_frame_mask_for_loss(frame_mask, num_frames)
+
+        # Compute velocity from raw positions BEFORE normalization
+        vel_gt = vel_pred = None
+        if self.vel_weight > 0 and num_frames > 1:
+            vel_gt = (root_pos_gt[..., 1:, :] - root_pos_gt[..., :-1, :]) / self.dt
+            vel_pred = (root_pos_pred[..., 1:, :] - root_pos_pred[..., :-1, :]) / self.dt
+
+        # Apply running normalization if enabled
+        if self._normalizer is not None:
+            root_pos_gt, root_pos_pred = _apply_normalizer(
+                self._normalizer, root_pos_gt, root_pos_pred, pos_mask
+            )
+
+        pos_loss = _compute_masked_loss(root_pos_pred, root_pos_gt, pos_mask, self.loss_type)
+
+        # Velocity loss (finite difference along future frames)
+        if self.vel_weight > 0 and vel_gt is not None:
+            vel_mask = _build_vel_frame_mask(frame_mask, num_frames)
+            if self._vel_normalizer is not None:
+                vel_gt, vel_pred = _apply_normalizer(
+                    self._vel_normalizer, vel_gt, vel_pred, vel_mask
+                )
+            vel_loss = _compute_masked_loss(vel_pred, vel_gt, vel_mask, self.loss_type)
+            return pos_loss + self.vel_weight * vel_loss
+
+        return pos_loss
+
+
+class G1RootRotationLoss(nn.Module):
+    """Loss on root (pelvis) rotation in robot frame.
+
+    Supports both Frobenius (on 6D representation) and geodesic (on 3x3 matrices) losses.
+    """
+
+    def __init__(
+        self,
+        loss_type="frobenius",
+        vel_weight=0.0,
+        dt=0.1,
+        normalize=False,
+        **kwargs,  # noqa: ARG002
+    ):  # noqa: ARG002
+        super().__init__()
+        self.loss_type = loss_type
+        self.vel_weight = vel_weight
+        self.dt = dt
+        # Normalization only makes sense for frobenius (element-wise); geodesic operates on SO(3)
+        self._normalizer = (
+            batch_normalizer.BatchNormNormalizer((6,))
+            if (normalize and loss_type == "frobenius")
+            else None
+        )
+        self._vel_normalizer = (
+            batch_normalizer.BatchNormNormalizer((6,))
+            if (normalize and loss_type == "frobenius")
+            else None
+        )
+
+    def forward(self, loss_inputs):
+        tokenizer_obs = loss_inputs["tokenizer_obs"]
+        decoders_cfg = loss_inputs["decoders_cfg"]
+        decoded_outputs = loss_inputs["decoded_outputs"]
+        frame_mask = loss_inputs.get("frame_mask", None)
+
+        # Get 6D rotations from decoder output
+        _, root_rot_6d_gt = decoder_output_to_root_transforms(tokenizer_obs, decoders_cfg["g1_kin"])
+        _, root_rot_6d_pred = decoder_output_to_root_transforms(
+            decoded_outputs["g1_kin"], decoders_cfg["g1_kin"]
+        )
+
+        # TCN/conv decoders may produce fewer frames; truncate gt to match pred (time dim -2)
+        if root_rot_6d_pred.shape[-2] < root_rot_6d_gt.shape[-2]:
+            root_rot_6d_gt = root_rot_6d_gt[..., : root_rot_6d_pred.shape[-2], :]
+
+        num_frames = root_rot_6d_gt.shape[-2]
+
+        if self.loss_type == "frobenius":
+            rot_mask = _build_frame_mask_for_loss(frame_mask, num_frames)
+
+            # Compute velocity from off-diagonal elements of rotation matrix BEFORE normalization
+            vel_gt = vel_pred = None
+            if self.vel_weight > 0 and num_frames > 1:
+                # Convert 6D -> 3x3 rotation matrix
+                R_gt = rotations.rot6d_to_mat_first_two_cols(root_rot_6d_gt)  # [..., F, 3, 3]
+                R_pred = rotations.rot6d_to_mat_first_two_cols(root_rot_6d_pred)  # [..., F, 3, 3]
+                # Extract 6 off-diagonal elements
+                idx_r = [1, 2, 0, 2, 0, 1]
+                idx_c = [0, 0, 1, 1, 2, 2]
+                off_gt = R_gt[..., idx_r, idx_c]  # [..., F, 6]
+                off_pred = R_pred[..., idx_r, idx_c]  # [..., F, 6]
+                # Finite difference velocity
+                vel_gt = (off_gt[..., 1:, :] - off_gt[..., :-1, :]) / self.dt
+                vel_pred = (off_pred[..., 1:, :] - off_pred[..., :-1, :]) / self.dt
+
+            # Apply running normalization if enabled
+            if self._normalizer is not None:
+                root_rot_6d_gt, root_rot_6d_pred = _apply_normalizer(
+                    self._normalizer, root_rot_6d_gt, root_rot_6d_pred, rot_mask
+                )
+
+            rot_loss = _compute_masked_loss(root_rot_6d_pred, root_rot_6d_gt, rot_mask, "mse")
+
+            # Velocity loss on off-diagonal elements of rotation matrix
+            if self.vel_weight > 0 and vel_gt is not None:
+                vel_mask = _build_vel_frame_mask(frame_mask, num_frames)
+                if self._vel_normalizer is not None:
+                    vel_gt, vel_pred = _apply_normalizer(
+                        self._vel_normalizer, vel_gt, vel_pred, vel_mask
+                    )
+                vel_loss = _compute_masked_loss(vel_pred, vel_gt, vel_mask, "mse")
+                return rot_loss + self.vel_weight * vel_loss
+
+            return rot_loss
+
+        elif self.loss_type == "geodesic":
+            # Convert 6D to 3x3 rotation matrices for geodesic loss
+            root_rot_gt = rotations.rot6d_to_mat_first_two_cols(root_rot_6d_gt)
+            root_rot_pred = rotations.rot6d_to_mat_first_two_cols(root_rot_6d_pred)
+
+            geo_mask = _build_frame_mask_for_loss(frame_mask, num_frames)
+            rot_loss = _masked_geodesic_angle(root_rot_pred, root_rot_gt, mask=geo_mask)
+
+            # Angular velocity loss (relative rotation between consecutive frames)
+            if self.vel_weight > 0 and root_rot_gt.shape[-3] > 1:  # [..., num_future, 3, 3]
+                rot_vel_gt = torch.matmul(
+                    root_rot_gt[..., 1:, :, :],
+                    root_rot_gt[..., :-1, :, :].transpose(-1, -2),
+                )
+                rot_vel_pred = torch.matmul(
+                    root_rot_pred[..., 1:, :, :],
+                    root_rot_pred[..., :-1, :, :].transpose(-1, -2),
+                )
+                vel_geo_mask = _build_vel_frame_mask(frame_mask, num_frames)
+                vel_loss = _masked_geodesic_angle(
+                    rot_vel_pred, rot_vel_gt, mask=vel_geo_mask, dt=self.dt
+                )
+                return rot_loss + self.vel_weight * vel_loss
+
+            return rot_loss
+
+        else:
+            raise ValueError(f"Unknown loss_type: {self.loss_type}")
+
+
+class G1FootContactLoss(nn.Module):
+    """Foot skating / foot contact loss.
+
+    Penalizes foot velocity when the foot is in contact with the ground.
+    Contact is determined by a height threshold: feet below the threshold
+    are considered in contact and should have zero horizontal velocity.
+
+    Note: This loss requires the new egocentric observation format
+    (command_multi_future_egocentric_joint_transforms + root_transforms).
+    With the old qpos-based format (command_multi_future_nonflat), root is
+    identity so z-coordinates are not in world frame — contact detection
+    will produce incorrect results.
+    """
+
+    def __init__(
+        self,
+        dt=0.05,
+        foot_joint_names=("left_ankle_link", "right_ankle_link"),
+        contact_height_threshold=0.05,
+        vel_threshold=0.15,
+        dt_warning_threshold=0.05,
+        **kwargs,
+    ):
+        super().__init__()
+        self._skeleton_name = kwargs.get("skeleton_name", "motion_g1_extended_toe")
+        self.dt = dt
+        self.foot_joint_names = list(foot_joint_names)
+        self.contact_height_threshold = contact_height_threshold
+        self.vel_threshold = vel_threshold
+        self._humanoid = create_humanoid(self._skeleton_name)
+        self._dof_converter = order_converter.G1Converter()
+        self._foot_indices = [
+            self._humanoid.body_names_augment.index(name) for name in self.foot_joint_names
+        ]
+
+        if dt > dt_warning_threshold:
+            import warnings
+
+            warnings.warn(
+                f"G1FootContactLoss: dt={dt:.3f}s is large (>{dt_warning_threshold}s). "
+                f"Foot sliding loss may be less effective with coarse temporal resolution.",
+                stacklevel=2,
+            )
+
+    def forward(self, loss_inputs):
+        tokenizer_obs = loss_inputs["tokenizer_obs"]
+        decoders_cfg = loss_inputs["decoders_cfg"]
+        decoded_outputs = loss_inputs["decoded_outputs"]
+        frame_mask = loss_inputs.get("frame_mask", None)
+
+        device = _get_device_from_loss_inputs(loss_inputs)
+        if self._humanoid.device != device:
+            self._humanoid = self._humanoid.to(device=device)
+
+        # Get joint positions in world frame for GT and predicted (include extended for toe joints)
+        world_pos_gt, _ = decoder_output_to_world_transforms(
+            tokenizer_obs,
+            decoders_cfg["g1_kin"],
+            self._humanoid,
+            self._dof_converter,
+            include_extended=True,
+        )
+        world_pos_pred, _ = decoder_output_to_world_transforms(
+            decoded_outputs["g1_kin"],
+            decoders_cfg["g1_kin"],
+            self._humanoid,
+            self._dof_converter,
+            include_extended=True,
+        )
+        feet_pos_gt = world_pos_gt[..., self._foot_indices, :]
+        feet_pos_pred = world_pos_pred[..., self._foot_indices, :]
+
+        # TCN/conv decoders may produce fewer frames; truncate gt to match pred (time dim -3)
+        if feet_pos_pred.shape[-3] < feet_pos_gt.shape[-3]:
+            feet_pos_gt = feet_pos_gt[..., : feet_pos_pred.shape[-3], :, :]
+
+        # Need at least 2 frames for velocity
+        if feet_pos_pred.shape[-3] < 2:
+            return torch.tensor(0.0, device=feet_pos_pred.device)
+
+        num_frames = feet_pos_gt.shape[-3]
+
+        # Contact mask from GT: foot height (z) below threshold AND velocity below threshold
+        gt_foot_height = feet_pos_gt[..., :-1, :, 2]  # [..., num_future-1, num_feet]
+        gt_foot_vel = (
+            torch.norm(feet_pos_gt[..., 1:, :, :] - feet_pos_gt[..., :-1, :, :], dim=-1) / self.dt
+        )
+        contact_mask = (
+            (gt_foot_height < self.contact_height_threshold) & (gt_foot_vel < self.vel_threshold)
+        ).float()
+
+        # Apply variable frame mask to contact_mask (velocity uses frame pairs)
+        # contact_mask shape: [..., F-1, num_feet] — expand vel mask to match
+        vel_frame_mask = _build_vel_frame_mask(frame_mask, num_frames)
+        if vel_frame_mask is not None:
+            contact_mask = contact_mask * vel_frame_mask.unsqueeze(-1)
+
+        # Predicted foot velocity: ||pos[t+1] - pos[t]|| / dt
+        pred_foot_vel = (
+            torch.norm(feet_pos_pred[..., 1:, :, :] - feet_pos_pred[..., :-1, :, :], dim=-1)
+            / self.dt
+        )
+
+        # Penalize predicted velocity when GT says foot is in contact
+        vel_err = pred_foot_vel * contact_mask
+
+        # Mean over contacting frames (avoid division by zero)
+        total_contact = contact_mask.sum() + 1e-6
+        skating_loss = vel_err.sum() / total_contact
+
+        return skating_loss
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/modules/__init__.py b/GR00T-WholeBodyControl/gear_sonic/trl/modules/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/modules/actor_critic_modules.py b/GR00T-WholeBodyControl/gear_sonic/trl/modules/actor_critic_modules.py
new file mode 100644
index 0000000000000000000000000000000000000000..0ab27d1835746aa61f3b87a77d042a64f550a58d
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/modules/actor_critic_modules.py
@@ -0,0 +1,656 @@
+"""Actor and Critic modules for PPO-based reinforcement learning."""
+
+from __future__ import annotations
+
+from copy import deepcopy
+
+from tensordict import TensorDict
+import torch
+from torch.distributions import Normal
+import torch.nn as nn
+
+from gear_sonic.trl.utils.common import Timer, custom_instantiate
+from gear_sonic.trl.utils.rl import compute_episode_attnmask
+from gear_sonic.utils.batch_normalizer import BatchNormNormalizer
+from gear_sonic.utils.running_mean_std import RunningMeanStd
+
+
+class Actor(nn.Module):
+    """Policy network that maps observations to action distributions.
+
+    Wraps an arbitrary backbone network and adds a diagonal Gaussian action
+    distribution on top. Supports both direct ``std`` and ``log_std``
+    parameterizations for exploration noise, with optional clamping.
+
+    The actor maintains an observation buffer for temporal models (e.g.
+    transformers) that require a history of past observations. During rollout,
+    observations are appended to this buffer up to ``max_rollout_history``
+    steps, and an episode attention mask is computed from done signals so the
+    backbone can attend only within episode boundaries.
+    """
+
+    def __init__(
+        self,
+        env_config,
+        algo_config,
+        backbone,
+        obs_dim_dict=None,
+        module_dim_dict={},
+        running_mean_std=False,
+        use_batch_norm=False,
+        max_rollout_history=1,
+        input_key="actor_obs",
+        input_obs_dict=False,
+        has_aux_loss=False,
+        output_original_obs_dict=False,
+        backbone_kwargs={},
+    ):
+        """Initialize the Actor.
+
+        Args:
+            env_config: Environment configuration containing robot specs.
+            algo_config: Algorithm configuration (noise std, clamping, etc.).
+            backbone: Hydra-style config for the backbone network to instantiate.
+            obs_dim_dict: Mapping from observation keys to their dimensions.
+                Defaults to ``env_config.robot.algo_obs_dim_dict``.
+            module_dim_dict: Additional dimension info passed to the backbone.
+            running_mean_std: Whether to normalize inputs with running statistics.
+            use_batch_norm: Whether to normalize inputs with batch normalization.
+            max_rollout_history: Number of past timesteps to keep in the
+                observation buffer for temporal models.
+            input_key: Key in ``obs_dict`` used as network input.
+            input_obs_dict: If True, pass the entire obs dict to the backbone
+                instead of a single tensor.
+            has_aux_loss: Whether the backbone produces auxiliary losses
+                (e.g. commitment loss from VQ).
+            output_original_obs_dict: If True, include the observation dict
+                in the output TensorDict.
+            backbone_kwargs: Extra keyword arguments forwarded to the backbone
+                constructor.
+        """
+        super().__init__()
+
+        self.algo_config = algo_config
+        self.env_config = env_config
+        if obs_dim_dict is None:
+            obs_dim_dict = env_config.robot.algo_obs_dim_dict
+        self.input_key = input_key
+        self.input_obs_dict = input_obs_dict
+        self.has_aux_loss = has_aux_loss
+        self.aux_losses = None
+        self.aux_loss_coef = None
+        self.output_original_obs_dict = output_original_obs_dict
+        self.max_rollout_history = max_rollout_history
+        self.actor_module = custom_instantiate(
+            backbone,
+            env_config=env_config,
+            algo_config=algo_config,
+            obs_dim_dict=obs_dim_dict,
+            module_dim_dict=module_dim_dict,
+            _resolve=False,
+            **backbone_kwargs,
+        )
+        self.use_batch_norm = use_batch_norm
+
+        self.use_running_mean_std = running_mean_std
+
+        self.running_mean_std = None
+        if running_mean_std:
+            self.running_mean_std = RunningMeanStd(
+                (obs_dim_dict[self.input_key],), per_channel=True
+            )
+
+        if use_batch_norm:
+            self.running_mean_std = BatchNormNormalizer((obs_dim_dict[self.input_key],))
+
+        assert not (
+            running_mean_std and use_batch_norm
+        ), "running_mean_std and use_batch_norm cannot be both True"
+
+        # Action noise
+        self.num_actions = self.env_config.robot.actions_dim
+        init_noise_std = algo_config.init_noise_std
+
+        # Support both std and log_std parameterization
+        # Using log_std ensures exp(log_std) > 0, which is numerically more stable
+        self.use_log_std = algo_config.get("use_log_std", False)
+        if self.use_log_std:
+            self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(self.num_actions)))
+        else:
+            self.std = nn.Parameter(init_noise_std * torch.ones(self.num_actions))
+
+        if algo_config.get("freeze_noise_std", False):
+            if self.use_log_std:
+                self.log_std.requires_grad = False
+            else:
+                self.std.requires_grad = False
+
+        self.clamp_noise_std = algo_config.get("clamp_noise_std", False)
+        if self.clamp_noise_std:
+            self.max_noise_std = algo_config.get("max_noise_std", 1.0)
+
+        self.distribution = None
+        # disable args validation for speedup
+        Normal.set_default_validate_args(False)
+
+        # Initialize observation buffer for rollout
+        self.obs_dict_buffer = TensorDict()
+        self.dones_buffer = None
+        self.steps = 0
+        self.is_eval_mode = False
+
+    def reset(self, dones=None):
+        pass
+
+    @property
+    def get_std(self):
+        """Get the standard deviation, handling both std and log_std parameterizations."""
+        if self.use_log_std:
+            # First, handle NaN or inf in log_std
+            if torch.any(torch.isnan(self.log_std)) or torch.any(torch.isinf(self.log_std)):
+                print("[ERROR] log_std contains NaN or Inf! Resetting to safe values.")
+                with torch.no_grad():
+                    self.log_std.data = torch.log(torch.ones_like(self.log_std) * 0.5)
+
+            # Apply clamping if configured before computing std
+            if self.algo_config.get("use_clampped_std", False):
+                std_min = self.algo_config.std_clamp_min
+                std_max = self.algo_config.std_clamp_max
+                log_std_clamped = torch.clamp(
+                    self.log_std,
+                    min=torch.log(
+                        torch.tensor(std_min, dtype=self.log_std.dtype, device=self.log_std.device)
+                    ),
+                    max=torch.log(
+                        torch.tensor(std_max, dtype=self.log_std.dtype, device=self.log_std.device)
+                    ),
+                )
+                std = torch.exp(log_std_clamped)
+                std = torch.clamp(std, min=std_min, max=std_max)
+                return std
+
+            if self.clamp_noise_std:
+                log_std_clamped = torch.clamp(
+                    self.log_std,
+                    max=torch.log(
+                        torch.tensor(
+                            self.max_noise_std, dtype=self.log_std.dtype, device=self.log_std.device
+                        )
+                    ),
+                )
+                std = torch.exp(log_std_clamped)
+                std = torch.clamp(std, min=1e-6)
+                return std
+
+            # Default case: clamp log_std to prevent extreme values
+            log_std_clamped = torch.clamp(self.log_std, min=-20, max=2)
+            std = torch.exp(log_std_clamped)
+            std = torch.clamp(std, min=1e-6)
+            return std
+        else:
+            # Original std parameterization with in-place clamping
+            if self.algo_config.get("use_clampped_std", False):
+                with torch.no_grad():
+                    self.std.clamp_(
+                        min=self.algo_config.std_clamp_min, max=self.algo_config.std_clamp_max
+                    )
+            if self.clamp_noise_std:
+                with torch.no_grad():
+                    self.std.clamp_(max=self.max_noise_std)
+            return self.std
+
+    def forward(self, obs_dict, is_training=False, **kwargs):
+        """Compute action means from observations.
+
+        Optionally normalizes input observations and collects auxiliary losses
+        from the backbone when training with VQ or similar modules.
+
+        Args:
+            obs_dict: Dictionary mapping observation keys to tensors.
+            is_training: If True and ``has_aux_loss``, request auxiliary losses
+                from the backbone.
+            **kwargs: Forwarded to the backbone (e.g. ``episode_attnmask``).
+
+        Returns:
+            Action mean tensor of shape ``(batch, act_dim)`` or
+            ``(batch, seq, act_dim)`` for temporal models.
+        """
+        obs_dict = obs_dict.copy()
+        if self.running_mean_std is not None:
+            if self.use_batch_norm:
+                obs_dict[self.input_key] = self.running_mean_std(obs_dict[self.input_key])
+            else:
+                with torch.no_grad():
+                    obs_dict[self.input_key] = self.running_mean_std(obs_dict[self.input_key])
+
+        with Timer("actor_module", instance_enabled=self.training):
+            if self.input_obs_dict:
+                net_input = obs_dict
+            else:
+                net_input = obs_dict[self.input_key]
+            net_kwargs = kwargs.copy()
+            if self.has_aux_loss and is_training:
+                net_kwargs["compute_aux_loss"] = True
+            output = self.actor_module(net_input, **net_kwargs)
+        if self.has_aux_loss and is_training:
+            # output needs to be a dict
+            action_mean = output["action_mean"]
+            self.aux_losses = output["aux_losses"]
+            self.aux_loss_coef = output["aux_loss_coef"]
+        else:
+            action_mean = output
+        return action_mean
+
+    @property
+    def has_normalized_actions(self):
+        return False
+
+    @property
+    def action_mean(self):
+        return self.distribution.mean
+
+    @property
+    def action_std(self):
+        return self.distribution.stddev
+
+    @property
+    def entropy(self):
+        return self.distribution.entropy().sum(dim=-1)
+
+    def update_distribution(
+        self, obs_dict, episode_attnmask=None, last_step_only=False, is_training=False, **kwargs
+    ):
+        """Compute forward pass and update the internal Gaussian distribution.
+
+        Args:
+            obs_dict: Observation dictionary. For temporal models this has shape
+                ``{key: (batch, seq, dim)}``.
+            episode_attnmask: Optional causal attention mask of shape
+                ``(batch, seq, seq)`` for transformer backbones.
+            last_step_only: If True, use only the last timestep's mean for
+                the distribution (used during rollout with temporal models).
+            is_training: Forwarded to ``forward`` to enable aux loss collection.
+            **kwargs: Forwarded to ``forward``.
+        """
+        mean = self.forward(
+            obs_dict, episode_attnmask=episode_attnmask, is_training=is_training, **kwargs
+        )
+        if last_step_only:
+            mean = mean[:, -1]
+
+        # Get std using the property that handles both parameterizations
+        std = self.get_std
+        # Safety check for NaN or negative values
+        if torch.any(torch.isnan(std)) or torch.any(std <= 0):
+            print(f"[WARNING] Invalid std detected! std: {std}")
+            std = torch.clamp(std, min=1e-6)
+        self.distribution = Normal(mean, (mean * 0.0 + std).clamp(min=1e-6))
+
+    def act(self, obs_dict, episode_attnmask=None, **kwargs):
+        """Sample actions from the current policy for a single timestep.
+
+        Update the action distribution and sample from it. Used during the
+        PPO learning phase (not rollout) where the full observation sequence
+        is available.
+
+        Args:
+            obs_dict: Observation dictionary with shape
+                ``{key: (batch, seq, dim)}``.
+            episode_attnmask: Optional attention mask of shape
+                ``(batch, seq, seq)``.
+            **kwargs: Forwarded to ``update_distribution``.
+
+        Returns:
+            TensorDict with keys ``actions`` ``(batch, act_dim)``,
+            ``action_mean``, ``action_sigma``, and optionally ``obs_dict``.
+        """
+        # try:
+        self.update_distribution(
+            obs_dict, episode_attnmask=episode_attnmask, is_training=True, **kwargs
+        )
+        # except Exception as e:
+        #     import ipdb; ipdb.set_trace()
+        #     raise e
+        actions = self.distribution.sample()
+        return TensorDict(
+            {
+                "actions": actions,
+                "action_mean": self.action_mean,
+                "action_sigma": self.action_std,
+                "obs_dict": obs_dict if self.output_original_obs_dict else None,
+            }
+        )
+
+    def update_dones_buffer_and_compute_episode_attnmask(self, cur_dones):
+        """Update the done-signal buffer and derive an episode attention mask.
+
+        Maintains a sliding window of done flags (length
+        ``max_rollout_history - 1``) and converts them into a causal attention
+        mask that prevents attending across episode boundaries.
+
+        Args:
+            cur_dones: Done flags for the current step, shape ``(batch,)``.
+
+        Returns:
+            Episode attention mask of shape
+            ``(batch, history_len, history_len)``.
+        """
+        if self.steps > 0 and self.max_rollout_history > 1:
+            if self.dones_buffer is None:
+                self.dones_buffer = cur_dones.clone().unsqueeze(1)
+            else:
+                self.dones_buffer = torch.cat(
+                    [self.dones_buffer, cur_dones.clone().unsqueeze(1)], dim=1
+                )
+            if self.dones_buffer.shape[1] > self.max_rollout_history - 1:
+                self.dones_buffer = self.dones_buffer[:, -self.max_rollout_history + 1 :]
+            dones = torch.cat(
+                [self.dones_buffer, torch.zeros_like(self.dones_buffer[:, :1])], dim=1
+            )
+        else:
+            dones = torch.zeros_like(cur_dones.unsqueeze(1))
+        episode_attnmask_from_dones = compute_episode_attnmask(dones)
+        return episode_attnmask_from_dones
+
+    def _update_obs_buffer(self, obs_dict, episode_attnmask=None, cur_dones=None):
+        """Append new observations to the rolling history buffer.
+
+        Grows the buffer up to ``max_rollout_history`` timesteps, then slides
+        the window. When ``cur_dones`` is provided, also updates the done
+        buffer and computes the episode attention mask.
+
+        Args:
+            obs_dict: Single-step observations ``{key: (batch, dim)}``.
+            episode_attnmask: Optional externally provided attention mask.
+                If both this and ``cur_dones`` are given, consistency is
+                asserted.
+            cur_dones: Optional done flags for the current step,
+                shape ``(batch,)``.
+
+        Returns:
+            Episode attention mask of shape
+            ``(batch, history_len, history_len)`` or None.
+        """
+        update_episode_attnmask = False
+
+        for key in obs_dict.keys():
+            if key not in self.obs_dict_buffer:
+                self.obs_dict_buffer[key] = obs_dict[key].unsqueeze(1)
+            else:
+                self.obs_dict_buffer[key] = torch.cat(
+                    [self.obs_dict_buffer[key], obs_dict[key].unsqueeze(1)], dim=1
+                )
+            if self.obs_dict_buffer[key].shape[1] > self.max_rollout_history:
+                update_episode_attnmask = True
+                self.obs_dict_buffer[key] = self.obs_dict_buffer[key][
+                    :, -self.max_rollout_history :
+                ]
+
+        if episode_attnmask is not None and update_episode_attnmask:
+            episode_attnmask = episode_attnmask[
+                :, -self.max_rollout_history :, -self.max_rollout_history :
+            ]
+
+        if cur_dones is not None:
+            episode_attnmask_from_dones = self.update_dones_buffer_and_compute_episode_attnmask(
+                cur_dones
+            )
+            if episode_attnmask is not None:
+                assert (episode_attnmask == episode_attnmask_from_dones).all()
+            if episode_attnmask is None:
+                episode_attnmask = episode_attnmask_from_dones
+
+        return episode_attnmask
+
+    def rollout(self, obs_dict, episode_attnmask=None, cur_dones=None, **kwargs):
+        """Execute one rollout step: buffer observations, sample actions.
+
+        Appends the current observations to the history buffer, runs the
+        forward pass over the full buffer, and samples from the resulting
+        distribution using only the last timestep's output.
+
+        Args:
+            obs_dict: Single-step observations ``{key: (batch, dim)}``.
+            episode_attnmask: Optional attention mask.
+            cur_dones: Done flags from the previous step, shape ``(batch,)``.
+            **kwargs: Forwarded to ``update_distribution``.
+
+        Returns:
+            TensorDict with keys ``actions`` ``(batch, act_dim)``,
+            ``action_mean``, ``action_sigma``, and optionally ``obs_dict``.
+        """
+        episode_attnmask = self._update_obs_buffer(obs_dict, episode_attnmask, cur_dones)
+        self.update_distribution(
+            obs_dict=self.obs_dict_buffer,
+            episode_attnmask=episode_attnmask,
+            last_step_only=True,
+            **kwargs,
+        )
+        self.steps += 1
+        return TensorDict(
+            {
+                "actions": self.distribution.sample(),
+                "action_mean": self.action_mean,
+                "action_sigma": self.action_std,
+                "obs_dict": self.obs_dict_buffer if self.output_original_obs_dict else None,
+            }
+        )
+
+    def rollout_with_tokens(
+        self, obs_dict, external_tokens, episode_attnmask=None, cur_dones=None, **kwargs
+    ):
+        """
+        Rollout with externally provided tokens (bypasses encoder).
+
+        This is used when an external model (e.g., kinematic diffusion) provides
+        pre-computed FSQ tokens, allowing the encoder to be bypassed while still
+        using the decoder for action generation.
+
+        Args:
+            obs_dict: Observation dict (for proprioception)
+            external_tokens: Pre-computed FSQ tokens, shape (B, 2, 32)
+            episode_attnmask: Optional attention mask
+            cur_dones: Optional done flags
+
+        Returns:
+            TensorDict with actions, action_mean, action_sigma
+        """
+        # Update observation buffer (for proprioception history)
+        episode_attnmask = self._update_obs_buffer(obs_dict, episode_attnmask, cur_dones)
+
+        # Get action mean using external tokens (bypasses encoder)
+        # actor_module should have forward_with_external_tokens method
+        if not hasattr(self.actor_module, "forward_with_external_tokens"):
+            raise NotImplementedError(
+                "actor_module does not have forward_with_external_tokens method. "
+                "This is required for token bypass mode."
+            )
+
+        # Use the last step of obs buffer for proprioception
+        obs_dict_last = {k: v[:, -1:] for k, v in self.obs_dict_buffer.items()}
+
+        action_mean = self.actor_module.forward_with_external_tokens(
+            input_data=obs_dict_last, external_tokens=external_tokens, **kwargs
+        )
+
+        # Update distribution
+        self.distribution = Normal(action_mean, (action_mean * 0.0 + self.std).clamp(min=1e-6))
+
+        self.steps += 1
+        return TensorDict(
+            {
+                "actions": self.distribution.sample(),
+                "action_mean": self.action_mean,
+                "action_sigma": self.action_std,
+            }
+        )
+
+    def get_actions_log_prob(self, actions):
+        """Compute log-probability of actions under the current distribution.
+
+        Args:
+            actions: Action tensor of shape ``(batch, act_dim)``.
+
+        Returns:
+            Log-probability scalar per batch element, shape ``(batch,)``.
+        """
+        return self.distribution.log_prob(actions).sum(dim=-1)
+
+    def act_inference(self, obs_dict, episode_attnmask=None, cur_dones=None, **kwargs):
+        """Compute deterministic actions for inference during rollout.
+
+        Similar to ``rollout`` but returns only the action mean (no sampling),
+        suitable for evaluation where stochastic exploration is not desired.
+
+        Args:
+            obs_dict: Single-step observations ``{key: (batch, dim)}``.
+            episode_attnmask: Optional attention mask.
+            cur_dones: Done flags from the previous step, shape ``(batch,)``.
+            **kwargs: Forwarded to ``forward``.
+
+        Returns:
+            Deterministic action tensor of shape ``(batch, act_dim)``.
+        """
+        episode_attnmask = self._update_obs_buffer(obs_dict, episode_attnmask, cur_dones)
+        actions_mean = self.forward(
+            obs_dict=self.obs_dict_buffer, episode_attnmask=episode_attnmask, **kwargs
+        )
+        # last step only
+        actions_mean = actions_mean[:, -1]
+        self.steps += 1
+        return actions_mean
+
+    def act_pure_inference(self, obs_dict, episode_attnmask=None, **kwargs):
+        """
+        Pure inference mode for temporal models like transformer.
+        Need to construct obs_dict and episode_attnmask outside the model.
+        """
+        actions_mean = self.forward(obs_dict=obs_dict, episode_attnmask=episode_attnmask, **kwargs)
+        # last step only
+        actions_mean = actions_mean[:, -1]
+        return actions_mean
+
+    def to_cpu(self):
+        """Move the actor and its normalizers to CPU."""
+        if self.running_mean_std is not None:
+            self.running_mean_std.to("cpu")
+        self.actor = deepcopy(self.actor).to("cpu")
+        if self.use_log_std:
+            self.log_std.to("cpu")
+        else:
+            self.std.to("cpu")
+
+    def init_rollout(self):
+        """Initialize the observation buffer for rollout phase."""
+        self.obs_dict_buffer = TensorDict()
+        self.dones_buffer = None
+        self.steps = 0
+
+    def clear_rollout(self):
+        """Clear the observation buffer after rollout phase."""
+        self.obs_dict_buffer = TensorDict()
+        self.dones_buffer = None
+        self.steps = 0
+        del self.distribution
+        if self.has_aux_loss:
+            del self.aux_losses
+            del self.aux_loss_coef
+
+    def eval_mode(self):
+        self.is_eval_mode = True
+
+    def train_mode(self):
+        self.is_eval_mode = False
+
+
+class Critic(nn.Module):
+    """Value function network that estimates state values for PPO.
+
+    Wraps a backbone network and optionally normalizes critic observations
+    via running mean/std or batch normalization. The backbone receives the
+    full observation dict (keyed by ``critic_obs``) and outputs a scalar
+    value estimate per environment.
+    """
+
+    def __init__(
+        self,
+        env_config,
+        algo_config,
+        backbone,
+        obs_dim_dict=None,
+        module_dim_dict={},
+        running_mean_std=False,
+        use_batch_norm=False,
+        backbone_kwargs={},
+    ):
+        """Initialize the Critic.
+
+        Args:
+            env_config: Environment configuration containing robot specs.
+            algo_config: Algorithm configuration.
+            backbone: Hydra-style config for the backbone network to instantiate.
+            obs_dim_dict: Mapping from observation keys to their dimensions.
+                Defaults to ``env_config.robot.algo_obs_dim_dict``.
+            module_dim_dict: Additional dimension info passed to the backbone.
+            running_mean_std: Whether to normalize inputs with running statistics.
+            use_batch_norm: Whether to normalize inputs with batch normalization.
+            backbone_kwargs: Extra keyword arguments forwarded to the backbone
+                constructor.
+        """
+        super().__init__()
+
+        if obs_dim_dict is None:
+            obs_dim_dict = env_config.robot.algo_obs_dim_dict
+        self.critic_module = custom_instantiate(
+            backbone,
+            env_config=env_config,
+            algo_config=algo_config,
+            obs_dim_dict=obs_dim_dict,
+            module_dim_dict=module_dim_dict,
+            _resolve=False,
+            **backbone_kwargs,
+        )
+        self.use_batch_norm = use_batch_norm
+        self.use_running_mean_std = running_mean_std
+
+        self.running_mean_std = None
+        if running_mean_std:
+            self.running_mean_std = RunningMeanStd((obs_dim_dict["critic_obs"],), per_channel=True)
+
+        if use_batch_norm:
+            self.running_mean_std = BatchNormNormalizer((obs_dim_dict["critic_obs"],))
+
+        assert not (
+            running_mean_std and use_batch_norm
+        ), "running_mean_std and use_batch_norm cannot be both True"
+
+    @property
+    def critic(self):
+        return self.critic_module
+
+    def reset(self, dones=None):
+        pass
+
+    def evaluate(self, obs_dict, **kwargs):
+        """Compute the value estimate for given observations.
+
+        Normalizes the ``critic_obs`` entry if a normalizer is configured,
+        then forwards through the critic backbone.
+
+        Args:
+            obs_dict: Observation dictionary containing at least
+                ``critic_obs`` of shape ``(batch, critic_obs_dim)``.
+            **kwargs: Forwarded to the critic backbone.
+
+        Returns:
+            Value estimate tensor of shape ``(batch, 1)``.
+        """
+        obs_dict = obs_dict.copy()
+        if self.running_mean_std is not None:
+            if self.use_batch_norm:
+                obs_dict["critic_obs"] = self.running_mean_std(obs_dict["critic_obs"])
+            else:
+                with torch.no_grad():
+                    obs_dict["critic_obs"] = self.running_mean_std(obs_dict["critic_obs"])
+        value = self.critic(obs_dict, **kwargs)
+        return value
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/modules/base_module.py b/GR00T-WholeBodyControl/gear_sonic/trl/modules/base_module.py
new file mode 100644
index 0000000000000000000000000000000000000000..a7dacb535aca62e30e46dec37f2963d5b8cba016
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/modules/base_module.py
@@ -0,0 +1,639 @@
+"""Config-driven neural network modules for actor-critic policy networks."""
+
+import inspect
+
+import torch.nn as nn
+import torchvision.models as models
+
+
+def get_norm(norm_type, dim):
+    """Build a normalization layer from a string identifier.
+
+    Args:
+        norm_type: Normalization type string ("layer_norm") or None to skip.
+        dim: Feature dimension for the normalization layer.
+
+    Returns:
+        An ``nn.Module`` normalization layer, or None if ``norm_type`` is None.
+    """
+    if norm_type == "layer_norm":
+        return nn.LayerNorm(dim)
+    elif norm_type is None:
+        return None
+    else:
+        raise ValueError(f"Unsupported norm type: {norm_type}")
+
+
+class ResidualBlock(nn.Module):
+    """Single pre-activation residual block: ``x + Linear -> Norm -> Act(x)``.
+
+    Uses a skip connection around a linear-norm-activation sequence so gradients
+    flow unimpeded through the identity path. The linear layer preserves
+    dimensionality (``dim -> dim``).
+
+    Args:
+        dim: Feature dimension (input and output are the same size).
+        norm_type: Normalization type applied after the linear layer.
+        activation: Name of an ``nn.Module`` activation class (e.g. "SiLU").
+    """
+
+    def __init__(self, dim, norm_type="layer_norm", activation="SiLU"):
+        super().__init__()
+        layers = [nn.Linear(dim, dim)]
+        norm = get_norm(norm_type, dim)
+        if norm:
+            layers.append(norm)
+        layers.append(getattr(nn, activation)())
+        self.block = nn.Sequential(*layers)
+
+    def forward(self, x):
+        """Forward with additive residual connection.
+
+        Args:
+            x: Input tensor of shape ``(*, dim)``.
+
+        Returns:
+            Tensor of same shape with residual added.
+        """
+        return x + self.block(x)
+
+
+class ResidualMLP(nn.Module):
+    """MLP with stacked residual blocks between input and output projections.
+
+    Architecture: ``Linear(in->hidden) -> Norm -> Act -> [ResidualBlock]*depth
+    -> Linear(hidden->out)``. The residual blocks maintain gradient flow through
+    deep networks, while the input/output projections handle dimension changes.
+
+    Args:
+        input_dim: Size of the input feature vector.
+        hidden_dim: Width of all residual blocks.
+        output_dim: Size of the output feature vector.
+        depth: Number of stacked ``ResidualBlock`` layers.
+        norm: Normalization type passed to each block.
+        activation: Activation function name for all layers.
+    """
+
+    def __init__(
+        self, input_dim, hidden_dim, output_dim, depth, norm="layer_norm", activation="SiLU"
+    ):
+        super().__init__()
+
+        # Input projection
+        input_layers = [nn.Linear(input_dim, hidden_dim)]
+        norm_layer = get_norm(norm, hidden_dim)
+        if norm_layer:
+            input_layers.append(norm_layer)
+        input_layers.append(getattr(nn, activation)())
+        self.input_layer = nn.Sequential(*input_layers)
+
+        # Residual blocks
+        self.res_blocks = nn.Sequential(
+            *[
+                ResidualBlock(hidden_dim, norm_type=norm, activation=activation)
+                for _ in range(depth)
+            ]
+        )
+
+        # Output projection
+        self.output_layer = nn.Linear(hidden_dim, output_dim)
+
+    def forward(self, x):
+        """Forward through input projection, residual stack, and output projection.
+
+        Args:
+            x: Input tensor of shape ``(*, input_dim)``.
+
+        Returns:
+            Output tensor of shape ``(*, output_dim)``.
+        """
+        x = self.input_layer(x)
+        x = self.res_blocks(x)
+        return self.output_layer(x)
+
+
+class BaseModule(nn.Module):
+    """Config-driven network module that auto-builds layers from a dictionary spec.
+
+    Resolves input/output dimensions from observation dictionaries or explicit
+    overrides, then dispatches to a layer builder (MLP, CNN, GRU, ResidualMLP,
+    ResNet) based on ``module_config_dict.layer_config.type``. Temporal
+    dimensions are handled by flattening on input and reshaping on output.
+
+    Args:
+        obs_dim_dict: Mapping of observation name to its flat dimension.
+            Falls back to ``env_config.robot.algo_obs_dim_dict`` if None.
+        module_config_dict: Config with ``input_dim``, ``output_dim``, and
+            ``layer_config`` keys that drive network construction.
+        module_dim_dict: Mapping of named module outputs to their dimensions,
+            used to resolve symbolic references in ``input_dim``/``output_dim``.
+        env_config: Environment configuration (provides observation dims, camera
+            settings, and robot action dims).
+        algo_config: Algorithm configuration (stored but not used directly).
+        process_output_dim: If True, replace ``"robot_action_dim"`` sentinel
+            values in ``output_dim`` with the actual action dimension.
+        input_dim: Explicit input dimension override (skips calculation).
+        output_dim: Explicit output dimension override (skips calculation).
+        num_input_temporal_dims: If set, input's last two dims are flattened
+            (``temporal * feature -> input_dim``).
+        num_output_temporal_dims: If set, output is reshaped to
+            ``(*, num_output_temporal_dims, feature_per_step)``.
+    """
+
+    def __init__(
+        self,
+        obs_dim_dict=None,
+        module_config_dict=None,
+        module_dim_dict={},
+        env_config=None,
+        algo_config=None,
+        process_output_dim=False,
+        input_dim=None,
+        output_dim=None,
+        num_input_temporal_dims=None,
+        num_output_temporal_dims=None,
+    ):
+        super().__init__()
+
+        self.env_config = env_config
+        self.algo_config = algo_config
+        if obs_dim_dict is None:
+            if env_config is not None:
+                self.obs_dim_dict = env_config.robot.algo_obs_dim_dict
+        else:
+            self.obs_dim_dict = obs_dim_dict
+
+        self.module_config_dict = module_config_dict
+        if process_output_dim:
+            self.module_config_dict = self._process_module_config(
+                self.module_config_dict, self.env_config.robot.actions_dim
+            )
+
+        self.module_dim_dict = module_dim_dict
+
+        if input_dim is None:
+            self._calculate_input_dim()
+        else:
+            self.input_dim = input_dim
+        if output_dim is None:
+            self._calculate_output_dim()
+        else:
+            self.output_dim = output_dim
+        self.num_input_temporal_dims = num_input_temporal_dims
+        self.num_output_temporal_dims = num_output_temporal_dims
+        if num_input_temporal_dims is not None:
+            self.input_dim *= num_input_temporal_dims
+        if num_output_temporal_dims is not None:
+            self.output_dim *= num_output_temporal_dims
+
+        self._build_network_layer(self.module_config_dict.layer_config)
+
+    def _process_module_config(self, module_config_dict, num_actions):
+        """Replace ``"robot_action_dim"`` sentinels with the actual action count.
+
+        Args:
+            module_config_dict: Module config containing an ``output_dim`` list.
+            num_actions: Number of robot action dimensions to substitute.
+
+        Returns:
+            The mutated ``module_config_dict`` with sentinels replaced.
+        """
+        output_dim_list = module_config_dict["output_dim"]
+        if isinstance(output_dim_list, int):
+            output_dim_list = [output_dim_list]
+
+        for idx, output_dim in enumerate(output_dim_list):
+            if output_dim == "robot_action_dim":
+                module_config_dict["output_dim"][idx] = num_actions
+        return module_config_dict
+
+    def _calculate_input_dim(self):
+        """Calculate total input dimension by summing over ``module_config_dict["input_dim"]``.
+
+        Each entry is resolved as an observation name (looked up in
+        ``obs_dim_dict``), a numeric literal, or a named module output
+        (looked up in ``module_dim_dict``). Sets ``self.input_dim``.
+        """
+        # calculate input dimension based on the input specifications
+        input_dim = 0
+        for each_input in self.module_config_dict["input_dim"]:
+            if each_input in self.obs_dim_dict:
+                # atomic observation type
+                input_dim += self.obs_dim_dict[each_input]
+            elif isinstance(each_input, int | float):
+                # direct numeric input
+                input_dim += each_input
+            elif each_input in self.module_dim_dict:
+                input_dim += self.module_dim_dict[each_input]
+            else:
+                current_function_name = inspect.currentframe().f_code.co_name
+                raise ValueError(f"{current_function_name} - Unknown input type: {each_input}")
+
+        self.input_dim = input_dim
+
+    def _calculate_output_dim(self):
+        """Calculate total output dimension by summing over ``module_config_dict["output_dim"]``.
+
+        Each entry is resolved as a numeric literal or a named module output.
+        Sets ``self.output_dim``.
+        """
+        output_dim = 0
+        output_dim_list = self.module_config_dict["output_dim"]
+        if isinstance(output_dim_list, int) or isinstance(output_dim_list, str):
+            output_dim_list = [output_dim_list]
+
+        for each_output in output_dim_list:
+            if isinstance(each_output, int | float):
+                output_dim += each_output
+            elif each_output in self.module_dim_dict:
+                output_dim += self.module_dim_dict[each_output]
+            else:
+                current_function_name = inspect.currentframe().f_code.co_name
+                raise ValueError(f"{current_function_name} - Unknown output type: {each_output}")
+
+        self.output_dim = output_dim
+
+    def _build_network_layer(self, layer_config):
+        """Dispatch to the appropriate layer builder based on ``layer_config["type"]``.
+
+        Supported types: ``"MLP"``, ``"CNN"``, ``"GRU"``, ``"ResidualMLP"``,
+        ``"ResNet"``. The built network is stored as ``self.module``.
+
+        Args:
+            layer_config: Dict with a ``"type"`` key and type-specific params.
+        """
+        if layer_config["type"] == "MLP":
+            self._build_mlp_layer(layer_config)
+        elif layer_config["type"] == "CNN":
+            self._build_cnn_layer(layer_config)
+        elif layer_config["type"] == "GRU":
+            self._build_gru_layer(layer_config)
+        elif layer_config["type"] == "ResidualMLP":
+            self._build_residual_mlp_layer(layer_config)
+        elif layer_config["type"] == "ResNet":
+            self._build_resnet_layer(layer_config)
+        else:
+            raise NotImplementedError(f"Unsupported layer type: {layer_config['type']}")
+
+    def _build_mlp_layer(self, layer_config):
+        """Build a plain MLP with configurable hidden dims and activation.
+
+        Architecture: ``input -> [hidden_i -> Act]* -> output``. No residual
+        connections or normalization (use ``ResidualMLP`` for those).
+
+        Args:
+            layer_config: Dict with ``"hidden_dims"`` (list of ints) and
+                ``"activation"`` (nn.Module class name).
+        """
+        layers = []
+        hidden_dims = layer_config["hidden_dims"]
+        output_dim = self.output_dim
+        activation = getattr(nn, layer_config["activation"])()
+
+        layers.append(nn.Linear(self.input_dim, hidden_dims[0]))
+        layers.append(activation)
+
+        for l in range(len(hidden_dims)):
+            if l == len(hidden_dims) - 1:
+                layers.append(nn.Linear(hidden_dims[l], output_dim))
+            else:
+                layers.append(nn.Linear(hidden_dims[l], hidden_dims[l + 1]))
+                layers.append(activation)
+
+        self.module = nn.Sequential(*layers)
+
+    def _build_cnn_layer(self, layer_config):
+        """Build a CNN encoder from env camera config and layer specifications.
+
+        Constructs conv/pool layers from ``layer_config["layers"]``, resolves
+        input spatial dims and channel count from ``env_config.simulator``, and
+        appends a flatten + linear projection to ``self.output_dim``.
+
+        NOTE: Input is a flattened vision observation vector that gets reshaped
+        to ``(width, height, channels)`` based on camera config. The assertion
+        verifies the flat dimension matches the expected spatial size.
+
+        Args:
+            layer_config: Dict with ``"channel_dims"`` (list of ints),
+                ``"activation"``, and ``"layers"`` (list of layer dicts with
+                ``"type"`` of ``"conv"`` or ``"pool"``).
+        """
+        layers = []
+        channel_dims = layer_config["channel_dims"]
+        activation = getattr(nn, layer_config["activation"])()
+
+        # Get input dimensions from env_config camera settings
+        camera_config = self.env_config.simulator.config.cameras
+        input_height = camera_config.camera_resolutions[0]
+        input_width = camera_config.camera_resolutions[1]
+
+        # Determine number of channels from camera types
+        input_channels = 0
+        for camera_type in camera_config.camera_types:
+            if camera_type.get("rgb", False):
+                input_channels += 3
+            if camera_type.get("depth", False):
+                input_channels += 1
+
+        # If no channels found, default to 1
+        if input_channels == 0:
+            input_channels = 1
+
+        vision_obs_dim = [input_width, input_height, input_channels]
+        print("vision_obs_dim", vision_obs_dim)
+        assert (
+            vision_obs_dim[0] * vision_obs_dim[1] * vision_obs_dim[2]
+            == self.obs_dim_dict["vision_obs"]
+        )
+        if len(vision_obs_dim) != 3:
+            raise ValueError(
+                f"vision_obs dimension should be (width, height, channels), got {vision_obs_dim}"
+            )
+        input_width, input_height, input_channels = vision_obs_dim
+
+        # Get layer configurations
+        layer_configs = layer_config.get("layers", [])
+        use_batch_norm = layer_config.get("norm_config", {}).get("use_batch_norm", False)
+
+        # Track spatial dimensions and channels
+        current_height, current_width = input_height, input_width
+        current_channels = input_channels
+        conv_idx = 0  # Track which conv layer we're on for channel dimensions
+
+        for layer_cfg in layer_configs:
+            layer_type = layer_cfg["type"]
+
+            if layer_type == "conv":
+                # Get conv parameters
+                kernel_size = layer_cfg.get("kernel_size", 3)
+                stride = layer_cfg.get("stride", 1)
+                padding = layer_cfg.get("padding", 1)
+
+                # Determine output channels
+                if conv_idx < len(channel_dims):
+                    out_channels = channel_dims[conv_idx]
+                else:
+                    out_channels = self.output_dim
+
+                # Add conv layer
+                layers.append(
+                    nn.Conv2d(
+                        current_channels,
+                        out_channels,
+                        kernel_size=kernel_size,
+                        stride=stride,
+                        padding=padding,
+                    )
+                )
+
+                if use_batch_norm:
+                    layers.append(nn.BatchNorm2d(out_channels))
+                layers.append(activation)
+
+                # Update dimensions
+                current_channels = out_channels
+                current_height = (current_height - kernel_size + 2 * padding) // stride + 1
+                current_width = (current_width - kernel_size + 2 * padding) // stride + 1
+                conv_idx += 1
+
+            elif layer_type == "pool":
+                # Get pool parameters
+                kernel_size = layer_cfg.get("kernel_size", 2)
+                stride = layer_cfg.get("stride", 2)
+
+                # Add pooling layer if dimensions allow
+                if current_height >= kernel_size and current_width >= kernel_size:
+                    layers.append(nn.MaxPool2d(kernel_size=kernel_size, stride=stride))
+                    current_height = current_height // stride
+                    current_width = current_width // stride
+
+        # Add global average pooling if spatial dimensions are too small
+        # if current_height * current_width > 1:
+        #     # import ipdb; ipdb.set_trace()
+        #     layers.append(nn.AdaptiveAvgPool2d(1))
+
+        layers.append(nn.Flatten())
+
+        layers.append(nn.Linear(current_channels * current_height * current_width, self.output_dim))
+
+        self.module = nn.Sequential(*layers)
+
+    def forward_without_hidden_state(self, input):
+        """Forward pass for stateless modules (MLP, CNN, ResidualMLP, ResNet).
+
+        Args:
+            input: Input tensor of shape ``(batch, input_dim)``.
+
+        Returns:
+            Output tensor of shape ``(batch, output_dim)``.
+        """
+        return self.module(input)
+
+    def forward_with_hidden_state(self, input, hidden_state):
+        """Forward pass for recurrent modules (GRU).
+
+        Args:
+            input: Input tensor of shape ``(batch, seq_len, input_dim)``.
+            hidden_state: Previous hidden state tensor.
+
+        Returns:
+            Tuple of (output, updated_hidden_state).
+        """
+        # import ipdb; ipdb.set_trace()
+        output, hidden_state = self.module(input, hidden_state)
+        return output, hidden_state
+
+    # def forward(self, input, hidden_state=None):
+    #     if hidden_state is None:
+    #         return self.forward_without_hidden_state(input)
+    #     else:
+    #         return self.forward_with_hidden_state(input, hidden_state)
+
+    def _build_gru_layer(self, layer_config):
+        """Build a GRU recurrent layer.
+
+        Args:
+            layer_config: Dict with ``"hidden_dim"`` and ``"num_layers"``.
+        """
+        self.module = nn.GRU(
+            input_size=self.input_dim,
+            hidden_size=layer_config["hidden_dim"],
+            num_layers=layer_config["num_layers"],
+            batch_first=True,
+        )
+
+    def _build_resnet_layer(self, layer_config):
+        """Build a torchvision ResNet backbone with global avg pool and linear head.
+
+        Removes the original classification head (avgpool + fc) and replaces
+        it with ``AdaptiveAvgPool2d(1) -> Flatten -> Linear(feat, output_dim)``.
+
+        Args:
+            layer_config: Dict with ``"resnet_type"`` (e.g. ``"resnet18"``),
+                ``"pretrained"`` (bool), and ``"trainable"`` (bool, freezes
+                backbone params when False).
+        """
+        print("Building ResNet layer")
+        resnet_type = layer_config.get("resnet_type", "resnet18")  # Default to resnet18
+        pretrained = layer_config.get("pretrained", True)
+        trainable = layer_config.get("trainable", True)
+
+        if resnet_type == "resnet18":
+            resnet = models.resnet18(pretrained=pretrained)
+        elif resnet_type == "resnet34":
+            resnet = models.resnet34(pretrained=pretrained)
+        elif resnet_type == "resnet50":
+            resnet = models.resnet50(pretrained=pretrained)
+        elif resnet_type == "resnet101":
+            resnet = models.resnet101(pretrained=pretrained)
+        elif resnet_type == "resnet152":
+            resnet = models.resnet152(pretrained=pretrained)
+        else:
+            raise ValueError(f"Unsupported ResNet type: {resnet_type}")
+
+        resnet_features = nn.Sequential(*list(resnet.children())[:-2])  # Remove avgpool and fc
+
+        if resnet_type in ["resnet18", "resnet34"]:
+            resnet_feature_dim = 512
+        else:  # resnet50, resnet101, resnet152
+            resnet_feature_dim = 2048
+
+        # Freeze ResNet parameters if not trainable
+        if not trainable:
+            for param in resnet_features.parameters():
+                param.requires_grad = False
+
+        # Add a final linear layer to match output_dim
+        layers = [
+            resnet_features,
+            nn.AdaptiveAvgPool2d(1),  # Global average pooling
+            nn.Flatten(),
+            nn.Linear(resnet_feature_dim, self.output_dim),
+        ]
+
+        self.module = nn.Sequential(*layers)
+
+    def _build_residual_mlp_layer(self, layer_config):
+        """Build a ``ResidualMLP`` from layer config.
+
+        Args:
+            layer_config: Dict with ``"hidden_dim"``, ``"depth"``, and
+                optional ``"norm"`` and ``"activation"``.
+        """
+        self.module = ResidualMLP(
+            input_dim=self.input_dim,
+            hidden_dim=layer_config["hidden_dim"],
+            output_dim=self.output_dim,
+            depth=layer_config["depth"],
+            norm=layer_config.get("norm", "layer_norm"),
+            activation=layer_config.get("activation", "SiLU"),
+        )
+
+    def forward(self, input, **kwargs):
+        """Forward pass with automatic temporal dim handling.
+
+        Accepts either a tensor or a dict keyed by observation name. When
+        ``num_input_temporal_dims`` is set, flattens the last two dims before
+        the network. When ``num_output_temporal_dims`` is set, reshapes the
+        output to ``(*, num_output_temporal_dims, feature_per_step)``.
+
+        Args:
+            input: Tensor of shape ``(batch, input_dim)`` or
+                ``(batch, temporal, feature)``, or a dict mapping observation
+                names to tensors (first key in ``input_dim`` config is used).
+            **kwargs: Passed through (unused by base; allows subclass compat).
+
+        Returns:
+            Output tensor of shape ``(batch, output_dim)`` or
+            ``(batch, num_output_temporal_dims, feature_per_step)``.
+        """
+        if isinstance(input, dict):
+            input_obs_key = self.module_config_dict["input_dim"][0]
+            input = input[input_obs_key]
+        if self.num_input_temporal_dims is not None:
+            input = input.view(*input.shape[:-2], self.input_dim)
+        output = self.module(input)
+        if self.num_output_temporal_dims is not None:
+            output = output.view(
+                *output.shape[:-1],
+                self.num_output_temporal_dims,
+                self.output_dim // self.num_output_temporal_dims,
+            )
+        return output
+
+
+class BaseModuleAux(BaseModule):
+    """BaseModule extended with auxiliary loss computation.
+
+    Wraps the parent forward pass so that, when requested, auxiliary loss
+    functions are evaluated on the network output and returned alongside
+    their coefficients. This enables the PPO aux-loss trainer to add
+    regularization terms (e.g. action smoothness) without modifying the
+    core module.
+
+    When ``compute_aux_loss=True``, returns a dict with ``"action_mean"``,
+    ``"aux_losses"``, and ``"aux_loss_coef"``. Otherwise returns a plain
+    tensor identical to ``BaseModule.forward``.
+    """
+
+    def __init__(self, aux_loss_func={}, aux_loss_coef={}, **kwargs):
+        """Initialize with auxiliary loss functions and their coefficients.
+
+        Args:
+            aux_loss_func: Mapping of loss name to either an ``nn.Module``
+                instance or a Hydra config dict (instantiated via
+                ``hydra.utils.instantiate``). Each function receives a dict
+                with an ``"action_mean"`` key.
+            aux_loss_coef: Mapping of loss name to its scalar coefficient,
+                returned alongside computed losses for the trainer to weight.
+            **kwargs: Forwarded to ``BaseModule.__init__``.
+        """
+        super().__init__(**kwargs)
+        self.aux_loss_coef = aux_loss_coef
+
+        # Instantiate loss functions
+        self.aux_loss_func = nn.ModuleDict()
+        for name, func_cfg in aux_loss_func.items():
+            if isinstance(func_cfg, nn.Module):
+                self.aux_loss_func[name] = func_cfg
+            else:
+                # Hydra instantiation
+                from hydra.utils import instantiate
+
+                self.aux_loss_func[name] = instantiate(func_cfg)
+
+    def forward(self, input, compute_aux_loss=False, **kwargs):
+        """Forward pass with optional auxiliary loss computation.
+
+        Args:
+            input: Input tensor or dict (same as ``BaseModule.forward``).
+            compute_aux_loss: If True, evaluate all registered aux loss
+                functions and return a dict instead of a plain tensor.
+            **kwargs: Forwarded to ``BaseModule.forward``.
+
+        Returns:
+            When ``compute_aux_loss=False``: output tensor of shape
+            ``(batch, output_dim)``.
+            When ``compute_aux_loss=True``: dict with keys
+            ``"action_mean"`` (tensor), ``"aux_losses"`` (dict of scalar
+            tensors), and ``"aux_loss_coef"`` (dict of floats).
+        """
+        # Call parent forward to get output tensor
+        output = super().forward(input, **kwargs)
+
+        if not compute_aux_loss:
+            return output
+
+        # Build loss inputs dict
+        loss_inputs = {"action_mean": output}
+
+        # Compute each loss
+        aux_losses = {}
+        for name, func in self.aux_loss_func.items():
+            aux_losses[name] = func(loss_inputs)
+
+        return {
+            "action_mean": output,
+            "aux_losses": aux_losses,
+            "aux_loss_coef": self.aux_loss_coef,
+        }
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/modules/data_utils.py b/GR00T-WholeBodyControl/gear_sonic/trl/modules/data_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..dc2ee505327d294044108594b793bc4ef457569e
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/modules/data_utils.py
@@ -0,0 +1,185 @@
+# Partially based on NVLabs ProtoMotions (Apache 2.0):
+# https://github.com/NVlabs/ProtoMotions/blob/94059259ba2b596bf908828cc04e8fc6ff901114/phys_anim/agents/utils/data_utils.py
+"""Data-management utilities for on-policy PPO rollouts.
+
+Provides :class:`RolloutStorage`, a flexible ``nn.Module``-based buffer that
+stores arbitrary per-transition tensors (observations, actions, rewards,
+values, …) and exposes them as randomised mini-batches for PPO training.
+"""
+
+import torch
+from torch import Tensor, nn
+
+
+class RolloutStorage(nn.Module):
+    """On-policy rollout buffer for PPO with dynamic key registration.
+
+    Stores up to ``num_transitions_per_env`` steps for each of the
+    ``num_envs`` parallel environments.  Keys (e.g. "obs", "actions",
+    "rewards") are registered on demand via :meth:`register_key`, which
+    allocates a zero-filled ``nn.Module`` buffer of the appropriate shape so
+    that the storage follows ``.to(device)`` calls.
+
+    Args:
+        num_envs: Number of parallel simulation environments.
+        num_transitions_per_env: Maximum rollout length (horizon).
+        device: PyTorch device string for all allocated buffers.
+    """
+
+    def __init__(self, num_envs, num_transitions_per_env, device="cpu"):
+
+        super().__init__()
+
+        self.device = device
+
+        self.num_transitions_per_env = num_transitions_per_env
+        self.num_envs = num_envs
+
+        # rnn
+        # self.saved_hidden_states_a = None
+        # self.saved_hidden_states_c = None
+
+        self.step = 0
+        self.stored_keys = []
+
+    def register_key(self, key: str, shape=(), dtype=torch.float):
+        """Allocate and register a new storage buffer for the given key.
+
+        The buffer has shape ``(num_transitions_per_env, num_envs, *shape)``
+        and is registered as a non-persistent ``nn.Module`` buffer so it moves
+        with ``.to(device)`` but is not saved in ``state_dict``.
+
+        Args:
+            key: Unique name for the data field (e.g. ``"obs"``, ``"rewards"``).
+            shape: Per-transition, per-environment shape tuple.
+            dtype: Tensor dtype for the buffer.
+
+        Raises:
+            AssertionError: If ``key`` is already registered or ``shape`` is
+                not a list or tuple.
+        """
+        # This class was partially copied from https://github.com/NVlabs/ProtoMotions/blob/94059259ba2b596bf908828cc04e8fc6ff901114/phys_anim/agents/utils/data_utils.py
+        assert not hasattr(self, key), key
+        assert isinstance(shape, list | tuple), f"shape must be a list or tuple, got {type(shape)}"
+        buffer = torch.zeros(
+            (self.num_transitions_per_env, self.num_envs) + shape, dtype=dtype, device=self.device
+        )
+        self.register_buffer(key, buffer, persistent=False)
+        self.stored_keys.append(key)
+
+    def increment_step(self):
+        """Advance the write cursor by one transition step."""
+        self.step += 1
+
+    def update_key(self, key: str, data: Tensor):
+        """Write ``data`` into the buffer at the current step for ``key``.
+
+        Args:
+            key: Previously registered buffer key.
+            data: Tensor of shape ``(num_envs, *key_shape)``.  Must not
+                require gradients.
+
+        Raises:
+            AssertionError: If ``data.requires_grad`` is True or the buffer
+                is full (step >= num_transitions_per_env).
+        """
+        # This class was partially copied from https://github.com/NVlabs/ProtoMotions/blob/94059259ba2b596bf908828cc04e8fc6ff901114/phys_anim/agents/utils/data_utils.py
+        assert not data.requires_grad
+        assert self.step < self.num_transitions_per_env, "Rollout buffer overflow"
+        getattr(self, key)[self.step].copy_(data)
+
+    def batch_update_data(self, key: str, data: Tensor):
+        """Overwrite the entire buffer for ``key`` with ``data``.
+
+        Useful for writing pre-computed values (e.g. advantages, returns)
+        after a full rollout has been collected.
+
+        Args:
+            key: Previously registered buffer key.
+            data: Tensor of shape
+                ``(num_transitions_per_env, num_envs, *key_shape)``.  Must
+                not require gradients.
+        """
+        # This class was partially copied from https://github.com/NVlabs/ProtoMotions/blob/94059259ba2b596bf908828cc04e8fc6ff901114/phys_anim/agents/utils/data_utils.py
+        assert not data.requires_grad
+        getattr(self, key)[:] = data
+        # self.store_dict[key] += self.total_sum()
+
+    def _save_hidden_states(self, hidden_states):
+        assert NotImplementedError
+        if hidden_states is None or hidden_states == (None, None):
+            return
+        # make a tuple out of GRU hidden state sto match the LSTM format
+        hid_a = hidden_states[0] if isinstance(hidden_states[0], tuple) else (hidden_states[0],)
+        hid_c = hidden_states[1] if isinstance(hidden_states[1], tuple) else (hidden_states[1],)
+
+        # initialize if needed
+        if self.saved_hidden_states_a is None:
+            self.saved_hidden_states_a = [
+                torch.zeros(self.observations.shape[0], *hid_a[i].shape, device=self.device)
+                for i in range(len(hid_a))
+            ]
+            self.saved_hidden_states_c = [
+                torch.zeros(self.observations.shape[0], *hid_c[i].shape, device=self.device)
+                for i in range(len(hid_c))
+            ]
+        # copy the states
+        for i in range(len(hid_a)):
+            self.saved_hidden_states_a[i][self.step].copy_(hid_a[i])
+            self.saved_hidden_states_c[i][self.step].copy_(hid_c[i])
+
+    def clear(self):
+        """Reset the write cursor to the start of the buffer."""
+        self.step = 0
+
+    def get_statistics(self):
+        """Return buffer statistics (not implemented)."""
+        raise NotImplementedError
+
+    def query_key(self, key: str):
+        """Return the full buffer tensor for ``key``.
+
+        Args:
+            key: Previously registered buffer key.
+
+        Returns:
+            Tensor of shape ``(num_transitions_per_env, num_envs, *key_shape)``.
+
+        Raises:
+            AssertionError: If ``key`` has not been registered.
+        """
+        assert hasattr(self, key), key
+        return getattr(self, key)
+
+    def mini_batch_generator(self, num_mini_batches, num_epochs=8):
+        """Yield randomly shuffled mini-batches over all stored transitions.
+
+        The full buffer (flattened across transitions and environments) is
+        shuffled once per call and then sliced into ``num_mini_batches``
+        equal-sized chunks, repeated ``num_epochs`` times.
+
+        Args:
+            num_mini_batches: Number of mini-batches per epoch.
+            num_epochs: Number of times to iterate over the shuffled data.
+
+        Yields:
+            Dict mapping each registered key to a mini-batch tensor of shape
+            ``(mini_batch_size, *key_shape)``.
+        """
+        batch_size = self.num_envs * self.num_transitions_per_env
+        mini_batch_size = batch_size // num_mini_batches
+        indices = torch.randperm(
+            num_mini_batches * mini_batch_size, requires_grad=False, device=self.device
+        )
+
+        _buffer_dict = {key: getattr(self, key)[:].flatten(0, 1) for key in self.stored_keys}
+
+        for epoch in range(num_epochs):
+            for i in range(num_mini_batches):
+
+                start = i * mini_batch_size
+                end = (i + 1) * mini_batch_size
+                batch_idx = indices[start:end]
+
+                _batch_buffer_dict = {key: _buffer_dict[key][batch_idx] for key in self.stored_keys}
+                yield _batch_buffer_dict
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/modules/universal_token_modules.py b/GR00T-WholeBodyControl/gear_sonic/trl/modules/universal_token_modules.py
new file mode 100644
index 0000000000000000000000000000000000000000..2755d5b73ac5db0490e370d9f489c00564d2c00f
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/modules/universal_token_modules.py
@@ -0,0 +1,1256 @@
+from __future__ import annotations
+
+from loguru import logger
+import numpy as np
+import torch
+from torch import nn
+
+from gear_sonic.trl.utils import common
+
+
+def set_fuzzy_config_params(config, candidate_keys, value):
+    """Set matching config keys to a given value.
+
+    Scans ``config`` for any key that appears in ``candidate_keys`` and
+    overwrites its value.  Used to wire dynamically-computed dimensions
+    (e.g. ``input_dim``, ``output_dim``) into encoder / decoder param
+    dicts before instantiation.
+
+    Args:
+        config: Mutable mapping of parameter names to values.
+        candidate_keys: Iterable of key names to look for in ``config``.
+        value: Replacement value to assign to each matched key.
+
+    Returns:
+        The same ``config`` mapping with matching keys updated in-place.
+    """
+    for key in candidate_keys:
+        if key in config:
+            config[key] = value
+    return config
+
+
+class UniversalTokenModule(nn.Module):
+    """SONIC-style action transform module (ATM) with FSQ token bottleneck.
+
+    Implements the encoder → FSQ quantizer → decoder pipeline described in the
+    SONIC paper.  Multiple named encoders (e.g. ``g1``, ``smpl``, ``teleop``)
+    convert different flavors of tokenizer observations into a shared latent
+    space.  An optional Finite Scalar Quantizer (FSQ) discretises the latent,
+    producing a compact token representation.  One or more decoders then map
+    the tokens (plus proprioception) back to joint-space actions.
+
+    Token flow::
+
+        tokenizer_obs ──► encoder(s) ──► [+ additive encoders]
+                                      ──► [latent_residual (pre_quantization)]
+                                      ──► FSQ quantizer
+                                      ──► [+ latent_residual (post_quantization)]
+                                      ──► decoder(s)
+                                      ──► action_mean
+
+    Latent residual modes allow an external HOI policy to inject corrections
+    into the latent space without retraining the base ATM:
+
+    * ``"post_quantization"`` - residual added to FSQ-quantized tokens (default).
+    * ``"pre_quantization"`` - residual added *before* FSQ; the sum gets quantized.
+    * ``"pre_quantization_replace"`` - latent is replaced entirely by the residual.
+
+    Attributes:
+        encoders: ``nn.ModuleDict`` of named encoder networks.
+        quantizer: FSQ quantizer (``None`` when quantization is disabled).
+        decoders: ``nn.ModuleDict`` of named decoder networks.
+        aux_loss_func: ``nn.ModuleDict`` of auxiliary loss callables.
+        token_dim: Dimensionality of a single token (equals ``num_fsq_levels``).
+        max_num_tokens: Number of tokens produced per timestep.
+        token_total_dim: ``token_dim * max_num_tokens`` - flat token size.
+    """
+
+    def __init__(
+        self,
+        env_config,
+        algo_config,
+        obs_dim_dict=None,
+        module_dim_dict=None,
+        proprioception_features=[],
+        num_fsq_levels=5,
+        fsq_level_list=16,
+        max_num_tokens=None,
+        down_t=2,
+        num_future_frames=1,
+        quantizer=None,
+        encoders=None,
+        decoders=None,
+        aux_loss_func={},
+        aux_loss_coef={},
+        encoder_sample_probs=None,
+        reencode_smpl_g1_recon=False,
+        meta_action_dim=None,  # For hierarchical policies with split actions
+        body_action_dim=None,  # For separate body decoder (e.g., 29 body joints)
+        hand_action_dim=None,  # For separate hand decoder (e.g., 14 hand joints)
+        freeze_encoders=False,
+        freeze_decoders=False,
+        freeze_quantizer=False,
+        stiff_compliance_threshold=0.01,  # Threshold for stiff mode filtering
+        optimize_encoders_ratio_for_CHIP=False,  # CHIP compliance training optimization
+        active_encoders=None,  # Optional list of encoder names to activate (None = all)
+        active_decoders=None,  # Optional list of decoder names to activate (None = all)
+        **kwargs,  # noqa: ARG002
+    ):
+        """Initialise encoders, FSQ quantizer, decoders, and auxiliary losses.
+
+        Args:
+            env_config: Environment configuration object exposing
+                ``obs.group_obs_dims``, ``obs.group_obs_names``, and
+                ``robot.actions_dim``.
+            algo_config: Algorithm configuration (stored but not directly
+                accessed during construction).
+            obs_dim_dict: Mapping from observation name to dimension.  Falls
+                back to ``env_config.robot.algo_obs_dim_dict`` when ``None``.
+            module_dim_dict: Optional extra dimension overrides for named
+                intermediate features.
+            proprioception_features: List of observation keys whose concatenation
+                forms the proprioception input fed to every decoder.
+            num_fsq_levels: Number of FSQ levels (equals ``token_dim``).
+            fsq_level_list: Per-level codebook size.  An ``int`` is broadcast
+                to ``[fsq_level_list] * num_fsq_levels``.
+            max_num_tokens: Explicit token count per timestep.  When ``None``
+                this is derived as ``max(1, num_future_frames // 2**down_t)``.
+            down_t: Temporal downsampling factor for token count derivation.
+            num_future_frames: Number of future motion frames in the tokenizer
+                input window.
+            quantizer: Hydra-instantiable config for the FSQ quantizer.
+                ``None`` disables quantization (identity passthrough).
+            encoders: Dict of encoder configs (Hydra DictConfig).  Each entry
+                specifies ``inputs``, ``outputs``, ``params``, and optional
+                ``additive_to`` / ``sub_encoders`` / ``mask`` / ``freeze``.
+            decoders: Dict of decoder configs (Hydra DictConfig).  Each entry
+                specifies ``inputs``, ``outputs``, ``conds``, ``params``, and
+                ``has_temporal_dim``.
+            aux_loss_func: Dict of auxiliary loss configs to instantiate as
+                ``nn.Module`` callables (stored in ``self.aux_loss_func``).
+            aux_loss_coef: Dict mapping loss name → scalar coefficient passed
+                back to the trainer.
+            encoder_sample_probs: Ordered dict mapping encoder name → sampling
+                probability; defines the column order of ``encoder_index`` in
+                tokenizer observations.
+            reencode_smpl_g1_recon: When ``True``, re-encode the G1 kinematic
+                decoder output to compute cycle-consistency auxiliary losses.
+            meta_action_dim: Action dimensionality for hierarchical policies
+                where the top-level action differs from the full joint count.
+            body_action_dim: Dimensionality of body joints for split
+                body / hand decoding (default 29).
+            hand_action_dim: Dimensionality of hand joints for split
+                body / hand decoding (default 14).
+            freeze_encoders: Freeze all encoder parameters after init.
+            freeze_decoders: Freeze all decoder parameters after init.
+            freeze_quantizer: Freeze quantizer parameters after init.
+            stiff_compliance_threshold: Absolute compliance value below which
+                an environment is treated as "stiff" for auxiliary loss gating.
+            optimize_encoders_ratio_for_CHIP: When ``True``, enforce one-hot
+                encoder selection (CHIP training mode) instead of the default
+                multi-hot selection where SMPL-native envs activate both SMPL
+                and G1 encoders.
+            active_encoders: Optional list of encoder names to instantiate and
+                run.  All encoders are active when ``None``.
+            active_decoders: Optional list of decoder names to run during
+                forward.  All decoders are run when ``None``.
+            **kwargs: Absorbed for forward-compatibility.
+        """
+        super().__init__()
+
+        # Store freeze flags
+        self.freeze_encoders = freeze_encoders
+        self.freeze_decoders = freeze_decoders
+        self.stiff_compliance_threshold = stiff_compliance_threshold
+        self.freeze_quantizer = freeze_quantizer
+        self.optimize_encoders_ratio_for_CHIP = optimize_encoders_ratio_for_CHIP
+
+        if self.optimize_encoders_ratio_for_CHIP:
+            logger.info(
+                "[CHIP Mode] optimize_encoders_ratio_for_CHIP=True: "
+                "Native encoder selection is now one-hot (no G1 auto-activation for SMPL). "
+                "G1 latents for SMPL-native envs computed only in aux losses when compliance=0."
+            )
+
+        # Cache for last encoded tokens (for external access, e.g., by callbacks)
+        # These are populated during forward() and can be read afterwards
+        self._last_encoded_tokens = None  # dict[encoder_name -> Tensor]
+        self._last_encoded_latents = None  # dict[encoder_name -> Tensor] (pre-quantization)
+
+        self.env_config = env_config
+        self.algo_config = algo_config
+        self.encoders_cfg = encoders
+        self.decoders_cfg = decoders
+        self.encoder_sample_probs = encoder_sample_probs
+        self.reencode_smpl_g1_recon = reencode_smpl_g1_recon
+        self.tokenizer_obs_dims = self.env_config.obs.group_obs_dims["tokenizer"]
+        self.tokenizer_obs_names = self.env_config.obs.group_obs_names["tokenizer"]
+        # The tokenizer observation schema is fixed after module construction.
+        tokenizer_obs_specs = []
+        tokenizer_obs_offset = 0
+        for name in self.tokenizer_obs_names:
+            dims = tuple(self.tokenizer_obs_dims[name])
+            flat_dim = int(np.prod(dims))
+            tokenizer_obs_specs.append(
+                (name, tokenizer_obs_offset, tokenizer_obs_offset + flat_dim, dims)
+            )
+            tokenizer_obs_offset += flat_dim
+        self.tokenizer_obs_specs = tuple(tokenizer_obs_specs)
+        self.tokenizer_obs_total_dim = tokenizer_obs_offset
+        self.actions_dim = self.env_config.robot.actions_dim
+        self.meta_action_dim = (
+            meta_action_dim  # Can be different from actions_dim for hierarchical policies
+        )
+
+        if obs_dim_dict is None:
+            obs_dim_dict = getattr(env_config.robot, "algo_obs_dim_dict", {})
+
+        if module_dim_dict is None:
+            module_dim_dict = {}
+
+        self.obs_dim_dict = obs_dim_dict
+        self.module_dim_dict = module_dim_dict
+        self.proprioception_features = proprioception_features
+        self.num_future_frames = num_future_frames
+
+        # Initialize auxiliary loss functions (nn.ModuleDict so sub-module buffers
+        # are part of the model state_dict and visible to DDP)
+        self.aux_loss_func = nn.ModuleDict()
+        for name, loss_func in aux_loss_func.items():
+            self.aux_loss_func[name] = common.custom_instantiate(loss_func, _resolve=False)
+        self.aux_loss_coef = aux_loss_coef
+
+        module_input_candidate_params = ["input_dim"]
+
+        module_output_candidate_params = ["output_dim"]
+
+        """
+        Initialize quantizer
+        """
+        if isinstance(fsq_level_list, int):
+            fsq_level_list = [fsq_level_list] * num_fsq_levels
+        if quantizer is not None:
+            self.quantizer = common.custom_instantiate(
+                quantizer, levels=fsq_level_list, _resolve=False
+            )
+        else:
+            self.quantizer = None
+        self.num_fsq_levels = num_fsq_levels
+        self.fsq_level_list = fsq_level_list
+
+        # quantizer levels determine the embedding dimension
+        self.token_dim = self.num_fsq_levels
+        self.down_t = down_t  # default
+        if max_num_tokens is not None:
+            self.max_num_tokens = max_num_tokens
+        else:
+            self.max_num_tokens = max(1, self.num_future_frames // (2**self.down_t))
+        self.token_total_dim = self.token_dim * self.max_num_tokens
+        logger.info(
+            f"Motion Encoder and Quantizer initialized with embedding dim: {self.token_total_dim} (num_tokens={self.max_num_tokens}, token_dim={self.token_dim})"  # noqa: E501
+        )
+
+        """
+        Initialize motion encoders dynamically
+        """
+        self.encoders = nn.ModuleDict()
+        self.sub_encoders = {}
+        self.encoder_input_features = {}
+        self.encoder_mask_features = {}
+        self.encoders_to_iterate = []
+        self.additive_encoders = {}  # Map: base_encoder_name -> list of additive encoder names
+
+        # First pass: identify additive encoders
+        for encoder_name, encoder_config in self.encoders_cfg.items():
+            additive_to = encoder_config.get("additive_to", None)
+            if additive_to is not None:
+                if additive_to not in self.additive_encoders:
+                    self.additive_encoders[additive_to] = []
+                self.additive_encoders[additive_to].append(encoder_name)
+
+        # New config structure - initialize all encoders from config
+        for encoder_name, encoder_config in self.encoders_cfg.items():
+            # Skip non-active encoders early (avoids KeyError on missing obs dims)
+            if active_encoders is not None and encoder_name not in active_encoders:
+                continue
+
+            sub_encoder_only = encoder_config.get("sub_encoder_only", False)
+            is_additive = encoder_config.get("additive_to", None) is not None
+            # Don't iterate additive encoders independently - they're called within their base encoder
+            if not sub_encoder_only and not is_additive:
+                self.encoders_to_iterate.append(encoder_name)
+            self.sub_encoders[encoder_name] = encoder_config.get("sub_encoders", [])
+
+            input_features = encoder_config.get("inputs", [])
+            output_features = encoder_config.get("outputs", [])
+            input_feature_dim = sum([self.tokenizer_obs_dims[key][-1] for key in input_features])
+            if len(output_features) > 0:
+                output_feature_dim = sum(
+                    [self.tokenizer_obs_dims[key][-1] for key in output_features]
+                )
+            else:
+                output_feature_dim = self.token_dim
+            self.encoder_input_features[encoder_name] = input_features
+            self.encoder_mask_features[encoder_name] = list(encoder_config.get("mask", []))
+
+            if len(self.sub_encoders[encoder_name]) > 0:
+                continue
+
+            encoder_params = encoder_config["params"].copy()
+            # Set dynamic parameters
+            set_fuzzy_config_params(
+                encoder_params, module_input_candidate_params, input_feature_dim
+            )
+            set_fuzzy_config_params(
+                encoder_params, module_output_candidate_params, output_feature_dim
+            )
+
+            # Instantiate encoder
+            encoder = common.custom_instantiate(encoder_params, _resolve=False)
+            self.encoders[encoder_name] = encoder
+
+            # Per-encoder freeze support
+            if encoder_config.get("freeze", False):
+                for param in encoder.parameters():
+                    param.requires_grad = False
+                logger.info(f"Froze encoder: {encoder_name}")
+
+            logger.info(f"Initialized {encoder_name} encoder with input features: {input_features}")
+
+        # Log additive encoder relationships
+        if self.additive_encoders:
+            for base_encoder, additive_list in self.additive_encoders.items():
+                logger.info(
+                    f"Additive encoders for '{base_encoder}': {additive_list} (outputs will be summed)"
+                )
+
+        """
+        Initialize motion decoders dynamically
+        """
+        self.decoders = nn.ModuleDict()
+        self.decoder_input_features = {}
+        self.decoder_output_features = {}
+        self.decoder_output_feature_dims = {}
+        self.decoder_cond_features = {}  # For root-disentangled decoders (external_cond)
+        self.decoder_mask_features = {}
+
+        # Support custom meta_action_dim for hierarchical policies
+        meta_action_dim = getattr(self, "meta_action_dim", None) or self.actions_dim
+
+        # Compute proprioception dim from actual features (not hardcoded actor_obs)
+        proprioception_dim = sum(
+            self.obs_dim_dict[key]
+            for key in self.proprioception_features
+            if key in self.obs_dim_dict
+        ) or self.obs_dim_dict.get("actor_obs", 0)
+
+        decoder_feature_dims_map = {
+            "proprioception": proprioception_dim,
+            "token": self.token_dim,
+            "token_flattened": self.token_dim * self.max_num_tokens,
+            "action": self.actions_dim,
+            "meta_action": meta_action_dim,  # For hierarchical policies
+            "hand_action": hand_action_dim or 14,  # 14 hand joints default
+            "body_action": body_action_dim or 29,  # 29 body joints default
+        }
+        for k, v in self.tokenizer_obs_dims.items():
+            if k not in decoder_feature_dims_map:
+                decoder_feature_dims_map[k] = v[-1]
+        self.decoder_feature_dims_map = decoder_feature_dims_map
+
+        for decoder_name, decoder_config in self.decoders_cfg.items():
+            # Skip non-active decoders early (avoids KeyError on missing obs dims)
+            if active_decoders is not None and decoder_name not in active_decoders:
+                continue
+
+            decoder_params = decoder_config["params"].copy()
+            input_features = decoder_config.get("inputs", [])
+            output_features = decoder_config.get("outputs", [])
+            cond_features = decoder_config.get("conds", [])
+            has_temporal_dim = decoder_config.has_temporal_dim
+            input_feature_dim = 0
+            for key in input_features:
+                input_feature_dim += self.decoder_feature_dims_map[key]
+            output_feature_dim = 0
+            self.decoder_output_feature_dims[decoder_name] = {}
+            for key in output_features:
+                if key in decoder_feature_dims_map:
+                    feature_dim = decoder_feature_dims_map[key]
+                else:
+                    assert has_temporal_dim
+                    feature_dim = self.tokenizer_obs_dims[key][-1]
+                self.decoder_output_feature_dims[decoder_name][key] = feature_dim
+                output_feature_dim += feature_dim
+            set_fuzzy_config_params(
+                decoder_params, module_input_candidate_params, input_feature_dim
+            )
+            set_fuzzy_config_params(
+                decoder_params, module_output_candidate_params, output_feature_dim
+            )
+
+            # For decoders with conds: compute external_cond_dim from cond feature dimensions.
+            if cond_features:
+                external_cond_dim = sum(self.tokenizer_obs_dims[key][-1] for key in cond_features)
+                decoder_params["external_cond_dim"] = external_cond_dim
+                logger.info(
+                    f"Decoder '{decoder_name}' has conds={cond_features} "
+                    f"with external_cond_dim={external_cond_dim}"
+                )
+            self.decoder_cond_features[decoder_name] = cond_features
+            self.decoder_mask_features[decoder_name] = list(decoder_config.get("mask", []))
+
+            # Instantiate decoder
+            decoder = common.custom_instantiate(decoder_params, _resolve=False)
+            self.decoders[decoder_name] = decoder
+            self.decoder_input_features[decoder_name] = input_features
+            self.decoder_output_features[decoder_name] = output_features
+
+            logger.info(
+                f"Initialized {decoder_name} decoder with input features: {input_features} and output features: {output_features}"  # noqa: E501
+            )
+
+        # Filter active encoders/decoders if specified (for kinematic-only training)
+        if active_encoders is not None:
+            self.encoders_to_iterate = [e for e in self.encoders_to_iterate if e in active_encoders]
+            logger.info(f"Active encoders filtered to: {self.encoders_to_iterate}")
+        if active_decoders is not None:
+            # Keep all decoders in ModuleDict (for checkpoint compat), but only iterate active ones
+            self._active_decoders = set(active_decoders)
+            logger.info(f"Active decoders filtered to: {list(self._active_decoders)}")
+        else:
+            self._active_decoders = None
+
+        # Apply freeze logic
+        if self.freeze_encoders:
+            for encoder in self.encoders.values():
+                for param in encoder.parameters():
+                    param.requires_grad = False
+            logger.info(f"Froze encoders: {list(self.encoders.keys())}")
+
+        if self.freeze_decoders:
+            for decoder in self.decoders.values():
+                for param in decoder.parameters():
+                    param.requires_grad = False
+            logger.info(f"Froze decoders: {list(self.decoders.keys())}")
+
+        if self.freeze_quantizer and self.quantizer is not None:
+            for param in self.quantizer.parameters():
+                param.requires_grad = False
+            logger.info("Froze quantizer")
+
+        # Variable frame support: enabled when any encoder or decoder has mask observations
+        self.variable_frames_enabled = any(
+            feats
+            for feats in (
+                *self.encoder_mask_features.values(),
+                *self.decoder_mask_features.values(),
+            )
+        )
+        if self.variable_frames_enabled:
+            # Cache frame indices for mask creation (avoids allocation every forward)
+            self.register_buffer(
+                "_frame_indices",
+                torch.arange(self.num_future_frames).unsqueeze(0),
+                persistent=False,
+            )
+
+    def _create_frame_and_token_masks(self, tokenizer_obs):
+        """Create frame and token masks from num_frames observation.
+
+        Returns (frame_mask, token_mask) or (None, None) if disabled.
+            frame_mask: [B*seq, max_frames] bool, True=valid
+            token_mask: [B*seq, max_tokens] bool, True=valid
+        """
+        if "command_num_frames" not in tokenizer_obs:
+            return None, None
+
+        num_frames = tokenizer_obs["command_num_frames"]  # [B, seq, 1]
+        num_frames_flat = num_frames.reshape(-1, 1)  # [B*seq, 1]
+        frame_mask = self._frame_indices < num_frames_flat  # [B*seq, max_frames]
+
+        frames_per_token = 2**self.down_t
+        token_frame_mask = frame_mask.reshape(-1, self.max_num_tokens, frames_per_token)
+        token_mask = token_frame_mask.all(dim=-1)  # [B*seq, max_tokens]
+
+        return frame_mask, token_mask
+
+    def parse_tokenizer_obs(self, input_data):
+        """Split the flat tokenizer observation tensor into a named dict.
+
+        The tokenizer observation is stored as a single concatenated vector in
+        ``input_data["tokenizer"]``.  This method reshapes each slice back to
+        its original multi-dimensional form using the dimension metadata from
+        ``env_config``.
+
+        Args:
+            input_data: Dict containing ``"tokenizer"`` key with shape
+                ``(..., total_tokenizer_dim)``.
+
+        Returns:
+            Dict mapping observation name to tensor with shape
+            ``(*batch_dims, *obs_dims)`` for each registered tokenizer
+            observation.
+        """
+        tokenizer_obs = input_data["tokenizer"]
+        tokenizer_obs_dict = {}
+        for name, start, end, dims in self.tokenizer_obs_specs:
+            tokenizer_obs_dict[name] = tokenizer_obs[..., start:end].reshape(
+                tokenizer_obs.shape[:-1] + dims
+            )
+        assert self.tokenizer_obs_total_dim == tokenizer_obs.shape[-1], (
+            f"{self.tokenizer_obs_total_dim=}, {tokenizer_obs.shape[-1]=}"
+        )
+        return tokenizer_obs_dict
+
+    def create_encoder_masks(self, tokenizer_obs):
+        """Create encoder masks for each encoder.
+
+        When optimize_encoders_ratio_for_CHIP=False (legacy):
+            - encoder_index is multi-hot: SMPL-native envs have both SMPL=1 and G1=1
+            - encoder_masks["g1"] includes G1-native AND SMPL-native envs
+            - encoder_masks["g1_has_smpl"] identifies which G1-masked envs are SMPL-native
+
+        When optimize_encoders_ratio_for_CHIP=True (CHIP mode):
+            - encoder_index is one-hot: each env has exactly one native encoder
+            - encoder_masks are disjoint (no overlap)
+            - encoder_masks["g1_has_smpl"] is empty (G1-native envs don't have SMPL)
+            - G1 latents for SMPL-native envs are computed separately in aux losses
+        """
+        if len(self.encoders) == 1:
+            return {list(self.encoders.keys())[0]: None}
+
+        encoder_masks = {}
+        encoder_index = tokenizer_obs["encoder_index"]
+
+        for i, encoder_name in enumerate(self.encoder_sample_probs.keys()):
+            encoder_masks[encoder_name] = encoder_index[..., i].bool().flatten()
+
+        # Re-organized creation of intersection encoder masks.
+        # NOTE: No need to use compliance encoders in this version -- directly use the normal encoders
+
+        # For encoder pairs, create combined masks reflecting shared samples.
+        # The pairs and their intersection masks:
+        encoder_mask_combinations = [
+            # (key1, key2, intersection_mask_names)
+            ("g1", "smpl", [("g1_has_smpl", "smpl", "g1")]),
+            (
+                "teleop",
+                "smpl",
+                [("teleop_has_smpl", "smpl", "teleop"), ("smpl_has_teleop", "teleop", "smpl")],
+            ),
+            ("g1", "teleop", [("g1_has_teleop", "teleop", "g1")]),
+            ("teleop", "g1", [("teleop_has_g1", "g1", "teleop")]),
+            ("g1", "soma", [("g1_has_soma", "soma", "g1")]),
+        ]
+        for key1, key2, mask_defs in encoder_mask_combinations:
+            if key1 in encoder_masks and key2 in encoder_masks:
+                for mask_name, mask_src, mask_cond in mask_defs:
+                    encoder_masks[mask_name] = encoder_masks[mask_src][encoder_masks[mask_cond]]
+        return encoder_masks
+
+    def assemble_all_tokens(self, encoded_tokens, encoder_masks, batch_size, seq_len):
+        """Scatter per-encoder tokens into a single batch-aligned tensor.
+
+        Each encoder processes only the environments assigned to it (via
+        ``encoder_masks``).  This method writes each encoder's output back to
+        the correct rows of a zero-initialised buffer, then reshapes to
+        ``(batch, seq, num_tokens, token_dim)``.
+
+        Args:
+            encoded_tokens: Dict mapping encoder name to tensor of shape
+                ``(num_masked, num_tokens, token_dim)``.
+            encoder_masks: Dict mapping encoder name to a boolean index tensor
+                of length ``batch*seq``.  ``None`` means all samples.
+            batch_size: Batch dimension ``B``.
+            seq_len: Sequence dimension ``S``.
+
+        Returns:
+            Tensor of shape ``(B, S, num_tokens, token_dim)`` with each row
+            filled by its assigned encoder's output.
+        """
+        first_token = encoded_tokens[list(encoded_tokens.keys())[0]]
+        if len(encoded_tokens) == 1:
+            return first_token.view(batch_size, seq_len, *first_token.shape[1:])
+
+        all_tokens = torch.zeros(
+            (batch_size * seq_len,) + first_token.shape[1:],
+            dtype=first_token.dtype,
+            device=first_token.device,
+        )
+        for encoder_name in encoded_tokens.keys():
+            all_tokens[encoder_masks[encoder_name]] = encoded_tokens[encoder_name]
+        all_tokens = all_tokens.view(batch_size, seq_len, *all_tokens.shape[1:])
+        return all_tokens
+
+    def _encode_single(self, encoder_name, tokenizer_obs, encoder_mask=None, frame_mask=None):
+        """Encode using a single encoder without quantization or additive composition.
+        This is the core encoding logic used by both regular and additive encoders.
+
+        Args:
+            encoder_name: Name of the encoder to use
+            tokenizer_obs: Dictionary of tokenizer observations
+            encoder_mask: Optional mask for selecting specific samples
+            frame_mask: Optional [B*seq, max_frames] bool mask for variable frame support
+
+        Returns:
+            latent: The encoded latent representation
+        """  # noqa: D205
+        encoder = sub_encoders = None
+        if encoder_name in self.encoders:
+            encoder = self.encoders[encoder_name]
+        else:
+            sub_encoders = [
+                self.encoders[sub_encoder_name]
+                for sub_encoder_name in self.sub_encoders[encoder_name]
+            ]
+
+        input_features = self.encoder_input_features[encoder_name]
+
+        obs_list = [tokenizer_obs[key] for key in input_features]
+        # Only handle dimension broadcasting for the "smpl" encoder_name and for obs with last dim 9 or 12
+        if "smpl" in encoder_name:
+            max_ndim = max(obs.ndim for obs in obs_list)
+            # Check if any obs needs broadcasting: (last dim is 9 or 12 and ndim < max_ndim)
+            need_broadcast = any(
+                (obs.shape[-1] in (9, 12) and obs.ndim < max_ndim) for obs in obs_list
+            )
+            if need_broadcast:
+                aligned_obs_list = []
+                # Use the first obs of max_ndim as the temporal reference
+                target_obs = next(o for o in obs_list if o.ndim == max_ndim)
+                temporal_dim = target_obs.shape[-2]
+                for obs in obs_list:
+                    if obs.shape[-1] in (9, 12) and obs.ndim < max_ndim:
+                        # Add missing dimension(s)
+                        while obs.ndim < max_ndim:
+                            obs = obs.unsqueeze(-2)  # noqa: PLW2901
+                        # Expand temporal dimension
+                        expand_shape = list(obs.shape)
+                        expand_shape[-2] = temporal_dim
+                        obs = obs.expand(*expand_shape)  # noqa: PLW2901
+                    aligned_obs_list.append(obs)
+                encoder_input = torch.cat(aligned_obs_list, dim=-1)
+            else:
+                encoder_input = torch.cat(obs_list, dim=-1)
+        else:
+            encoder_input = torch.cat(obs_list, dim=-1)
+        encoder_input = encoder_input.view(-1, *encoder_input.shape[2:])
+
+        # Apply encoder_mask to both input and frame_mask
+        frame_mask_enc = None
+        if encoder_mask is not None:
+            encoder_input = encoder_input[encoder_mask]
+            if frame_mask is not None:
+                frame_mask_enc = frame_mask[encoder_mask]
+        elif frame_mask is not None:
+            frame_mask_enc = frame_mask
+
+        # encode using the provided encoder (with gradients enabled for end-to-end training)
+        if encoder is not None:
+            if frame_mask_enc is not None and self.encoder_mask_features.get(encoder_name):
+                latent = encoder(encoder_input, frame_mask=frame_mask_enc)
+            else:
+                latent = encoder(encoder_input)
+        else:
+            latent = encoder_input
+            for sub_encoder in sub_encoders:
+                latent = sub_encoder(latent)
+
+        return latent
+
+    def encode(
+        self, encoder_name, tokenizer_obs, encoder_mask=None, no_quantization=False, frame_mask=None
+    ):
+        """Encode using the specified encoder, including any additive encoders.
+
+        Args:
+            encoder_name: Name of the base encoder to use
+            tokenizer_obs: Dictionary of tokenizer observations
+            encoder_mask: Optional mask for selecting specific samples
+            no_quantization: If True, skip quantization and return raw latent
+            frame_mask: Optional [B*seq, max_frames] bool mask for variable frame support
+
+        Returns:
+            If no_quantization: latent tensor
+            Otherwise: (encoded_tokens, latent) tuple
+        """
+        # Get base encoder latent
+        latent = self._encode_single(
+            encoder_name, tokenizer_obs, encoder_mask, frame_mask=frame_mask
+        )
+
+        # Add contributions from additive encoders
+        if encoder_name in self.additive_encoders:
+            for additive_encoder_name in self.additive_encoders[encoder_name]:
+                additive_latent = self._encode_single(
+                    additive_encoder_name, tokenizer_obs, encoder_mask, frame_mask=frame_mask
+                )
+                latent = latent + additive_latent
+
+        if no_quantization:
+            return latent
+
+        # quantize using quantizer
+        if self.quantizer is not None:
+            quantized_codes, _ = self.quantizer(latent)
+            encoded_tokens = quantized_codes.contiguous()
+        else:
+            encoded_tokens = latent
+        return encoded_tokens, latent
+
+    def decode(self, decoder_name, decode_input_dict, token_mask=None):
+        """Run a single named decoder and split its output by feature.
+
+        Concatenates the decoder's declared input features from
+        ``decode_input_dict``, passes them through the decoder network
+        (optionally supplying ``external_cond`` and ``token_mask``), then
+        slices the output back into a per-feature dict.
+
+        Args:
+            decoder_name: Key into ``self.decoders`` and associated metadata.
+            decode_input_dict: Dict containing at least the keys declared in
+                the decoder's ``inputs`` and ``conds`` config entries.
+                Common keys include ``"token"``, ``"token_flattened"``, and
+                ``"proprioception"``.
+            token_mask: Optional ``(B*seq, max_tokens)`` bool mask for
+                variable-length token sequences.  Only forwarded to decoders
+                that declare mask features.
+
+        Returns:
+            Dict mapping output feature name to the corresponding slice of the
+            decoder output tensor, e.g.
+            ``{"action": Tensor(..., action_dim)}``.
+        """
+        decoder = self.decoders[decoder_name]
+        input_features = self.decoder_input_features[decoder_name]
+        output_feature_dims = self.decoder_output_feature_dims[decoder_name]
+        cond_features = self.decoder_cond_features.get(decoder_name, [])
+        decoder_input = torch.cat([decode_input_dict[key] for key in input_features], dim=-1)
+
+        # Build optional kwargs for the decoder call
+        kwargs = {}
+        if cond_features:
+            kwargs["external_cond"] = torch.cat(
+                [decode_input_dict[key] for key in cond_features], dim=-1
+            )
+        if token_mask is not None and self.decoder_mask_features.get(decoder_name):
+            kwargs["token_mask"] = token_mask
+        output = decoder(decoder_input, **kwargs)
+
+        # parse output
+        output_dict = {}
+        index = 0
+        for key, dim in output_feature_dims.items():
+            output_dict[key] = output[..., index : index + dim]
+            index += dim
+        assert index == output.shape[-1], f"{index=}, {output.shape[-1]=}"
+
+        return output_dict
+
+    def forward(  # noqa: D417
+        self,
+        input_data,
+        compute_aux_loss=False,
+        return_dict=False,
+        latent_residual=None,
+        latent_residual_mode="post_quantization",
+        **kwargs,  # noqa: ARG002
+    ):
+        """Run the full encode → quantize → decode pipeline.
+
+        Parses tokenizer observations, routes each environment to its assigned
+        encoder, optionally applies an external latent residual, then decodes
+        to joint-space actions.  When ``compute_aux_loss=True``, also computes
+        all registered auxiliary losses (e.g. G1-SMPL alignment, cycle
+        consistency) and returns a rich result dict.
+
+        Args:
+            input_data: Dict of named observation tensors, all with leading
+                shape ``(B, S, ...)``.  Must contain at minimum ``"actor_obs"``
+                (for batch/seq inference) and ``"tokenizer"`` (flat tokenizer
+                obs of shape ``(B, S, total_tokenizer_dim)``).
+            compute_aux_loss: When ``True``, evaluate all auxiliary loss
+                functions and include them in the returned dict.
+            return_dict: When ``True`` and ``compute_aux_loss=False``, return
+                the full output dict instead of just ``action_mean``.
+            latent_residual: Optional additive correction in latent token
+                space.  Shape ``(B, token_total_dim)`` where
+                ``token_total_dim = max_num_tokens * token_dim``.  Allows HOI
+                policies to steer motion without modifying the base ATM.
+            latent_residual_mode: Controls *when* ``latent_residual`` is
+                applied:
+
+                * ``"post_quantization"`` - add after FSQ (default, residual
+                  stays continuous).
+                * ``"pre_quantization"`` - add before FSQ (residual gets
+                  quantized together with the encoder latent).
+                * ``"pre_quantization_replace"`` - replace the encoder latent
+                  entirely with ``latent_residual`` before quantization.
+            **kwargs: Passed through; currently unused.
+
+        Returns:
+            When ``compute_aux_loss=True`` or ``return_dict=True``: a dict
+            with keys:
+
+            * ``"action_mean"`` - joint targets, shape ``(B, S, action_dim)``.
+            * ``"aux_losses"`` - dict of scalar loss tensors (empty when
+              ``compute_aux_loss=False``).
+            * ``"aux_loss_coef"`` - per-loss coefficient dict from config.
+            * ``"decoded_outputs"`` - raw decoder output dicts keyed by
+              decoder name.
+            * ``"tokenizer_obs"`` - parsed tokenizer observation dict.
+            * ``"encoder_masks"`` - per-encoder boolean index tensors.
+            * ``"encoded_tokens"`` - post-quantization tokens per encoder.
+            * ``"encoded_latents"`` - pre-quantization latents per encoder.
+            * ``"encoders_cfg"`` / ``"decoders_cfg"`` - config references.
+
+            When ``compute_aux_loss=False`` and ``return_dict=False``:
+            ``action_mean`` tensor directly.
+        """
+        # parse tokenizer obs
+        batch_size, seq_len = input_data["actor_obs"].shape[:2]
+        tokenizer_obs = self.parse_tokenizer_obs(input_data)
+        proprioception_input = torch.cat(
+            [input_data[key] for key in self.proprioception_features], dim=-1
+        )
+
+        # Reshape residual if provided: (batch, token_total_dim) -> (batch, 1, max_num_tokens, token_dim)
+        residual_reshaped = None
+        if latent_residual is not None:
+            residual_reshaped = latent_residual.view(
+                batch_size, 1, self.max_num_tokens, self.token_dim
+            )
+
+        # Variable frame masks
+        frame_mask, token_mask = self._create_frame_and_token_masks(tokenizer_obs)
+
+        # encode motion using all available encoders
+        encoder_masks = self.create_encoder_masks(tokenizer_obs)
+        encoded_tokens = {}
+        encoded_latents = {}
+        if latent_residual is not None and latent_residual_mode in [
+            "pre_quantization",
+            "pre_quantization_replace",
+        ]:
+            # PRE-QUANTIZATION MODES: add or replace residual before FSQ
+            # Reshape residual to (batch*seq, num_tokens, token_dim) for masking
+            residual_flat = latent_residual.view(
+                batch_size * seq_len, self.max_num_tokens, self.token_dim
+            )
+
+            for encoder_name in self.encoders_to_iterate:
+                encoder_mask = encoder_masks[encoder_name]
+
+                # Get raw latent (including additive encoders) without quantization
+                latent = self.encode(
+                    encoder_name,
+                    tokenizer_obs,
+                    encoder_mask,
+                    no_quantization=True,
+                    frame_mask=frame_mask,
+                )
+                # Apply same mask to residual before adding
+                if encoder_mask is not None:
+                    masked_residual = residual_flat[encoder_mask]
+                else:
+                    masked_residual = residual_flat
+
+                # Apply residual before quantization (only if we have samples)
+                if latent.shape[0] > 0:
+                    if latent_residual_mode == "pre_quantization":
+                        # Add residual to encoder latent
+                        latent = latent + masked_residual
+                    elif latent_residual_mode == "pre_quantization_replace":
+                        # Replace encoder latent with residual (zero out encoder)
+                        latent = masked_residual
+                    else:
+                        raise ValueError(f"Unknown latent_residual_mode: {latent_residual_mode}")
+
+                # Now quantize
+                if self.quantizer is not None:
+                    quantized_codes, _ = self.quantizer(latent)
+                    encoded_tokens[encoder_name] = quantized_codes.contiguous()
+                else:
+                    encoded_tokens[encoder_name] = latent
+                encoded_latents[encoder_name] = latent
+        else:
+            # STANDARD MODE: encode normally
+            for encoder_name in self.encoders_to_iterate:
+                encoded_tokens[encoder_name], encoded_latents[encoder_name] = self.encode(
+                    encoder_name,
+                    tokenizer_obs,
+                    encoder_masks[encoder_name],
+                    frame_mask=frame_mask,
+                )
+        all_tokens = self.assemble_all_tokens(encoded_tokens, encoder_masks, batch_size, seq_len)
+
+        # POST-QUANTIZATION MODE: add residual after FSQ tokens (default)
+        if latent_residual is not None and latent_residual_mode == "post_quantization":
+            all_tokens = all_tokens + residual_reshaped
+
+        # Cache tokens for external access (e.g., by callbacks)
+        self._last_encoded_tokens = {k: v.detach().cpu() for k, v in encoded_tokens.items()}
+        self._last_encoded_latents = {k: v.detach().cpu() for k, v in encoded_latents.items()}
+
+        # Cache flattened full latent on device for reward computation (token smoothness)
+        # all_tokens is the post-quantization token that gets sent to decoder
+        # Shape: (batch, seq, num_tokens, token_dim) -> (batch, seq, latent_dim)
+        self._last_full_latent_flat = all_tokens.detach().view(*all_tokens.shape[:-2], -1)
+
+        # decode action and motion
+        decode_input_dict = {
+            "token": all_tokens,
+            "token_flattened": all_tokens.view(*all_tokens.shape[:-2], -1),
+            "proprioception": proprioception_input,
+        }
+        decode_input_dict.update(tokenizer_obs)
+
+        decoded_outputs = {}
+        decoders_to_run = (
+            self._active_decoders if self._active_decoders is not None else self.decoders.keys()
+        )
+        for decoder_name in decoders_to_run:
+            decoded_outputs[decoder_name] = self.decode(
+                decoder_name, decode_input_dict, token_mask=token_mask
+            )
+
+        # Support "body_action", "meta_action", and "action" outputs (g1_dyn may not exist in kinematic-only mode)
+        if "g1_dyn" in decoded_outputs:
+            g1_dyn_out = decoded_outputs["g1_dyn"]
+            if "body_action" in g1_dyn_out:
+                action_mean = g1_dyn_out["body_action"]
+            elif "meta_action" in g1_dyn_out:
+                action_mean = g1_dyn_out["meta_action"]
+            else:
+                action_mean = g1_dyn_out["action"]
+
+            # Concatenate hand decoder output if present
+            if "hand_dyn" in decoded_outputs and "hand_action" in decoded_outputs["hand_dyn"]:
+                hand_action = decoded_outputs["hand_dyn"]["hand_action"]
+                action_mean = torch.cat([action_mean, hand_action], dim=-1)
+        else:
+            action_mean = None
+
+        # compute aux losses
+        if compute_aux_loss:
+            aux_losses = {}
+
+            # Initialize all paired latents
+            reencoded_smpl_g1_latents = None
+            paired_g1_smpl_latents = None
+            paired_compliance_latents = None
+            original_g1_latents_for_reencode = None
+
+            # Determine encoder names
+            smpl_encoder_name = "smpl" if "smpl" in self.encoders_to_iterate else None
+            teleop_encoder_name = "teleop" if "teleop" in self.encoders_to_iterate else None
+
+            # =========================================================================
+            # STIFF-MODE OPTIMIZATION: Check for stiff samples FIRST before computing
+            # expensive G1 latents. G1-SMPL loss and cycle consistency loss ONLY apply
+            # in stiff mode (compliance ≈ 0). Skip computation if all samples are compliant.
+            # =========================================================================
+            has_stiff_samples = False
+            compliance_values = None
+            smpl_latents = None
+            smpl_mask = None
+
+            if smpl_encoder_name is not None and "g1" in encoded_latents:
+                smpl_mask = encoder_masks.get(smpl_encoder_name)
+
+                if smpl_mask is not None and smpl_mask.sum() > 0:
+                    smpl_latents = encoded_latents[smpl_encoder_name]
+
+                    # Extract compliance values for these envs FIRST
+                    if "compliance" in tokenizer_obs:
+                        compliance_flat = tokenizer_obs["compliance"].view(
+                            -1, tokenizer_obs["compliance"].shape[-1]
+                        )
+                        compliance_values = compliance_flat[smpl_mask]
+
+                        # Check if ANY samples are stiff
+                        is_stiff = (compliance_values.abs() < self.stiff_compliance_threshold).all(
+                            dim=-1
+                        )
+                        has_stiff_samples = is_stiff.any().item()
+                    else:
+                        # No compliance info means all samples are "stiff" (default behavior)
+                        has_stiff_samples = True
+
+            # Only compute stiff-mode losses if there are stiff samples
+            if has_stiff_samples and smpl_mask is not None:
+                # Compute G1 latents for the same envs as SMPL (for proper pairing)
+                # OPTIMIZATION: Detach here since all downstream losses use g1 as
+                # fixed target (detach_g1_target=True). This avoids redundant gradient
+                # computation through the g1 encoder for these latents.
+                g1_latents_for_smpl = self.encode(
+                    encoder_name="g1",
+                    tokenizer_obs=tokenizer_obs,
+                    encoder_mask=smpl_mask,
+                    no_quantization=True,
+                ).detach()
+
+                paired_g1_smpl_latents = {
+                    "g1": g1_latents_for_smpl,
+                    "smpl": smpl_latents,
+                    "compliance": compliance_values,
+                }
+
+                # Store original G1 latents for cycle consistency loss (already detached)
+                original_g1_latents_for_reencode = g1_latents_for_smpl
+
+                # Cycle consistency: re-encode decoded G1 motion (only if stiff samples exist)
+                if self.reencode_smpl_g1_recon:
+                    reencoded_smpl_g1_latents = self.encode(
+                        encoder_name="g1",
+                        tokenizer_obs=decoded_outputs["g1_kin"],
+                        encoder_mask=smpl_mask,
+                        no_quantization=True,
+                    )
+
+            # Compute paired teleop-SMPL latents for TeleopSmplComplianceLatentLoss
+            # NOTE: This applies to ALL compliance values (not just stiff), so no stiff check
+            if (
+                teleop_encoder_name is not None
+                and smpl_latents is not None
+                and smpl_mask is not None
+            ):
+                # Compute teleop latents for the same envs as SMPL
+                # OPTIMIZATION: Detach here since TeleopSmplComplianceLatentLoss uses
+                # teleop as teacher (detach_teleop_target=True). Only SMPL learns.
+                teleop_latents_for_smpl = self.encode(
+                    encoder_name=teleop_encoder_name,
+                    tokenizer_obs=tokenizer_obs,
+                    encoder_mask=smpl_mask,
+                    no_quantization=True,
+                ).detach()
+
+                paired_compliance_latents = {
+                    "teleop": teleop_latents_for_smpl,
+                    "smpl": smpl_latents,
+                }
+
+            loss_inputs = {
+                "input_data": input_data,
+                "tokenizer_obs": tokenizer_obs,
+                "encoded_tokens": encoded_tokens,
+                "encoded_latents": encoded_latents,
+                "encoder_masks": encoder_masks,
+                "decoded_outputs": decoded_outputs,
+                "action_mean": action_mean,
+                "encoders_cfg": self.encoders_cfg,
+                "frame_mask": (
+                    frame_mask.view(batch_size, seq_len, -1) if frame_mask is not None else None
+                ),
+                "token_mask": (
+                    token_mask.view(batch_size, seq_len, -1) if token_mask is not None else None
+                ),
+                "decoders_cfg": self.decoders_cfg,
+                "reencoded_smpl_g1_latents": reencoded_smpl_g1_latents,
+                # New fields for compliance-aware losses
+                "paired_g1_smpl_latents": paired_g1_smpl_latents,
+                "paired_compliance_latents": paired_compliance_latents,
+                "original_g1_latents_for_reencode": original_g1_latents_for_reencode,
+            }
+
+            for loss_name, loss_func in self.aux_loss_func.items():
+                aux_losses[loss_name] = loss_func(loss_inputs)
+
+            output = {
+                "action_mean": action_mean,
+                "aux_losses": aux_losses,
+                "aux_loss_coef": self.aux_loss_coef,
+                # Exposed for recon trainer (existing RL callers ignore these keys)
+                "decoded_outputs": decoded_outputs,
+                "tokenizer_obs": tokenizer_obs,
+                "encoder_masks": encoder_masks,
+                "encoded_tokens": encoded_tokens,
+                "encoded_latents": encoded_latents,
+                "encoders_cfg": self.encoders_cfg,
+                "decoders_cfg": self.decoders_cfg,
+            }
+        elif return_dict:
+            output = {
+                "action_mean": action_mean,
+                "aux_losses": {},
+                "aux_loss_coef": self.aux_loss_coef,
+                "decoded_outputs": decoded_outputs,
+                "tokenizer_obs": tokenizer_obs,
+                "encoder_masks": encoder_masks,
+                "encoded_tokens": encoded_tokens,
+                "encoded_latents": encoded_latents,
+                "encoders_cfg": self.encoders_cfg,
+                "decoders_cfg": self.decoders_cfg,
+            }
+        else:
+            output = action_mean
+        return output
+
+    def get_token_info(self):
+        """Return a summary of the FSQ token configuration.
+
+        Returns:
+            Dict with keys:
+
+            * ``"token_dim"`` - dimensionality of one token.
+            * ``"total_dim"`` - ``token_dim * max_num_tokens``.
+            * ``"num_levels"`` - number of FSQ levels.
+            * ``"level_list"`` - per-level codebook sizes, or ``None`` if the
+              quantizer does not expose a ``levels`` attribute.
+            * ``"model_available"`` - always ``True``.
+        """
+        return {
+            "token_dim": self.token_dim,
+            "total_dim": self.token_total_dim,
+            "num_levels": self.num_fsq_levels,
+            "level_list": (
+                list(self.quantizer.levels) if hasattr(self.quantizer, "levels") else None
+            ),
+            "model_available": True,
+        }
+
+    def forward_with_external_tokens(  # noqa: D417
+        self, input_data: dict, external_tokens: torch.Tensor, **kwargs  # noqa: ARG002
+    ) -> torch.Tensor:
+        """Forward pass with externally provided tokens (bypasses encoder).
+
+        This is used when an external model (e.g., kinematic diffusion) provides
+        pre-computed FSQ tokens, allowing the encoder to be bypassed while still
+        using the decoder for action generation.
+
+        Args:
+            input_data: Dict with 'actor_obs' for proprioception
+            external_tokens: Pre-computed FSQ tokens from kinematic diffusion
+                            Shape: (B, 2, 32) or (B, seq, 2, 32)
+
+        Returns:
+            action_mean: (B, action_dim) or (B, seq, action_dim)
+        """
+        # Get proprioception input
+        proprioception_input = torch.cat(
+            [input_data[key] for key in self.proprioception_features], dim=-1
+        )
+
+        # Handle different input shapes
+        if external_tokens.dim() == 3:
+            # (B, 2, 32) -> add seq dim -> (B, 1, 2, 32)
+            external_tokens = external_tokens.unsqueeze(1)
+
+        # external_tokens: (B, seq, 2, 32) or (B, seq, num_tokens, token_dim)
+        batch_size = external_tokens.shape[0]
+        seq_len = external_tokens.shape[1]
+
+        # Flatten tokens for decoder: (B, seq, 2, 32) -> (B, seq, 64)
+        token_flattened = external_tokens.view(batch_size, seq_len, -1)
+
+        # Build decode input dict
+        decode_input_dict = {
+            "token": external_tokens,
+            "token_flattened": token_flattened,
+            "proprioception": proprioception_input,
+        }
+
+        # Decode actions using g1_dyn decoder
+        g1_dyn_out = self.decode("g1_dyn", decode_input_dict)
+        if "body_action" in g1_dyn_out:
+            action_mean = g1_dyn_out["body_action"]
+        elif "meta_action" in g1_dyn_out:
+            action_mean = g1_dyn_out["meta_action"]
+        else:
+            action_mean = g1_dyn_out["action"]
+
+        # Concatenate hand decoder output if present
+        if "hand_dyn" in self.decoders:
+            hand_dyn_out = self.decode("hand_dyn", decode_input_dict)
+            if "hand_action" in hand_dyn_out:
+                action_mean = torch.cat([action_mean, hand_dyn_out["hand_action"]], dim=-1)
+
+        # Squeeze seq dim if it was added
+        if action_mean.shape[1] == 1:
+            action_mean = action_mean.squeeze(1)
+
+        return action_mean
+
+    def get_example_input(self, encoder_name, batch_size=1, device="cpu"):
+        """Generate example input for ONNX export with specific encoder.
+
+        Args:
+            encoder_name: Name of the encoder to use
+            batch_size: Batch size for the example input
+            device: Device to create tensors on
+
+        Returns:
+            Tensor with shape (batch_size, feature_dim) including encoder features and proprioception
+        """
+        if encoder_name not in self.encoders:
+            raise ValueError(
+                f"Encoder '{encoder_name}' not found. Available encoders: {list(self.encoders.keys())}"
+            )
+
+        # Calculate total input feature dimension for the specified encoder
+        encoder_input_features = self.encoder_input_features[encoder_name]
+        total_feature_dim = 0
+
+        for feature_name in encoder_input_features:
+            if feature_name in self.tokenizer_obs_dims:
+                feature_dims = self.tokenizer_obs_dims[feature_name]
+                total_feature_dim += torch.prod(torch.tensor(feature_dims)).item()
+
+        # Add proprioception dimension
+        proprioception_dim = self.obs_dim_dict.get("actor_obs", 0)
+        total_feature_dim += proprioception_dim
+
+        # Create a single 2D tensor (B, F) including both encoder features and proprioception
+        example_input = torch.randn(batch_size, total_feature_dim, device=device)
+
+        return example_input
+
+    def get_all_example_input(self, batch_size=1, device="cpu"):
+        """Generate example inputs covering all tokenizer observations.
+
+        Constructs random tensors whose shapes match the full tokenizer
+        observation space plus proprioception.  Intended for ONNX tracing
+        or shape debugging.
+
+        Args:
+            batch_size: Number of examples in the batch dimension.
+            device: Target device for the returned tensors (e.g. ``"cpu"``
+                or ``"cuda"``).
+
+        Returns:
+            Tuple of:
+
+            * ``example_input`` - flat tensor of shape
+              ``(batch_size, total_feature_dim)`` where
+              ``total_feature_dim = tokenizer_feature_dim + proprioception_dim``.
+            * ``example_dict`` - dict with keys ``"tokenizer"`` and
+              ``"proprioception"``, each a random tensor with the
+              corresponding shape.
+        """
+        # Calculate total input feature dimension for the specified encoder
+        total_feature_dim = 0
+        for feature_name in self.tokenizer_obs_names:
+            feature_dims = self.tokenizer_obs_dims[feature_name]
+            print(  # noqa: T201
+                f"{feature_name}: {feature_dims}, start: {total_feature_dim} end: {total_feature_dim + torch.prod(torch.tensor(feature_dims)).item()}"  # noqa: E501
+            )
+            total_feature_dim += torch.prod(torch.tensor(feature_dims)).item()
+        tokenizer_feature_dim = total_feature_dim
+
+        # Add proprioception dimension
+        proprioception_dim = self.obs_dim_dict.get("actor_obs", 0)
+        total_feature_dim += proprioception_dim
+
+        print(f"tokenizer_feature_dim: {tokenizer_feature_dim}")  # noqa: T201
+        print(f"proprioception_dim: {proprioception_dim}")  # noqa: T201
+
+        # Create a single 2D tensor (B, F) including both encoder features and proprioception
+        example_input = torch.randn(batch_size, total_feature_dim, device=device)
+
+        example_dict = {
+            "tokenizer": torch.randn(batch_size, tokenizer_feature_dim, device=device),
+            "proprioception": torch.randn(batch_size, proprioception_dim, device=device),
+        }
+
+        return example_input, example_dict
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/trainer/__init__.py b/GR00T-WholeBodyControl/gear_sonic/trl/trainer/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/trainer/ppo_trainer.py b/GR00T-WholeBodyControl/gear_sonic/trl/trainer/ppo_trainer.py
new file mode 100644
index 0000000000000000000000000000000000000000..d0176ffe8cf9df0c82987b5ef983c7b3faf4e3e3
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/trainer/ppo_trainer.py
@@ -0,0 +1,2309 @@
+"""PPO trainer adapted from HuggingFace TRL for humanoid whole-body control."""
+
+import gc  # noqa: F401
+import math
+import os
+
+import accelerate
+from accelerate import utils as accelerate_utils
+import numpy as np
+import torch
+from torch import nn
+from transformers import utils as transformers_utils
+from transformers.trainer import *  # noqa: F403
+from trl import models
+from trl.experimental.ppo import ppo_trainer
+from trl.models import utils as models_utils
+from trl.trainer import utils as trainer_utils
+from trl.trainer.ppo_trainer import *  # noqa: F403
+import wandb
+
+# Conditional peft imports
+if accelerate_utils.is_peft_available():
+    import peft
+
+# Compatibility shim for loading old checkpoints saved with TRL < 0.28.0
+# These were moved in TRL 0.28.0, but old checkpoints reference the old paths
+import sys
+
+import trl.trainer.utils
+
+trl.trainer.utils.OnlineTrainerState = ppo_trainer.OnlineTrainerState
+trl.trainer.utils.exact_div = ppo_trainer.exact_div
+sys.modules["trl.trainer.utils"].OnlineTrainerState = ppo_trainer.OnlineTrainerState
+sys.modules["trl.trainer.utils"].exact_div = ppo_trainer.exact_div
+
+# Constants and utilities that may not be exported from trl
+INVALID_LOGPROB = 1.0  # Invalid log probability marker
+
+
+def masked_mean(values, mask, axis=None):
+    """Compute mean of values where mask is True."""
+    if axis is not None:
+        return (values * mask).sum(axis=axis) / mask.sum(axis=axis).clamp(min=1)
+    return (values * mask).sum() / mask.sum().clamp(min=1)
+
+
+from collections import deque  # noqa: E402
+
+import pandas as pd  # noqa: E402, F401
+from rich import (
+    console,
+    live,
+    panel,
+)
+from tqdm import tqdm  # noqa: E402, F401
+
+from gear_sonic.trl.callbacks import hv_callback_handler  # noqa: E402
+from gear_sonic.trl.modules import data_utils  # noqa: E402
+from gear_sonic.trl.utils import (
+    common,
+    rl,
+    scheduler,
+)
+from gear_sonic.utils import average_meters  # noqa: E402
+
+console_ = console.Console()
+import time  # noqa: E402
+
+
+class PolicyAndValueWrapper(nn.Module):
+    """Wrap policy, value, and optional discriminator models into a single nn.Module.
+
+    This wrapper enables a single ``forward`` call to dispatch to multiple model
+    components (policy, value, discriminator) in one DDP-safe pass, which is
+    required because calling ``forward`` on a DDP module more than once per
+    backward triggers gradient synchronization errors.
+    """
+
+    def __init__(self, policy, value_model, **kwargs) -> None:
+        super().__init__()
+        self.policy = policy
+        self.value_model = value_model
+        if "disc_model" in kwargs:
+            self.disc_model = kwargs["disc_model"]
+
+    @property
+    def is_gradient_checkpointing(self):
+        """
+        Whether gradient checkpointing is enabled for this model.
+        """  # noqa: D200, D212
+        if hasattr(self.policy, "is_gradient_checkpointing"):
+            return self.policy.is_gradient_checkpointing
+        return False
+
+    def gradient_checkpointing_enable(self, **kwargs):
+        if hasattr(self.policy, "gradient_checkpointing_enable"):
+            self.policy.gradient_checkpointing_enable(**kwargs)
+
+    def gradient_checkpointing_disable(self, **kwargs):
+        if hasattr(self.policy, "gradient_checkpointing_disable"):
+            self.policy.gradient_checkpointing_disable(**kwargs)
+
+    def set_mode(self, mode):
+        """Set the operating mode on the policy (e.g. ``"train"``, ``"eval"``, ``"train_rollout"``)."""
+        if hasattr(self.policy, "mode"):
+            self.policy.mode = mode
+
+    def transform_train(self):
+        """Enable training-time transforms (e.g. image augmentation) on the policy."""
+        if hasattr(self.policy, "transform_train"):
+            self.policy.transform_train()
+
+    def transform_eval(self):
+        """Disable training-time transforms on the policy for clean rollouts."""
+        if hasattr(self.policy, "transform_eval"):
+            self.policy.transform_eval()
+
+    def forward(self, modes, input_kwargs):
+        """Run forward passes for the requested component modes in a single call.
+
+        Args:
+            modes: List of mode strings to evaluate (e.g. ``["policy", "value"]``).
+            input_kwargs: Mapping from mode name to keyword arguments for that
+                component's forward pass.
+
+        Returns:
+            Dict mapping each mode name to its forward-pass output.
+        """
+        results = {}
+        for mode in modes:
+            results[mode] = self.forward_component(mode, **input_kwargs[mode])
+        return results
+
+    def forward_component(self, mode, actions=None, **kwargs):
+        """Dispatch a forward pass to a single component identified by *mode*.
+
+        Supported modes:
+            - ``"policy"`` -- standard policy forward with log-prob computation.
+            - ``"policy_distill"`` -- policy forward for distillation (no log-probs).
+            - ``"policy_distill_ppo"`` -- distillation forward that also computes
+              log-probs for PPO training.
+            - ``"policy_w_and_wo_imgaug"`` -- two policy forwards (with and without
+              image augmentation) for image-augmentation BC loss.
+            - ``"policy_deterministic"`` -- deterministic (mean) action only.
+            - ``"vae_policy_deterministic"`` -- VAE policy with prior evaluation.
+            - ``"value"`` -- value model evaluation.
+            - ``"eval_disc"`` / ``"train_disc"`` -- discriminator evaluation / training.
+
+        Args:
+            mode: Component mode string.
+            actions: Actions tensor for log-prob computation, ``(num_envs, num_steps, act_dim)``.
+            **kwargs: Forwarded to the underlying model component.
+
+        Returns:
+            Dict of component outputs (keys vary by mode).
+        """
+        if mode == "policy":
+            self.policy.act(**kwargs)
+            log_probs = self.policy.get_actions_log_prob(actions=actions)
+            results = {
+                "logprobs": log_probs,
+                "action_mean": self.policy.action_mean,
+                "action_std": self.policy.action_std,
+                "entropy": self.policy.entropy,
+            }
+            if self.policy.has_aux_loss:
+                results["aux_losses"] = self.policy.aux_losses
+                results["aux_loss_coef"] = self.policy.aux_loss_coef
+        elif mode == "policy_distill":
+            results = self.policy.act(**kwargs)
+        elif mode == "policy_distill_ppo":
+            policy_state_dict = self.policy.act(**kwargs)
+            log_probs = self.policy.get_actions_log_prob(actions=actions)
+            results = {
+                "actions": policy_state_dict["actions"],
+                "logprobs": log_probs,
+                "action_mean": policy_state_dict["action_mean"],
+                "action_std": policy_state_dict["action_sigma"],
+                "entropy": self.policy.entropy,
+            }
+            if "normalized_actions" in policy_state_dict:
+                results["normalized_actions"] = policy_state_dict["normalized_actions"]
+        elif mode == "policy_w_and_wo_imgaug":
+            # The first forward is without image augmentation
+            self.policy.transform_eval()
+            policy_state_dict = self.policy.act(**kwargs)
+            # Use the distribution without image augmentation to get the log_probs
+            log_probs = self.policy.get_actions_log_prob(actions=actions)
+            results = {
+                "actions": policy_state_dict["actions"],
+                "logprobs": log_probs,
+                "action_mean": policy_state_dict["action_mean"],
+                "action_std": policy_state_dict["action_sigma"],
+                "entropy": self.policy.entropy,
+            }
+
+            # The second forward is with image augmentation
+            self.policy.transform_train()
+            # The second time doesn't need deepcopy
+            policy_state_dict_w_imgaug = self.policy.act(**kwargs)
+            results["action_mean_w_imgaug"] = policy_state_dict_w_imgaug["action_mean"]
+            results["actions_w_imgaug"] = policy_state_dict_w_imgaug["actions"]
+            if "normalized_actions" in policy_state_dict_w_imgaug:
+                results["normalized_actions_w_imgaug"] = policy_state_dict_w_imgaug[
+                    "normalized_actions"
+                ]
+        elif mode == "policy_deterministic":
+            self.policy.act(**kwargs)
+            results = {
+                "action_mean": self.policy.action_mean,
+            }
+        elif mode == "vae_policy_deterministic":
+            self.policy.act(**kwargs)
+            prior_mu, prior_log_var = self.policy.eval_prior(**kwargs)
+            results = {
+                "action_mean": self.policy.action_mean,
+                "vae_mu": self.policy.z_mu,
+                "vae_log_var": self.policy.z_log_sigma,
+                "prior_mu": prior_mu,
+                "prior_log_var": prior_log_var,
+            }
+        elif mode == "value":
+            results = self.value_model.evaluate(**kwargs)
+        elif mode == "eval_disc":
+            results = self.disc_model.eval_disc(**kwargs)
+        elif mode == "train_disc":
+            results = self.disc_model.evaluate(**kwargs)
+        else:
+            raise ValueError(f"Invalid mode: {mode}")
+
+        return results
+
+
+class PrinterHVCallback(TrainerCallback):  # noqa: F405
+    """
+    A bare [`TrainerCallback`] that just prints the logs.
+    """  # noqa: D200, D212
+
+    def on_log(self, args, state, control, logs=None, **kwargs):  # noqa: ARG002
+        _ = logs.pop("total_flos", None)
+        if state.is_world_process_zero:
+            width = 80
+            pad = 35
+            print_str = f" \033[1m Learning iteration {state.global_step}  \033[0m "
+
+            log_string = (
+                f"""{print_str.center(width, ' ')}\n\n"""
+                f"""{'Computation:':>{pad}} {logs['fps']:.0f} steps/s (Collection: {logs['collection_time']:.3f}s, Learning {logs['learn_time']:.3f}s)\n"""  # noqa: E501
+                f"""{'Mean action noise std:':>{pad}} {logs['Policy/mean_noise_std']:.2f}\n"""
+            )
+
+            for k, v in logs.items():
+                if k.startswith("objective/"):
+                    # Keep the original logic
+                    if k.startswith("objective/kin_"):
+                        log_string += f"""{f'{k}:':>{pad}} {v:.5f}\n"""
+                    else:
+                        new_key = k.replace("objective/", "")
+                        log_string += f"""{f'Mean {new_key}:':>{pad}} {v:.5f}\n"""
+
+            env_log_string = ""
+            ep_string = ""
+            for k, v in logs.items():
+                if k.startswith("Env/"):
+                    entry = f"{f'{k}:':>{pad}} {v:.4f}"
+                    env_log_string += f"{entry}\n"
+                if k.startswith("Disc/"):
+                    entry = f"{f'{k}:':>{pad}} {v:.4f}"
+                    env_log_string += f"{entry}\n"
+                if k.startswith("Episode/"):
+                    new_key = k.replace("Episode/", "")
+                    ep_string += f"""{f'Mean episode {new_key}:':>{pad}} {v:.4f}\n"""
+
+            log_string += env_log_string
+            log_string += ep_string
+            log_string += (
+                f"""{'-' * width}\n"""
+                f"""{'Total episodes:':>{pad}} {logs['episode']}\n"""
+                f"""{'Total timesteps:':>{pad}} {logs['tot_timesteps']}\n"""
+                f"""{'Iteration time:':>{pad}} {logs['collection_time'] + logs['learn_time']:.2f}s\n"""
+                f"""{'Total time:':>{pad}} {logs['tot_time']:.2f}s\n"""
+                f"""{'ETA:':>{pad}} {logs['tot_time'] / logs['batch_idx'] * (logs['num_total_batches'] - logs['batch_idx']):.1f}s\n"""  # noqa: E501
+            )
+
+            log_string += f"Logging Directory: {logs['experiment_save_dir']}"
+            with live.Live(
+                panel.Panel(log_string, title="Training Log"),
+                refresh_per_second=4,
+                console=console_,
+            ):
+                # Your training loop or other operations
+                pass
+
+
+def process_ep_infos(ep_infos, device):
+    """Aggregate per-episode info dicts into a single dict of per-key means.
+
+    Args:
+        ep_infos: List of episode info dicts, each mapping metric names to
+            scalars or tensors.
+        device: Torch device to place intermediate tensors on.
+
+    Returns:
+        Dict mapping each metric name to its mean value across all episodes.
+    """
+    infos = {}
+    for key in ep_infos[0]:
+        # Buffer tensors in Python and concatenate once to avoid filling allocator buckets with mismatched sizes.
+        values = []
+        for ep_info in ep_infos:
+            v = ep_info[key]
+            if not isinstance(v, torch.Tensor):
+                v = torch.tensor([v], device=device)
+            else:
+                if v.dim() == 0: v = v.unsqueeze(0)
+                if v.device != device: v = v.to(device)
+            values.append(v)
+        infos[key] = torch.cat(values).mean()
+    return infos
+
+
+class TRLPPOTrainer(PPOTrainer):  # noqa: F405
+    """PPO trainer adapted from HuggingFace TRL for humanoid whole-body control.
+
+    Extends TRL's ``PPOTrainer`` to support IsaacLab-based vectorized
+    environments, multi-critic advantage estimation, symmetry augmentation,
+    adaptive KL-based learning rate scheduling, and gradient-checkpointed
+    policy/value models.
+
+    The training loop follows a standard on-policy PPO cycle:
+        1. Collect rollouts with the current policy (``_rollout_step``).
+        2. Compute GAE returns and advantages (``_compute_returns``).
+        3. Run multiple PPO epochs of mini-/micro-batch updates (``train``).
+        4. Synchronize running statistics across processes.
+        5. Log metrics and invoke callbacks.
+    """
+
+    _tag_names = ["trl", "humanoid_ppo"]  # noqa: RUF012
+
+    def __init__(
+        self,
+        args,
+        config,
+        env,
+        model,
+        ref_model=None,
+        reward_model=None,
+        processing_class=None,
+        value_model=None,
+        data_collator=None,
+        train_dataset=None,
+        eval_dataset=None,
+        log_dir=None,
+        # less commonly used
+        optimizers=(None, None),
+        callbacks=None,
+        peft_config=None,
+        use_ref_model=False,
+        checkpoint=None,
+        resume=False,
+        local_seed=None,
+        schedule_dict=None,
+        accelerator=None,
+        **kwargs,
+    ) -> None:
+        """Initialize the PPO trainer, models, optimizer, storage, and optionally load a checkpoint.
+
+        Args:
+            args: Training arguments (learning rate, batch sizes, clipping, etc.).
+            config: Algorithm config dict (PPO hyperparameters, num_steps_per_env, etc.).
+            env: Vectorized IsaacLab environment instance.
+            model: Policy (actor) model.
+            ref_model: Optional frozen reference policy for KL regularization.
+            reward_model: Optional learned reward model.
+            processing_class: HuggingFace processing class (unused, kept for API compat).
+            value_model: Critic (value) model.
+            data_collator: Optional data collator for the dataloader.
+            train_dataset: Optional training dataset (defaults to env-driven rollouts).
+            eval_dataset: Optional evaluation dataset.
+            log_dir: Directory for saving training logs.
+            optimizers: Tuple of ``(optimizer, lr_scheduler)``; created automatically if
+                ``(None, None)``.
+            callbacks: Additional ``TrainerCallback`` instances.
+            peft_config: Optional PEFT configuration for parameter-efficient fine-tuning.
+            use_ref_model: Whether to create/use a reference model for KL penalty.
+            checkpoint: Path to a checkpoint file to load on init.
+            resume: If True and *checkpoint* is provided, also restore optimizer,
+                scheduler, and trainer state for full training resumption.
+            local_seed: Per-process random seed for reproducibility under DDP.
+            schedule_dict: Dict defining parameter schedules over training steps.
+            accelerator: HuggingFace ``Accelerator`` instance for distributed training.
+            **kwargs: Extra keyword arguments forwarded to ``_init_trl`` (e.g.
+                ``disc_model`` for discriminator-based training).
+        """
+        self.accelerator = accelerator
+        self._init_trl(
+            args,
+            config,
+            env,
+            processing_class,
+            model,
+            ref_model,
+            reward_model,
+            train_dataset,
+            value_model,
+            data_collator,
+            eval_dataset,
+            optimizers,
+            callbacks,
+            peft_config,
+            use_ref_model,
+            local_seed,
+            log_dir,
+            schedule_dict=schedule_dict,
+            **kwargs,
+        )
+        self._init_config()
+        self._setup_storage()
+
+        if checkpoint is not None:
+            self.load_checkpoint(checkpoint, resume=resume)
+
+    def _init_trl(
+        self,
+        args,
+        config,
+        env,
+        processing_class,
+        model,
+        ref_model,
+        reward_model,
+        train_dataset,
+        value_model,
+        data_collator,
+        eval_dataset,
+        optimizers,
+        callbacks,
+        peft_config,
+        use_ref_model,
+        local_seed,
+        log_dir,
+        schedule_dict=None,
+        **kwargs,
+    ):
+        """Initialize TRL internals: models, optimizer, accelerator, dataloaders, and callbacks.
+
+        NOTE: This replicates much of TRL's ``PPOTrainer.__init__`` because the
+        upstream implementation assumes language-model rollouts, not vectorized
+        environment rollouts. Batch-size calculations, PEFT handling, DeepSpeed
+        preparation, and callback wiring are all customized here.
+        """
+        self.args = args
+        self.config = config
+        self.env = env
+        self.processing_class = processing_class
+        self.policy_model = model
+        self.learn_normalized_actions = model.has_normalized_actions
+        self.episode_env_tensors = average_meters.TensorAverageMeterDict()
+        self.ep_infos = []
+        self.eval_callbacks = []
+        self.log_dir = log_dir
+        self.schedule_dict = schedule_dict
+        self.scheduled_params_dict = {}
+
+        # peft support
+        if not accelerate_utils.is_peft_available() and peft_config is not None:
+            raise ImportError(
+                "PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models"  # noqa: E501
+            )
+        elif accelerate_utils.is_peft_available() and peft_config is not None:  # noqa: RET506
+            # if model is a peft model and we have a peft_confg, we merge and unload it first
+            if isinstance(self.policy_model, peft.PeftModel):
+                self.policy_model = self.policy_model.merge_and_unload()
+
+            # get peft model with the given config
+            self.policy_model = peft.get_peft_model(self.policy_model, peft_config)
+            if args.bf16 and getattr(self.policy_model, "is_loaded_in_4bit", False):
+                models_utils.peft_module_casting_to_bf16(self.policy_model)
+
+        self.is_peft_model = accelerate_utils.is_peft_available() and isinstance(
+            self.policy_model, peft.PeftModel
+        )
+        self.model_adapter_name = args.model_adapter_name
+        self.ref_adapter_name = args.ref_adapter_name
+
+        if use_ref_model:
+            if ref_model:
+                self.ref_model = ref_model
+            elif self.is_peft_model:
+                self.ref_model = None
+            else:
+                self.ref_model = models.create_reference_model(self.policy_model)
+        else:
+            self.ref_model = None
+
+        self.reward_model = reward_model
+        self.train_dataset = train_dataset
+        self.train_dataset_len = (
+            len(train_dataset) if train_dataset is not None else self.env.config.num_envs
+        )
+        self.value_model = value_model
+        self.data_collator = data_collator
+        self.eval_dataset = eval_dataset
+
+        self.optimizer, self.lr_scheduler = optimizers
+        self.optimizer_cls_and_kwargs = None  # needed for transformers >= 4.47
+
+        #########
+        # calculate various batch sizes
+        #########
+
+        accelerator = self.accelerator
+
+        self.device = accelerator.device
+        if "use_symmetry" in self.env.config and self.env.config.use_symmetry is not None:
+            self.use_symmetry = self.env.use_symmetry
+        else:
+            self.use_symmetry = False
+        args.global_rank = accelerator.process_index
+        args.world_size = accelerator.num_processes
+        args.is_main_process = accelerator.is_main_process
+        args.local_batch_size = self.env.config.num_envs
+        args.batch_size = int(args.local_batch_size * args.world_size)
+        try:
+            args.mini_batch_size = ppo_trainer.exact_div(
+                args.batch_size,
+                args.num_mini_batches,
+                "`batch_size` must be a multiple of `num_mini_batches`",
+            )
+            args.local_mini_batch_size = ppo_trainer.exact_div(
+                args.local_batch_size,
+                args.num_mini_batches,
+                "`local_batch_size` must be a multiple of `num_mini_batches`",
+            )
+        except Exception as e:  # noqa: BLE001
+            print(f"Error: {e}")  # noqa: T201
+            args.mini_batch_size = 1
+            args.local_mini_batch_size = 1
+
+        if args.per_device_train_batch_size is None:
+            args.per_device_train_batch_size = (
+                args.local_mini_batch_size
+            )  # same as mini-batch size, which implies no micro-batching (num_micro_batches = 1)
+        args.num_micro_batches = args.local_mini_batch_size // args.per_device_train_batch_size
+        args.micro_batch_size = int(args.per_device_train_batch_size * args.world_size)
+        # `per_rank_rollout_batch_size` is our `args.local_batch_size`
+        # `per_rank_minibatch_size` is our `args.local_mini_batch_size`
+        if args.total_episodes is None:
+            assert args.num_total_batches is not None
+            args.total_episodes = args.num_total_batches * args.batch_size
+        args.num_total_batches = math.ceil(
+            args.total_episodes / args.batch_size
+        )  # we may train for more than `total_episodes`
+        time_tensor = torch.tensor(int(time.time()), device=accelerator.device)
+        time_int = accelerate_utils.broadcast(
+            time_tensor, 0
+        ).item()  # avoid different timestamps across processes
+        args.run_name = f"{args.exp_name}__{args.seed}__{time_int}"
+        self.local_seed = local_seed
+        if args.num_sample_generations > 0:
+            self.sample_generations_freq = max(
+                1, args.num_total_batches // args.num_sample_generations
+            )
+        self.local_dataloader_batch_size = args.local_batch_size
+
+        #########
+        # setup model, optimizer, and others
+        #########
+        if self.config.get("disable_dropout", True):
+            for module in [self.policy_model, self.ref_model, self.value_model, self.reward_model]:
+                if module is not None:
+                    trainer_utils.disable_dropout_in_model(module)
+        addition_models = {}
+        if "reward_model" in kwargs:
+            addition_models["reward_model"] = kwargs["reward_model"]
+        if "disc_model" in kwargs:
+            addition_models["disc_model"] = kwargs["disc_model"]
+
+        self.model = PolicyAndValueWrapper(self.policy_model, self.value_model, **addition_models)
+        # self.model.config = self.policy_model.config  # needed for pushing to hub
+        self.create_optimizer_and_scheduler(
+            num_training_steps=args.num_total_batches
+        )  # note that we are calling `self.lr_scheduler.step()` manually only at the batch level
+
+        #########
+        ### trainer specifics
+        #########
+
+        default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(  # noqa: F405
+            self.args.report_to
+        )
+        self.callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks
+        self.callback_handler = hv_callback_handler.HVCallbackHandler(
+            self.callbacks,
+            self.model,
+            self.processing_class,
+            self.optimizer,
+            self.lr_scheduler,
+            self.env,
+            self.accelerator,
+        )
+        self.add_callback(
+            PrinterHVCallback if self.args.disable_tqdm else DEFAULT_PROGRESS_CALLBACK  # noqa: F405
+        )
+        self.control = TrainerControl()  # noqa: F405
+        self.state = ppo_trainer.OnlineTrainerState(
+            is_local_process_zero=self.is_local_process_zero(),
+            is_world_process_zero=self.is_world_process_zero(),
+            stateful_callbacks=[
+                cb
+                for cb in self.callback_handler.callbacks + [self.control]  # noqa: RUF005
+                if isinstance(cb, ExportableState)  # noqa: F405
+            ],
+        )
+        self.current_flos = 0
+        self.hp_search_backend = None
+        self.is_deepspeed_enabled = (
+            getattr(self.accelerator.state, "deepspeed_plugin", None) is not None
+        )
+        self.is_fsdp_enabled = getattr(self.accelerator.state, "fsdp_plugin", None) is not None
+        # Create distant repo and output directory if needed
+        self.hub_model_id = None
+        if self.args.push_to_hub:
+            self.init_hf_repo()
+        if self.args.should_save:
+            os.makedirs(self.args.output_dir, exist_ok=True)
+
+        # Add tags for models that have been loaded with the correct transformers version
+        if hasattr(self.model, "add_model_tags"):
+            self.model.add_model_tags(self._tag_names)
+
+        #########
+        ### setup dataloader
+        #########
+        if self.train_dataset is not None:
+            self.dataloader = DataLoader(  # noqa: F405
+                self.train_dataset,
+                batch_size=self.local_dataloader_batch_size,
+                shuffle=True,
+                collate_fn=self.data_collator,
+                drop_last=False,  # needed; otherwise the last batch will be of ragged shape
+            )
+        else:
+            self.dataloader = None
+        # sync random states for DataLoader(shuffle=True) before `accelerator.prepare`
+        # see https://gist.github.com/vwxyzjn/2581bff1e48e185e0b85b6dfe1def79c
+        torch.manual_seed(args.seed)
+
+        self.model, self.optimizer, self.dataloader = accelerator.prepare(
+            self.model, self.optimizer, self.dataloader
+        )
+        self.unwrapped_model = unwrap_model(self.model)  # noqa: F405
+        torch.manual_seed(self.local_seed)  # reset the local seed again
+
+        if self.eval_dataset is not None:
+            self.eval_dataloader = DataLoader(  # noqa: F405
+                self.eval_dataset,
+                batch_size=args.per_device_eval_batch_size,
+                collate_fn=self.data_collator,
+                drop_last=False,
+            )  # no need to shuffle eval dataset
+            self.eval_dataloader = accelerator.prepare(self.eval_dataloader)
+        else:
+            self.eval_dataloader = None
+
+        if self.is_deepspeed_enabled:
+            if self.reward_model is not None:
+                self.reward_model = trainer_utils.prepare_deepspeed(
+                    self.reward_model, args.per_device_train_batch_size, args.fp16, args.bf16
+                )
+
+            if self.ref_model is None:
+                if not self.is_peft_model:
+                    raise ValueError("No reference model and model is not a Peft model.")
+            else:
+                self.ref_model = trainer_utils.prepare_deepspeed(
+                    self.ref_model, args.per_device_train_batch_size, args.fp16, args.bf16
+                )
+        else:
+            if self.ref_model is None:
+                # if not self.is_peft_model:
+                #     raise ValueError("No reference model and model is not a Peft model.")
+                pass
+            else:
+                self.ref_model = self.ref_model.to(self.accelerator.device)
+            if self.reward_model is not None:
+                self.reward_model = self.reward_model.to(self.accelerator.device)
+        self.use_apex = False
+
+        self.train_with_evaluating_env = self.config.get("train_with_evaluating_env", False)
+
+        # Camera resolution
+        if "vision_obs" in self.env.config.obs.obs_dict:
+            if self.env.config.obs.obs_dict.vision_obs[0] in ["depth_image", "height_map"]:
+                num_channels = 1
+            elif self.env.config.obs.obs_dict.vision_obs[0] in ["rgb_image"]:
+                num_channels = 3
+            else:
+                raise ValueError(
+                    f"Invalid vision observation type: {self.env.config.obs.obs_dict.vision_obs[0]}"
+                )
+
+            if self.env.config.obs.obs_dict.vision_obs[0] == "height_map":
+                heightmap_resolution = self.env.config.simulator.config.heightmap.resolution
+                self.camera_resolution = [
+                    heightmap_resolution,
+                    heightmap_resolution,
+                ] + [  # noqa: RUF005
+                    num_channels
+                ]
+            else:
+                self.camera_resolution = [
+                    *self.env.config.simulator.config.cameras.camera_resolutions,
+                    num_channels,
+                ]
+        elif "camera_rgb" in self.env.config.robot.algo_obs_dim_dict:
+            # For manager_env with camera_rgb observation group
+            camera_res = self.env.config.get("camera_resolution", [144, 256])
+            self.camera_resolution = (camera_res[0], camera_res[1], 3)
+        else:
+            self.camera_resolution = None
+
+        self.num_critics = self.env.config.rewards.get("num_critics", 1)
+
+    def _init_config(self):
+        """Extract PPO hyperparameters and environment dimensions from config."""
+        # Env related Config
+        self.num_envs: int = self.env.config.num_envs
+        self.algo_obs_dim_dict = self.env.config.robot.algo_obs_dim_dict
+        self.num_act = self.policy_model.num_actions
+
+        self.num_steps_per_env = self.config.num_steps_per_env
+        self.use_padding_mask = self.config.get("use_padding_mask", False)
+        self.ppo_shuffle_every_epoch = self.config.get("ppo_shuffle_every_epoch", True)
+        self.empty_cache_every_n_ppo_epoch = self.config.get("empty_cache_every_n_ppo_epoch", -1)
+
+        self.entropy_coef = self.config.entropy_coef
+        self.desired_kl = self.config.desired_kl
+        self.gamma = self.args.gamma
+        self.lam = self.args.lam
+        self.adaptive_lr_min = self.config.get("adaptive_lr_min", 1e-5)
+        self.adaptive_lr_max = self.config.get("adaptive_lr_max", 1e-2)
+        self.sync_advantage_normalization = self.config.get("sync_advantage_normalization", True)
+        self.multi_critic_advantage_weights = self.config.get(
+            "multi_critic_advantage_weights", None
+        )
+
+        self.compute_imgaug_bc_loss = self.config.get("compute_imgaug_bc_loss", False)
+        self.imgaug_bc_loss_coef = self.config.get("imgaug_bc_loss_coef", 1.0)
+        self.imgaug_bc_loss_fn = torch.nn.MSELoss()
+
+    def _setup_storage(self):
+        """Allocate rollout storage buffers and episode tracking accumulators.
+
+        Registers observation, action, reward, done, value, return, and advantage
+        buffers in a ``RolloutStorage`` instance sized for
+        ``(num_envs, num_steps_per_env)``.
+        """
+        self.storage = data_utils.RolloutStorage(
+            self.env.num_envs, self.num_steps_per_env, device=self.accelerator.device
+        )
+        ## Register obs keys
+        for obs_key, obs_dim in self.algo_obs_dim_dict.items():
+            obs_shape = (obs_dim,) if isinstance(obs_dim, int) else tuple(obs_dim)
+            if obs_key in ["vision_obs", "camera_rgb"]:
+                # Vision observations are stored as [H, W, C] image, not flattened
+                self.storage.register_key(
+                    obs_key, shape=tuple(self.camera_resolution), dtype=torch.float
+                )
+            else:
+                self.storage.register_key(obs_key, shape=obs_shape, dtype=torch.float)
+            if obs_key == "critic_obs" and self.use_symmetry:
+                self.storage.register_key("next_" + obs_key, shape=obs_shape, dtype=torch.float)
+        ## Register others
+        reward_dim = self.num_critics
+        self.storage.register_key("actions", shape=(self.num_act,), dtype=torch.float)
+        self.storage.register_key("rewards", shape=(reward_dim,), dtype=torch.float)
+        self.storage.register_key("dones", shape=(1,), dtype=torch.bool)
+        self.storage.register_key("time_outs", shape=(1,), dtype=torch.bool)
+        self.storage.register_key("values", shape=(reward_dim,), dtype=torch.float)
+        self.storage.register_key("returns", shape=(reward_dim,), dtype=torch.float)
+        self.storage.register_key("advantages", shape=(reward_dim,), dtype=torch.float)
+        self.storage.register_key("actions_log_prob", shape=(1,), dtype=torch.float)
+        self.storage.register_key("action_mean", shape=(self.num_act,), dtype=torch.float)
+        self.storage.register_key("action_sigma", shape=(self.num_act,), dtype=torch.float)
+
+        if self.learn_normalized_actions:
+            self.storage.register_key(
+                "normalized_actions", shape=(self.num_act,), dtype=torch.float
+            )
+
+        self.state.rewbuffer = deque(maxlen=100)
+        self.state.lenbuffer = deque(maxlen=100)
+        self.cur_reward_sum = torch.zeros(
+            self.env.num_envs, self.num_critics, dtype=torch.float, device=self.accelerator.device
+        )
+        self.cur_episode_length = torch.zeros(
+            self.env.num_envs, dtype=torch.float, device=self.accelerator.device
+        )
+        self.state.cur_reward_sum = self.cur_reward_sum
+        self.state.cur_episode_length = self.cur_episode_length
+        self.ep_infos = []
+        self.state.tot_timesteps = 0
+        self.state.tot_time = 0
+        self.state.eval_step = 0
+        self.state.eval_render_step = 0
+
+    def policy_step(self, policy_model, obs_dict, cur_dones=None):
+        """Run the policy model for one rollout step, returning actions and log-probs.
+
+        Constructs an episode attention mask from the done history in storage
+        (for transformer-based policies), then calls ``policy_model.rollout``.
+
+        Args:
+            policy_model: The actor model to query.
+            obs_dict: Current observations, each value ``(num_envs, obs_dim)``.
+            cur_dones: Current done flags ``(num_envs,)``.  If None, the mask
+                is built from the full storage done history.
+
+        Returns:
+            Dict containing ``"actions"`` ``(num_envs, act_dim)``,
+            ``"actions_log_prob"`` ``(num_envs, 1)``, ``"action_mean"``,
+            ``"action_sigma"``, and any additional policy outputs.
+        """
+        actor_obs_dict = obs_dict.copy()
+
+        if cur_dones is None:
+            dones = (
+                self.storage.query_key("dones")
+                .to(self.accelerator.device)[: self.storage.step + 1]
+                .squeeze(-1)
+                .transpose(0, 1)
+            )
+            episode_attnmask = rl.compute_episode_attnmask(dones)
+        else:
+            episode_attnmask = None
+
+        policy_state_dict = policy_model.rollout(
+            obs_dict=actor_obs_dict, episode_attnmask=episode_attnmask, cur_dones=cur_dones
+        )
+        actions = policy_state_dict["actions"]
+        actions_log_prob = policy_model.get_actions_log_prob(actions).unsqueeze(1)
+        policy_state_dict["actions_log_prob"] = actions_log_prob
+
+        # assert len(actions.shape) == 2, f"{actions.shape=}"
+        # assert len(actions_log_prob.shape) == 2, f"{actions_log_prob.shape=}"
+        # assert len(action_mean.shape) == 2, f"{action_mean.shape=}"
+        # assert len(action_sigma.shape) == 2, f"{action_sigma.shape=}"
+
+        return policy_state_dict
+
+    def _chunked_value_evaluate(self, value_model, obs_dict, episode_attnmask, chunk_size=1024):
+        """Evaluate the value model in chunks to limit peak GPU memory.
+
+        Args:
+            value_model: Critic model with an ``evaluate`` method.
+            obs_dict: Observation dict, each value ``(batch, seq, obs_dim)``.
+            episode_attnmask: Attention mask ``(batch, seq, seq)`` or None.
+            chunk_size: Maximum batch dimension per chunk.
+
+        Returns:
+            Value predictions ``(batch, seq, num_critics)``.
+        """
+        batch_size = list(obs_dict.values())[0].shape[0]  # noqa: RUF015
+        if batch_size <= chunk_size:
+            return value_model.evaluate(obs_dict=obs_dict, episode_attnmask=episode_attnmask)
+
+        obs_chunks = {}
+        for key, value in obs_dict.items():
+            obs_chunks[key] = torch.split(value, chunk_size, dim=0)
+        if episode_attnmask is not None:
+            attnmask_chunks = torch.split(episode_attnmask, chunk_size, dim=0)
+        else:
+            attnmask_chunks = [None] * len(obs_chunks[list(obs_chunks.keys())[0]])  # noqa: RUF015
+
+        value_chunks = []
+        for i in range(len(attnmask_chunks)):
+            chunk_obs_dict = {key: obs_chunks[key][i] for key in obs_chunks}
+            chunk_values = value_model.evaluate(
+                obs_dict=chunk_obs_dict, episode_attnmask=attnmask_chunks[i]
+            )
+            value_chunks.append(chunk_values)
+        return torch.cat(value_chunks, dim=0)
+
+    def _rollout_step(self, model, obs_dict):
+        """Collect a full rollout of ``num_steps_per_env`` transitions and compute returns.
+
+        Performs the environment interaction loop under ``torch.no_grad()``,
+        stores transitions in ``self.storage``, then runs the value model over
+        the full trajectory to compute GAE returns and advantages.
+
+        Args:
+            model: The ``PolicyAndValueWrapper`` model (unwrapped from DDP).
+            obs_dict: Initial observations from the environment, each value
+                ``(num_envs, obs_dim)``.
+
+        Returns:
+            The final observation dict after the last environment step,
+            to be used as the starting point for the next rollout.
+        """
+        self._train_rollout_mode()
+        device = self.accelerator.device
+        policy_model = model.policy
+        value_model = model.value_model
+        policy_model.init_rollout()
+        self.storage.clear()
+
+        dones = torch.zeros(self.env.num_envs, device=device)
+        with torch.no_grad():
+            for i in range(self.num_steps_per_env):  # noqa: B007
+                # Compute the actions and values
+                # TODO: 1: unsqueeze to [B, 1, ...]  # noqa: TD002, TD003
+                policy_state_dict = self.policy_step(policy_model, obs_dict, cur_dones=dones)
+
+                # Append states to storage
+                for key, value in obs_dict.items():
+                    if key == "height_map":
+                        if getattr(self.storage, key).ndim != 5:
+                            # re-register height_map
+                            delattr(self.storage, key)
+                            self.storage.register_key(key, shape=value.shape[1:])
+                    elif key in ["vision_obs", "camera_rgb"]:
+                        # Vision observations have shape [B, H, W, C], verify storage matches
+                        if getattr(self.storage, key, None) is None:
+                            self.storage.register_key(key, shape=value.shape[1:], dtype=torch.float)
+                        elif getattr(self.storage, key).shape[2:] != value.shape[1:]:
+                            # re-register with correct shape
+                            delattr(self.storage, key)
+                            self.storage.register_key(key, shape=value.shape[1:], dtype=torch.float)
+                    self.storage.update_key(key, value)
+                for key, value in policy_state_dict.items():
+                    if key == "obs_dict":
+                        continue
+                    self.storage.update_key(key, value)
+
+                # Step the environment
+                if self.use_symmetry:
+                    obs_dict, rewards, dones, infos, termination_ids, termination_observations = (
+                        self.env.step(policy_state_dict)
+                    )
+                else:
+                    obs_dict, rewards, dones, infos = self.env.step(policy_state_dict)
+                for obs_key in obs_dict.keys():  # noqa: SIM118
+                    obs_dict[obs_key] = obs_dict[obs_key].to(device)
+                    if obs_key == "critic_obs" and self.use_symmetry:
+                        next_critic_obs = obs_dict[obs_key].clone()
+                        next_critic_obs[termination_ids.to(device)] = termination_observations.to(
+                            device
+                        )
+                        self.storage.update_key("next_" + obs_key, next_critic_obs)
+                rewards, dones = rewards.to(device), dones.to(device)
+                rewards_stored = rewards.clone()
+                if rewards.dim() == 1:
+                    rewards_stored = rewards_stored.unsqueeze(1)
+
+                assert len(rewards_stored.shape) == 2
+
+                self.ep_infos.append(infos["episode"])
+                self.storage.update_key("rewards", rewards_stored)
+                self.storage.update_key("dones", dones.unsqueeze(1))
+                self.storage.update_key("time_outs", infos["time_outs"].unsqueeze(1))
+                self.storage.increment_step()
+
+                self._process_env_step(rewards, dones, infos)
+                self.cur_reward_sum += rewards_stored
+                self.cur_episode_length += 1
+                new_ids = (dones > 0).nonzero(as_tuple=False)
+                self.state.rewbuffer.extend(self.cur_reward_sum[new_ids].cpu().numpy().tolist())
+                self.state.lenbuffer.extend(self.cur_episode_length[new_ids].cpu().numpy().tolist())
+                self.cur_reward_sum[new_ids] = 0
+                self.cur_episode_length[new_ids] = 0
+
+            policy_model.clear_rollout()
+            # gc.collect()
+            # torch.cuda.empty_cache()
+
+            if self.value_model is not None:
+                dones = self.storage.query_key("dones").to(device).squeeze(-1).transpose(0, 1)
+                dones = torch.cat([dones, torch.zeros_like(dones[:, :1])], dim=1)
+                episode_attnmask = rl.compute_episode_attnmask(dones)
+                all_obs_dict = {}
+                for key in obs_dict.keys():  # noqa: SIM118
+                    if key not in ["actor_obs"]:  # actor_obs not required by value model
+                        obs_value = self.storage.query_key(key).to(device)
+                        obs_value = torch.cat([obs_value, obs_dict[key].unsqueeze(0)], dim=0)
+                        all_obs_dict[key] = obs_value.transpose(0, 1)
+                all_values = self._chunked_value_evaluate(
+                    value_model, all_obs_dict, episode_attnmask
+                ).transpose(0, 1)
+                values, last_values = all_values[:-1], all_values[-1]
+
+                rewards = self.storage.query_key("rewards")
+
+                new_rewards = (
+                    rewards.to(device)
+                    + self.gamma * self.storage.query_key("time_outs").to(device) * values
+                )
+                self.storage.batch_update_data("rewards", new_rewards)
+
+                returns, advantages = self._compute_returns(
+                    values=values,
+                    last_values=last_values,
+                    policy_state_dict={
+                        "dones": self.storage.query_key("dones"),
+                        "rewards": self.storage.query_key("rewards"),
+                    },
+                )
+                self.storage.batch_update_data("values", values)
+                self.storage.batch_update_data("returns", returns)
+                self.storage.batch_update_data("advantages", advantages)
+
+        return obs_dict
+
+    def _flip_obs(self, obs, key):
+        """Mirror observations left-right for symmetry augmentation.
+
+        Args:
+            obs: Observation tensor ``(batch, seq, obs_dim)``.
+            key: Observation key (``"actor_obs"`` or ``"critic_obs"``), which
+                determines the history length and flip index mapping.
+
+        Returns:
+            Flipped observation tensor with the same shape as *obs*.
+        """
+        if key == "actor_obs":
+            proprioceptive_obs = obs.clone().view(
+                obs.shape[0], obs.shape[1], self.env.actor_history_length + 1, -1
+            )
+            flipper_proprioceptive_obs = torch.zeros_like(proprioceptive_obs)
+            flipper_proprioceptive_obs[:, :, :, :] = (
+                proprioceptive_obs[:, :, :, self.env.flip_actor_obs_info[:, 0]]
+                * self.env.flip_actor_obs_info[:, 1]
+            )
+        elif key == "critic_obs":
+            proprioceptive_obs = obs.clone().view(
+                obs.shape[0], obs.shape[1], self.env.critic_history_length + 1, -1
+            )
+            flipper_proprioceptive_obs = torch.zeros_like(proprioceptive_obs)
+            flipper_proprioceptive_obs[:, :, :, :] = (
+                proprioceptive_obs[:, :, :, self.env.flip_critic_obs_info[:, 0]]
+                * self.env.flip_critic_obs_info[:, 1]
+            )
+        else:
+            raise NotImplementedError
+        return flipper_proprioceptive_obs.view(obs.shape[0], obs.shape[1], -1)
+
+    def _flip_actions(self, actions):
+        """Mirror actions left-right for symmetry augmentation.
+
+        Args:
+            actions: Action tensor ``(batch, seq, act_dim)``.
+
+        Returns:
+            Flipped action tensor with the same shape.
+        """
+        flipped_actions = (
+            actions[:, :, self.env.flip_action_info[:, 0]] * self.env.flip_action_info[:, 1]
+        )
+        return flipped_actions
+
+    def _process_env_step(self, rewards, dones, infos):  # noqa: ARG002
+        """Handle post-step bookkeeping: reset models on done envs and log metrics.
+
+        Args:
+            rewards: Reward tensor ``(num_envs,)`` or ``(num_envs, num_critics)``.
+            dones: Done flags ``(num_envs,)``.
+            infos: Info dict from the environment step, containing ``"to_log"``
+                entries for metric tracking.
+        """
+        self.policy_model.reset(dones)
+        if self.value_model is not None:
+            self.value_model.reset(dones)
+        self.episode_env_tensors.add(infos["to_log"])
+
+    def _register_stats_buffer(self):
+        """Allocate per-epoch/mini-batch/micro-batch statistic tensors for logging.
+
+        Tensors have shape ``(num_ppo_epochs, num_mini_batches, num_micro_batches)``
+        and are overwritten each training iteration.
+        """
+        args = self.args
+        device = self.accelerator.device
+
+        stats_shape = (args.num_ppo_epochs, args.num_mini_batches, args.num_micro_batches)
+        approxkl_stats = torch.zeros(stats_shape, device=device)
+        pg_clipfrac_stats = torch.zeros(stats_shape, device=device)
+        pg_loss_stats = torch.zeros(stats_shape, device=device)
+        vf_loss_stats = torch.zeros(stats_shape, device=device)
+        entropy_stats = torch.zeros(stats_shape, device=device)
+        weighted_ppo_loss_stats = torch.zeros(stats_shape, device=device)
+        vf_clipfrac_stats = torch.zeros(stats_shape, device=device)
+        ratio_stats = torch.zeros(stats_shape, device=device)
+        advantage_mean_stats = torch.zeros(stats_shape, device=device)
+        advantage_std_stats = torch.zeros(stats_shape, device=device)
+        if self.compute_imgaug_bc_loss:
+            imgaug_bc_loss_stats = torch.zeros(stats_shape, device=device)
+            weighted_imgaug_bc_loss_stats = torch.zeros(stats_shape, device=device)
+
+        self.approxkl_stats = approxkl_stats
+        self.pg_clipfrac_stats = pg_clipfrac_stats
+        self.pg_loss_stats = pg_loss_stats
+        self.vf_loss_stats = vf_loss_stats
+        self.entropy_stats = entropy_stats
+        self.weighted_ppo_loss_stats = weighted_ppo_loss_stats
+        self.vf_clipfrac_stats = vf_clipfrac_stats
+        self.ratio_stats = ratio_stats
+        self.advantage_mean_stats = advantage_mean_stats
+        self.advantage_std_stats = advantage_std_stats
+        if self.use_symmetry:
+            estimation_loss_stats = torch.zeros(stats_shape, device=device)
+            swap_loss_stats = torch.zeros(stats_shape, device=device)
+            actor_sym_loss_stats = torch.zeros(stats_shape, device=device)
+            critic_sym_loss_stats = torch.zeros(stats_shape, device=device)
+            self.estimation_loss_stats = estimation_loss_stats
+            self.swap_loss_stats = swap_loss_stats
+            self.actor_sym_loss_stats = actor_sym_loss_stats
+            self.critic_sym_loss_stats = critic_sym_loss_stats
+
+        if self.compute_imgaug_bc_loss:
+            self.imgaug_bc_loss_stats = imgaug_bc_loss_stats
+            self.weighted_imgaug_bc_loss_stats = weighted_imgaug_bc_loss_stats
+
+    def _get_rollout_data(self, obs_keys):
+        """Transpose storage from ``(steps, envs, ...)`` to ``(envs, steps, ...)`` and apply augmentations.
+
+        Retrieves all rollout tensors from ``self.storage``, optionally doubles
+        the batch via left-right symmetry flipping, and builds padding masks
+        for episodes that terminated mid-rollout.
+
+        Args:
+            obs_keys: Observation keys to extract from storage.
+
+        Returns:
+            Dict containing ``"all_obs_dict"``, ``"actions"``, ``"logprobs"``,
+            ``"values"``, ``"rewards"``, ``"dones"``, ``"old_mu_batch"``,
+            ``"old_sigma_batch"``, ``"returns"``, ``"advantages"``, and
+            padding masks.
+        """
+        device = self.accelerator.device
+
+        all_obs_dict = {
+            key: self.storage.query_key(key).transpose(0, 1).to(device) for key in obs_keys
+        }
+        actions = self.storage.actions.transpose(0, 1).to(device)
+        logprobs = self.storage.actions_log_prob.transpose(0, 1).squeeze(-1).to(device)
+        values = self.storage.values.transpose(0, 1).to(device)  # noqa: PD011
+        rewards = self.storage.rewards.transpose(0, 1).to(device)
+        dones = self.storage.dones.transpose(0, 1).squeeze(-1).to(device)
+        old_mu_batch = self.storage.action_mean.transpose(0, 1).to(device)
+        old_sigma_batch = self.storage.action_sigma.transpose(0, 1).to(device)
+        returns = self.storage.returns.transpose(0, 1).to(device)
+        advantages = self.storage.advantages.transpose(0, 1).to(device)
+
+        if self.use_symmetry:
+            next_critic_obs = self.storage.next_critic_obs.transpose(0, 1).to(device)
+            all_obs_dict = {
+                key: torch.cat((all_obs_dict[key], self._flip_obs(all_obs_dict[key], key)), dim=0)
+                for key in all_obs_dict.keys()  # noqa: SIM118
+            }
+            next_critic_obs = torch.cat(
+                (next_critic_obs, self._flip_obs(next_critic_obs, "critic_obs")), dim=0
+            )
+            actions = torch.cat((actions, self._flip_actions(actions)), dim=0)
+            logprobs = logprobs.repeat(2, 1)
+            values = values.repeat(2, 1, 1)
+            rewards = rewards.repeat(2, 1, 1)
+            dones = dones.repeat(2, 1)
+            old_mu_batch = old_mu_batch.repeat(2, 1, 1)
+            old_sigma_batch = old_sigma_batch.repeat(2, 1, 1)
+            returns = returns.repeat(2, 1, 1)
+            advantages = advantages.repeat(2, 1, 1)
+
+        if self.use_padding_mask:
+            padding_mask = dones.clone()
+            padding_mask_p1 = padding_mask.clone()
+            for i in range(padding_mask.shape[0]):
+                true_indices = torch.where(padding_mask[i])[0]
+                if len(true_indices) > 0:
+                    padding_mask[i, true_indices[0]] = False
+                    padding_mask_p1[
+                        i, true_indices[0] : min(true_indices[0] + 2, padding_mask_p1.shape[1])
+                    ] = False
+            logprobs = torch.masked_fill(logprobs, padding_mask, INVALID_LOGPROB)
+            values = torch.masked_fill(values, padding_mask_p1, 0)
+        else:
+            padding_mask = torch.zeros_like(dones)
+            padding_mask_p1 = torch.zeros_like(dones)
+
+        rollout_data = {
+            "all_obs_dict": all_obs_dict,
+            "actions": actions,
+            "logprobs": logprobs,
+            "values": values,
+            "rewards": rewards,
+            "dones": dones,
+            "old_mu_batch": old_mu_batch,
+            "old_sigma_batch": old_sigma_batch,
+            "returns": returns,
+            "advantages": advantages,
+            "padding_mask": padding_mask,
+            "padding_mask_p1": padding_mask_p1,
+        }
+        if self.use_symmetry:
+            rollout_data["next_critic_obs"] = next_critic_obs
+        return rollout_data
+
+    def _get_mb_rollout_data(self, rollout_data, micro_batch_inds):
+        """Slice a micro-batch from the full rollout data and build its attention mask.
+
+        Args:
+            rollout_data: Full rollout dict from ``_get_rollout_data``.
+            micro_batch_inds: 1-D index tensor selecting environments for this
+                micro-batch.
+
+        Returns:
+            Dict of micro-batch tensors (prefixed ``"mb_"``) plus the
+            ``"episode_attnmask"`` for this subset.
+        """
+        mb_obs_dict = {
+            key: rollout_data["all_obs_dict"][key][micro_batch_inds]
+            for key in rollout_data["all_obs_dict"].keys()  # noqa: SIM118
+        }
+        mb_advantage = rollout_data["advantages"][micro_batch_inds]
+        mb_logprobs = rollout_data["logprobs"][micro_batch_inds]
+        mb_return = rollout_data["returns"][micro_batch_inds]
+        mb_values = rollout_data["values"][micro_batch_inds]
+        mb_dones = rollout_data["dones"][micro_batch_inds]
+        mb_actions = rollout_data["actions"][micro_batch_inds]
+        mb_old_mu = rollout_data["old_mu_batch"][micro_batch_inds]
+        mb_old_sigma = rollout_data["old_sigma_batch"][micro_batch_inds]
+        mb_padding_mask = rollout_data["padding_mask"][micro_batch_inds]
+        mb_padding_mask_p1 = rollout_data["padding_mask_p1"][micro_batch_inds]
+
+        episode_attnmask = rl.compute_episode_attnmask(mb_dones)
+
+        mb_rollout_data = {
+            "micro_batch_inds": micro_batch_inds,
+            "mb_obs_dict": mb_obs_dict,
+            "mb_advantage": mb_advantage,
+            "mb_logprobs": mb_logprobs,
+            "mb_return": mb_return,
+            "mb_values": mb_values,
+            "mb_dones": mb_dones,
+            "mb_actions": mb_actions,
+            "mb_old_mu": mb_old_mu,
+            "mb_old_sigma": mb_old_sigma,
+            "mb_padding_mask": mb_padding_mask,
+            "mb_padding_mask_p1": mb_padding_mask_p1,
+            "episode_attnmask": episode_attnmask,
+        }
+        if self.use_symmetry:
+            mb_next_critic_obs = rollout_data["next_critic_obs"][micro_batch_inds]
+            mb_rollout_data["mb_next_critic_obs"] = mb_next_critic_obs
+        return mb_rollout_data
+
+    def _forward_model(self, model, mb_rollout_data):
+        """Run a single combined policy + value forward pass on a micro-batch.
+
+        NOTE: Policy and value are forwarded together in one ``model.forward``
+        call so DDP only synchronizes gradients once per backward pass.
+
+        Args:
+            model: The ``PolicyAndValueWrapper`` (possibly DDP-wrapped).
+            mb_rollout_data: Micro-batch dict from ``_get_mb_rollout_data``.
+
+        Returns:
+            Dict with ``"policy_results"`` and ``"value_results"`` sub-dicts.
+        """
+        mb_obs_dict = mb_rollout_data["mb_obs_dict"]
+        mb_actions = mb_rollout_data["mb_actions"]
+        episode_attnmask = mb_rollout_data["episode_attnmask"]
+
+        # We should only do one forward pass for especially DDP model
+        if self.compute_imgaug_bc_loss:
+            results = model.forward(
+                modes=["policy_w_and_wo_imgaug", "value"],
+                input_kwargs={
+                    "policy_w_and_wo_imgaug": {
+                        "obs_dict": mb_obs_dict,
+                        "actions": mb_actions,
+                        "episode_attnmask": episode_attnmask,
+                    },
+                    "value": {"obs_dict": mb_obs_dict, "episode_attnmask": episode_attnmask},
+                },
+            )
+            policy_results = results["policy_w_and_wo_imgaug"]
+        else:
+            with common.Timer("wrapper_forward_model"):
+                results = model.forward(
+                    modes=["policy", "value"],
+                    input_kwargs={
+                        "policy": {
+                            "obs_dict": mb_obs_dict,
+                            "actions": mb_actions,
+                            "episode_attnmask": episode_attnmask,
+                        },
+                        "value": {"obs_dict": mb_obs_dict, "episode_attnmask": episode_attnmask},
+                    },
+                )
+                policy_results = results["policy"]
+        return {
+            "policy_results": policy_results,
+            "value_results": results["value"],
+        }
+
+    def _compute_loss(self, forward_results, mb_rollout_data):
+        """Compute the total loss as a weighted sum of PPO and optional auxiliary losses.
+
+        Args:
+            forward_results: Output of ``_forward_model``.
+            mb_rollout_data: Micro-batch dict from ``_get_mb_rollout_data``.
+
+        Returns:
+            Dict containing ``"loss"`` (scalar to backprop), ``"ppo_loss_dict"``,
+            and optionally ``"imgaug_bc_loss_dict"``.
+        """
+        ppo_loss_dict = self._compute_ppo_loss(forward_results, mb_rollout_data)
+
+        loss = ppo_loss_dict["ppo_loss"] * self.config.get("ppo_loss_coef", 1.0)
+
+        ret_dict = {
+            "ppo_loss_dict": ppo_loss_dict,
+        }
+
+        if self.compute_imgaug_bc_loss:
+            imgaug_bc_loss_dict = self._compute_imgaug_bc_loss(forward_results, mb_rollout_data)
+            loss += imgaug_bc_loss_dict["imgaug_bc_loss"] * self.config.imgaug_bc_loss_coef
+            ret_dict["imgaug_bc_loss_dict"] = imgaug_bc_loss_dict
+
+        ret_dict["loss"] = loss
+
+        return ret_dict
+
+    def _compute_ppo_loss(self, forward_results, mb_rollout_data):
+        """Compute the clipped PPO surrogate loss, value loss, and entropy bonus.
+
+        Implements standard clipped PPO with:
+        - Clipped surrogate policy gradient loss.
+        - Clipped value function loss.
+        - Entropy regularization.
+        - Adaptive learning rate adjustment based on KL divergence.
+        - Optional left-right symmetry consistency losses for actor and critic.
+
+        Args:
+            forward_results: Output of ``_forward_model``.
+            mb_rollout_data: Micro-batch dict from ``_get_mb_rollout_data``.
+
+        Returns:
+            Dict with ``"ppo_loss"`` (combined scalar), plus individual loss
+            components and diagnostic metrics (KL, clip fractions, ratios).
+        """
+        args = self.args
+        optimizer = self.optimizer
+
+        policy_results = forward_results["policy_results"]
+        value_results = forward_results["value_results"]
+
+        mb_obs_dict = mb_rollout_data["mb_obs_dict"]
+        mb_old_mu = mb_rollout_data["mb_old_mu"]
+        mb_old_sigma = mb_rollout_data["mb_old_sigma"]
+        mb_values = mb_rollout_data["mb_values"]
+        mb_return = mb_rollout_data["mb_return"]
+        mb_logprobs = mb_rollout_data["mb_logprobs"]
+        mb_advantage = mb_rollout_data["mb_advantage"]
+        padding_mask = mb_rollout_data["mb_padding_mask"]
+        padding_mask_p1 = mb_rollout_data["mb_padding_mask_p1"]
+        micro_batch_inds = mb_rollout_data["micro_batch_inds"]  # noqa: F841
+
+        new_logprobs = policy_results["logprobs"]
+        sigma_batch = policy_results["action_std"]
+        mu_batch = policy_results["action_mean"]
+        entropy_batch = policy_results["entropy"]
+        with torch.no_grad():
+            kl = torch.sum(
+                torch.log(sigma_batch / mb_old_sigma + 1.0e-5)
+                + (torch.square(mb_old_sigma) + torch.square(mb_old_mu - mu_batch))
+                / (2.0 * torch.square(sigma_batch))
+                - 0.5,
+                axis=-1,
+            )
+            local_kl_mean = torch.mean(kl)
+            kl_mean = self.accelerator.gather(local_kl_mean).mean()
+            self._adjust_learning_rate_based_on_kl(kl_mean, optimizer)
+
+        # Forward a DDP model twice will cause the error: "one of the variables needed for gradient computation has been modified by an inplace operation"  # noqa: E501
+        vpred = value_results
+        vpredclipped = torch.clamp(
+            vpred,
+            mb_values - args.cliprange_value,
+            mb_values + args.cliprange_value,
+        )
+        vf_losses1 = torch.square(vpred - mb_return)
+        vf_losses2 = torch.square(vpredclipped - mb_return)
+        vf_loss_max = torch.max(vf_losses1, vf_losses2).mean(dim=-1)
+        # vf_loss_max[vf_loss_max.isnan()] = 0.0
+        vf_loss = masked_mean(vf_loss_max, ~padding_mask_p1)
+        vf_clipfrac = masked_mean((vf_losses2 > vf_losses1).float(), ~padding_mask_p1.unsqueeze(-1))
+        logprobs_diff = new_logprobs - mb_logprobs
+        ratio = torch.exp(logprobs_diff).unsqueeze(-1)
+        if self.multi_critic_advantage_weights is not None:
+            mb_advantage = (
+                mb_advantage
+                * torch.tensor(self.multi_critic_advantage_weights).to(mb_advantage)[None, None, :]
+            )
+        pg_losses = -mb_advantage * ratio
+        pg_losses2 = -mb_advantage * torch.clamp(ratio, 1.0 - args.cliprange, 1.0 + args.cliprange)
+        pg_loss_max = torch.max(pg_losses, pg_losses2).sum(dim=-1)
+        # pg_loss_max[pg_loss_max.isnan()] = 0.0
+        pg_loss = masked_mean(pg_loss_max, ~padding_mask)
+
+        # entropy_batch[entropy_batch.isnan()] = 0.0
+        entropy_loss = -masked_mean(entropy_batch, ~padding_mask)
+        if self.use_symmetry:
+            actor_sym_loss = torch.mean(
+                torch.sum(
+                    torch.square(
+                        self._flip_actions(self.policy_model(mb_obs_dict))
+                        - self.policy_model(
+                            {"actor_obs": self._flip_obs(mb_obs_dict["actor_obs"], "actor_obs")}
+                        )
+                    ),
+                    dim=-1,
+                )
+            )
+            critic_sym_loss = torch.mean(
+                torch.sum(
+                    torch.square(
+                        self.value_model.critic(mb_obs_dict["critic_obs"])
+                        - self.value_model.critic(
+                            {"critic_obs": self._flip_obs(mb_obs_dict["critic_obs"], "critic_obs")}
+                        )
+                    ),
+                    dim=-1,
+                )
+            )
+            loss = (
+                pg_loss
+                + args.vf_coef * vf_loss
+                + self.entropy_coef * entropy_loss
+                + actor_sym_loss
+                + critic_sym_loss
+            )
+        else:
+            loss = pg_loss + args.vf_coef * vf_loss + self.entropy_coef * entropy_loss
+
+        if torch.isnan(loss) or torch.isinf(loss):
+            print(f"Invalid loss detected: {loss}")  # noqa: T201
+            print(
+                f"Ratio stats: min={ratio.min()}, max={ratio.max()}, mean={ratio.mean()}"
+            )  # noqa: T201
+            print(
+                f"Advantage stats: min={mb_advantage.min()}, max={mb_advantage.max()}"
+            )  # noqa: T201
+            # Skip this update or use previous valid parameters
+
+        loss_dict = {
+            "ppo_loss": loss,
+            # logging metrics
+            "local_kl_mean": local_kl_mean,
+            "pg_losses": pg_losses,
+            "pg_losses2": pg_losses2,
+            "pg_loss": pg_loss,
+            "vf_loss": vf_loss,
+            "entropy_loss": entropy_loss,
+            "ratio": ratio,
+            "vf_clipfrac": vf_clipfrac,
+        }
+        if self.use_symmetry:
+            loss_dict["actor_sym_loss"] = actor_sym_loss
+            loss_dict["critic_sym_loss"] = critic_sym_loss
+        return loss_dict
+
+    def _compute_imgaug_bc_loss(self, forward_results, mb_rollout_data):  # noqa: ARG002
+        """Compute behavior cloning loss between augmented and non-augmented action means.
+
+        Encourages the policy to produce similar actions regardless of image
+        augmentation, improving sim-to-real visual transfer.
+
+        Args:
+            forward_results: Output of ``_forward_model`` (must use
+                ``"policy_w_and_wo_imgaug"`` mode).
+            mb_rollout_data: Micro-batch dict (unused directly, kept for API
+                consistency).
+
+        Returns:
+            Dict with ``"imgaug_bc_loss"`` scalar.
+        """
+        policy_results = forward_results["policy_results"]
+        mu_batch = policy_results["action_mean"]
+
+        action_mean_w_imgaug = policy_results["action_mean_w_imgaug"]
+        imgaug_bc_loss = self.imgaug_bc_loss_fn(action_mean_w_imgaug, mu_batch.detach())
+
+        return {
+            "imgaug_bc_loss": imgaug_bc_loss,
+        }
+
+    def _update_stats_buffer(
+        self,
+        ppo_epoch_idx,
+        minibatch_idx,
+        microbatch_idx,
+        loss_dict,
+        forward_results,  # noqa: ARG002
+        mb_rollout_data,
+    ):
+        """Record per-update diagnostic statistics into the pre-allocated stat buffers.
+
+        Args:
+            ppo_epoch_idx: Current PPO epoch index.
+            minibatch_idx: Current mini-batch index within the epoch.
+            microbatch_idx: Current micro-batch index within the mini-batch.
+            loss_dict: Output of ``_compute_loss``.
+            forward_results: Output of ``_forward_model``.
+            mb_rollout_data: Micro-batch dict from ``_get_mb_rollout_data``.
+        """
+        local_kl_mean = loss_dict["ppo_loss_dict"]["local_kl_mean"]
+        pg_losses = loss_dict["ppo_loss_dict"]["pg_losses"].mean(dim=-1)
+        pg_losses2 = loss_dict["ppo_loss_dict"]["pg_losses2"].mean(dim=-1)
+        pg_loss = loss_dict["ppo_loss_dict"]["pg_loss"]
+        vf_loss = loss_dict["ppo_loss_dict"]["vf_loss"]
+        entropy_loss = loss_dict["ppo_loss_dict"]["entropy_loss"]
+        weighted_ppo_loss = loss_dict["ppo_loss_dict"]["ppo_loss"] * self.config.get(
+            "ppo_loss_coef", 1.0
+        )
+        ratio = loss_dict["ppo_loss_dict"]["ratio"]
+        vf_clipfrac = loss_dict["ppo_loss_dict"]["vf_clipfrac"]
+
+        padding_mask = mb_rollout_data["mb_padding_mask"]
+        micro_batch_inds = mb_rollout_data["micro_batch_inds"]  # noqa: F841
+
+        self.approxkl_stats[ppo_epoch_idx, minibatch_idx, microbatch_idx] = local_kl_mean
+        pg_clipfrac = masked_mean((pg_losses2 > pg_losses).float(), ~padding_mask)
+        self.pg_clipfrac_stats[ppo_epoch_idx, minibatch_idx, microbatch_idx] = pg_clipfrac
+        self.pg_loss_stats[ppo_epoch_idx, minibatch_idx, microbatch_idx] = pg_loss
+        self.vf_loss_stats[ppo_epoch_idx, minibatch_idx, microbatch_idx] = vf_loss
+        if self.compute_imgaug_bc_loss:
+            imgaug_bc_loss = loss_dict["imgaug_bc_loss_dict"]["imgaug_bc_loss"]
+            self.imgaug_bc_loss_stats[ppo_epoch_idx, minibatch_idx, microbatch_idx] = imgaug_bc_loss
+            self.weighted_imgaug_bc_loss_stats[ppo_epoch_idx, minibatch_idx, microbatch_idx] = (
+                self.config.imgaug_bc_loss_coef * imgaug_bc_loss
+            )
+        if self.use_symmetry:
+            self.actor_sym_loss_stats[ppo_epoch_idx, minibatch_idx, microbatch_idx] = loss_dict[
+                "ppo_loss_dict"
+            ]["actor_sym_loss"]
+            self.critic_sym_loss_stats[ppo_epoch_idx, minibatch_idx, microbatch_idx] = loss_dict[
+                "ppo_loss_dict"
+            ]["critic_sym_loss"]
+            self.estimation_loss_stats[ppo_epoch_idx, minibatch_idx, microbatch_idx] = loss_dict[
+                "ppo_loss_dict"
+            ]["estimation_loss"]
+            self.swap_loss_stats[ppo_epoch_idx, minibatch_idx, microbatch_idx] = loss_dict[
+                "ppo_loss_dict"
+            ]["swap_loss"]
+        self.entropy_stats[ppo_epoch_idx, minibatch_idx, microbatch_idx] = -entropy_loss
+        self.weighted_ppo_loss_stats[ppo_epoch_idx, minibatch_idx, microbatch_idx] = (
+            weighted_ppo_loss
+        )
+        self.vf_clipfrac_stats[ppo_epoch_idx, minibatch_idx, microbatch_idx] = vf_clipfrac
+        self.ratio_stats[ppo_epoch_idx, minibatch_idx, microbatch_idx] = ratio.mean()
+        self.advantage_mean_stats[ppo_epoch_idx, minibatch_idx, microbatch_idx] = mb_rollout_data[
+            "mb_advantage"
+        ].mean()
+        self.advantage_std_stats[ppo_epoch_idx, minibatch_idx, microbatch_idx] = mb_rollout_data[
+            "mb_advantage"
+        ].std()
+
+    def _get_train_metrics(self):
+        """Gather and aggregate training statistics from all processes.
+
+        Returns:
+            Dict of scalar metrics (approx KL, clip fractions, losses, entropy,
+            ratios, advantage stats) averaged across all GPUs and update steps.
+        """
+        metrics = {}
+
+        approxkl_avg = self.accelerator.gather_for_metrics(self.approxkl_stats).mean().item()
+
+        metrics["policy/approxkl_avg"] = approxkl_avg
+        metrics["policy/clipfrac_avg"] = (
+            self.accelerator.gather_for_metrics(self.pg_clipfrac_stats).mean().item()
+        )
+        metrics["loss/policy_avg"] = (
+            self.accelerator.gather_for_metrics(self.pg_loss_stats).mean().item()
+        )
+        if self.compute_imgaug_bc_loss:
+            metrics["loss/imgaug_bc_avg"] = (
+                self.accelerator.gather_for_metrics(self.imgaug_bc_loss_stats).mean().item()
+            )
+            metrics["loss/weighted_imgaug_bc_avg"] = (
+                self.accelerator.gather_for_metrics(self.weighted_imgaug_bc_loss_stats)
+                .mean()
+                .item()
+            )
+        if self.use_symmetry:
+            metrics["loss/actor_sym"] = (
+                self.accelerator.gather_for_metrics(self.actor_sym_loss_stats).mean().item()
+            )
+            metrics["loss/critic_sym"] = (
+                self.accelerator.gather_for_metrics(self.critic_sym_loss_stats).mean().item()
+            )
+            metrics["loss/estimation"] = (
+                self.accelerator.gather_for_metrics(self.estimation_loss_stats).mean().item()
+            )
+            metrics["loss/swap"] = (
+                self.accelerator.gather_for_metrics(self.swap_loss_stats).mean().item()
+            )
+        metrics["loss/value_avg"] = (
+            self.accelerator.gather_for_metrics(self.vf_loss_stats).mean().item()
+        )
+        metrics["loss/entropy_avg"] = (
+            self.accelerator.gather_for_metrics(self.entropy_stats).mean().item()
+        )
+        metrics["loss/weighted_ppo_loss_avg"] = (
+            self.accelerator.gather_for_metrics(self.weighted_ppo_loss_stats).mean().item()
+        )
+        metrics["val/clipfrac_avg"] = (
+            self.accelerator.gather_for_metrics(self.vf_clipfrac_stats).mean().item()
+        )
+        metrics["val/ratio"] = self.accelerator.gather_for_metrics(self.ratio_stats).mean().item()
+        metrics["val/ratio_var"] = (
+            self.accelerator.gather_for_metrics(self.ratio_stats).var().item()
+        )
+        metrics["val/advantage_mean"] = (
+            self.accelerator.gather_for_metrics(self.advantage_mean_stats).mean().item()
+        )
+        metrics["val/advantage_std"] = (
+            self.accelerator.gather_for_metrics(self.advantage_std_stats).mean().item()
+        )
+        metrics["objective/entropy"] = metrics["loss/entropy_avg"]
+
+        return metrics
+
+    def train(self):
+        """Run the full PPO training loop until ``num_total_batches`` iterations.
+
+        Each iteration:
+            1. Collect ``num_steps_per_env`` transitions via ``_rollout_step``.
+            2. Compute GAE returns and advantages.
+            3. Run ``num_ppo_epochs`` of mini-batch gradient updates.
+            4. Synchronize running mean/std and adaptive sampling across GPUs.
+            5. Log metrics and invoke ``on_step_end`` callbacks (which handle
+               checkpointing, evaluation, and early stopping).
+        """
+        args = self.args
+        accelerator = self.accelerator
+        optimizer = self.optimizer
+        model = self.model
+        dataloader = self.dataloader
+        device = accelerator.device
+
+        def repeat_generator():
+            while True:
+                if dataloader is not None:
+                    yield from dataloader
+                else:
+                    yield None
+
+        iter_dataloader = iter(repeat_generator())
+
+        accelerator.print("===training policy===")
+        start_time = time.time()
+        self._register_stats_buffer()
+        model.train()
+
+        # trainer state initialization
+        self.state.max_steps = args.num_total_batches
+        self.state.num_train_epochs = args.total_episodes / self.train_dataset_len
+        # Compute absolute values for logging, eval, and save if given as ratio
+        if args.logging_steps is not None:
+            if args.logging_steps < 1:
+                self.state.logging_steps = math.ceil(self.state.max_steps * args.logging_steps)
+            else:
+                self.state.logging_steps = args.logging_steps
+        if args.eval_steps is not None:
+            if args.eval_steps < 1:
+                self.state.eval_steps = math.ceil(self.state.max_steps * args.eval_steps)
+            else:
+                self.state.eval_steps = args.eval_steps
+        if args.save_steps is not None:
+            if args.save_steps < 1:
+                self.state.save_steps = math.ceil(self.state.max_steps * args.save_steps)
+            else:
+                self.state.save_steps = args.save_steps
+        self.control = self.callback_handler.on_train_begin(args, self.state, self.control)
+
+        # backward compatibility
+        if self.is_deepspeed_enabled:
+            self.deepspeed = self.model
+            self.model_wrapped = self.model
+
+        # env
+        obs_dict = self.env.reset_all()
+        for obs_key in obs_dict.keys():  # noqa: SIM118
+            obs_dict[obs_key] = obs_dict[obs_key].to(device)
+
+        for batch_idx in range(1, args.num_total_batches + 1):
+            batch_start_time = time.time()
+            self.state.episode += 1 * args.batch_size
+            data = next(iter_dataloader)  # noqa: F841
+
+            # update scheduled params
+            if self.schedule_dict is not None:
+                self.scheduled_params_dict = scheduler.update_scheduled_params(
+                    self, self.schedule_dict, self.state.global_step
+                )
+
+            reinit_dr_freq = self.env.config.get("reinit_dr_freq", 0)
+            if reinit_dr_freq > 0 and self.state.global_step % reinit_dr_freq == 0:
+                self.env.reinit_dr()
+                if self.env.config.get("reset_on_reinit_dr", False):
+                    obs_dict = self.env.reset_all()
+                    for obs_key in obs_dict.keys():  # noqa: SIM118
+                        obs_dict[obs_key] = obs_dict[obs_key].to(device)
+
+            with torch.no_grad():
+                with models_utils.unwrap_model_for_generation(
+                    self.model,
+                    self.accelerator,
+                    gather_deepspeed3_params=self.args.ds3_gather_for_generation,
+                ) as model:
+                    obs_dict = self._rollout_step(model, obs_dict)
+
+                end_collection_time = time.time()
+                collection_time = end_collection_time - batch_start_time
+
+            with common.Timer("get_rollout_data"):
+                rollout_data = self._get_rollout_data(obs_keys=obs_dict.keys())
+
+                model = self.model
+                self._train_mode()
+            with common.Timer("ppo_training"):
+                for ppo_epoch_idx in range(args.num_ppo_epochs):
+                    minibatch_idx = 0
+                    if self.ppo_shuffle_every_epoch or ppo_epoch_idx == 0:
+                        b_inds = torch.randperm(args.local_batch_size, device=device)
+                    for mini_batch_start in range(
+                        0, args.local_batch_size, args.local_mini_batch_size
+                    ):
+                        mini_batch_end = mini_batch_start + args.local_mini_batch_size
+                        mini_batch_inds = b_inds[mini_batch_start:mini_batch_end]
+                        microbatch_idx = 0
+                        for micro_batch_start in range(
+                            0, args.local_mini_batch_size, args.per_device_train_batch_size
+                        ):
+                            with common.Timer(  # noqa: SIM117
+                                f"ppo_microbatch_{ppo_epoch_idx}_{minibatch_idx}_{microbatch_idx}"
+                            ):
+                                with accelerator.accumulate(model):
+                                    with common.Timer("get_mb_rollout_data"):
+                                        micro_batch_end = (
+                                            micro_batch_start + args.per_device_train_batch_size
+                                        )
+                                        micro_batch_inds = mini_batch_inds[
+                                            micro_batch_start:micro_batch_end
+                                        ]
+                                        mb_rollout_data = self._get_mb_rollout_data(
+                                            rollout_data, micro_batch_inds
+                                        )
+
+                                    if self.use_symmetry:
+                                        estimation_loss, swap_loss = (
+                                            self.policy_model.update_estimator(
+                                                mb_rollout_data["mb_obs_dict"]["actor_obs"],
+                                                mb_rollout_data["mb_next_critic_obs"],
+                                                self.args.learning_rate,
+                                            )
+                                        )
+
+                                    with common.Timer("forward_model"):
+                                        forward_results = self._forward_model(
+                                            model, mb_rollout_data
+                                        )
+
+                                    with common.Timer("compute_loss"):
+                                        loss_dict = self._compute_loss(
+                                            forward_results, mb_rollout_data
+                                        )
+
+                                    with common.Timer("backward"):
+                                        accelerator.backward(loss_dict["loss"])
+
+                                    with common.Timer("gradient_clipping"):
+                                        grad_norm = self._gradient_clipping()
+
+                                    if grad_norm is not None:
+                                        with common.Timer("optimizer_step"):
+                                            optimizer.step()
+                                    else:
+                                        print("NaN in gradient! Skipped!!!!")  # noqa: T201
+
+                                    optimizer.zero_grad()
+                                    with torch.no_grad():
+                                        if self.use_symmetry:
+                                            loss_dict["ppo_loss_dict"][
+                                                "estimation_loss"
+                                            ] = estimation_loss
+                                            loss_dict["ppo_loss_dict"]["swap_loss"] = swap_loss
+                                        with common.Timer("update_stats_buffer"):
+                                            self._update_stats_buffer(
+                                                ppo_epoch_idx,
+                                                minibatch_idx,
+                                                microbatch_idx,
+                                                loss_dict,
+                                                forward_results,
+                                                mb_rollout_data,
+                                            )
+                                    del loss_dict, forward_results, mb_rollout_data
+                                    microbatch_idx += 1
+                        minibatch_idx += 1  # noqa: SIM113
+                        # del everything and empty cache
+                        # fmt: off
+                    # if self.empty_cache_every_n_ppo_epoch > 0 and (ppo_epoch_idx + 1) % self.empty_cache_every_n_ppo_epoch == 0:  # noqa: E501
+                    #     # print(f"Empty cache at ppo_epoch_idx {ppo_epoch_idx}")
+                    #     gc.collect()
+                    #     torch.cuda.empty_cache()
+            ######################################################### Sync Running Mean Std #########################################################  # noqa: E501
+            with common.Timer("sync_running_mean_std"):
+                self.sync_running_mean_std()
+
+            with common.Timer("sync_adaptive_sampling"):
+                self.sync_adaptive_sampling()
+
+            # print(self.accelerator.process_index, self.model.module.policy.running_mean_std.running_mean.mean(), self.model.module.policy.running_mean_std.running_var.mean(), self.model.module.policy.running_mean_std.count)  # noqa: E501
+            # print(self.accelerator.process_index, self.policy_model.running_mean_std.running_mean)
+            # print('--------------------------------')
+            # print('**********', self.policy_model.running_mean_std.running_mean)
+
+            # print(self.accelerator.process_index, self.policy_model.running_mean_std._normalizer.state_dict())
+            # print('--------------------------------')
+            # if self.state.global_step == 3:
+            #     if self.accelerator.is_main_process :
+            #         import ipdb; ipdb.set_trace()
+            #     else:
+            #         time.sleep(100000000)
+
+            ######################################################### Sync Running Mean Std #########################################################  # noqa: E501
+
+            with torch.no_grad():
+                learn_time = time.time() - end_collection_time
+                eps = int(self.state.episode / (time.time() - start_time))
+
+                metrics = {}
+                train_metrics = self._get_train_metrics()
+                metrics.update(train_metrics)
+                metrics["eps"] = eps
+                metrics["objective/rewards"] = (
+                    self.accelerator.gather_for_metrics(
+                        torch.tensor(np.mean(np.array(self.state.rewbuffer).sum(axis=-1))).to(
+                            device
+                        )
+                    )
+                    .mean()
+                    .item()
+                )
+                metrics["objective/length"] = (
+                    self.accelerator.gather_for_metrics(
+                        torch.tensor(np.mean(self.state.lenbuffer)).to(device)
+                    )
+                    .mean()
+                    .item()
+                )
+                metrics["lr"] = self.args.learning_rate
+                metrics["episode"] = self.state.episode
+                env_log_dict = self.episode_env_tensors.mean_and_clear()
+
+                ep_infos = process_ep_infos(self.ep_infos, device)
+                self.state.tot_timesteps += (
+                    self.num_steps_per_env * self.env.num_envs * accelerator.num_processes
+                )
+                self.state.tot_time += collection_time + learn_time
+                self.state.epoch = self.state.episode / self.train_dataset_len  # used by self.log
+                self.state.global_step += 1
+                log_dict = {
+                    "collection_time": collection_time,
+                    "learn_time": learn_time,
+                    "tot_timesteps": self.state.tot_timesteps,
+                    "tot_time": self.state.tot_time,
+                    "it": self.state.global_step,
+                    "fps": int(
+                        self.num_steps_per_env
+                        * self.env.num_envs
+                        * accelerator.num_processes
+                        / (collection_time + learn_time)
+                    ),
+                    "experiment_save_dir": self.args.output_dir,
+                    "batch_idx": batch_idx,
+                    "num_total_batches": args.num_total_batches,
+                }
+
+                for key, value in ep_infos.items():
+                    log_dict[f"Episode/{key}"] = value
+
+                # Add scheduled parameters to metrics
+                for param_name, param_value in self.scheduled_params_dict.items():
+                    log_dict[f"scheduled_params/{param_name}"] = param_value
+
+                if hasattr(self.policy_model, "std"):
+                    metrics["Policy/mean_noise_std"] = self.policy_model.std.mean().item()
+                else:
+                    metrics["Policy/mean_noise_std"] = 0.0
+                self.append_to_log_dict(log_dict)
+                metrics.update({f"Env/{k}": v for k, v in env_log_dict.items()})
+                metrics.update(env_log_dict)
+                metrics.update(log_dict)
+
+                self.log(metrics)
+                self.ep_infos.clear()
+
+            self.lr_scheduler.step()
+
+            del metrics, rollout_data
+            # Skip pre-iteration GC here; motion loading performs fallback cleanup,
+            # which avoids extra synchronization overhead and improves training speed.
+            # gc.collect()
+            # torch.cuda.empty_cache()
+
+            self.control = self.callback_handler.on_step_end(args, self.state, self.control)
+
+            if self.control.should_training_stop:
+                break
+
+        if self.control.should_training_stop:
+            return
+
+        # HF trainer specifics
+        self.control = self.callback_handler.on_train_end(args, self.state, self.control)
+        if self.control.should_save:
+            self._save_checkpoint(model, trial=None, metrics=None)
+            self.control = self.callback_handler.on_save(self.args, self.state, self.control)
+
+        if common.wandb_run_exists():
+            wandb.finish()
+
+    def sync_running_mean_std(self):
+        """Synchronize observation running mean/std normalizers across all GPU processes.
+
+        Syncs every step for the first 200 iterations (warm-up), then at
+        ``sync_running_mean_std_freq`` intervals.
+        """
+        sync_running_mean_std_freq = self.env.config.get("sync_running_mean_std_freq", 1)
+        if self.state.global_step < 200 or (
+            sync_running_mean_std_freq > 0
+            and (self.state.global_step + 1) % sync_running_mean_std_freq == 0
+        ):
+            if (
+                hasattr(self.policy_model, "use_running_mean_std")
+                and self.policy_model.use_running_mean_std
+            ):
+                # print(f"Syncing policy running mean std at global step {self.state.global_step}")
+                self.accelerator.wait_for_everyone()
+                self.policy_model.running_mean_std.sync_across_gpus(self.accelerator)
+            if (
+                hasattr(self.value_model, "use_running_mean_std")
+                and self.value_model.use_running_mean_std
+            ):
+                self.accelerator.wait_for_everyone()
+                self.value_model.running_mean_std.sync_across_gpus(self.accelerator)
+
+    def sync_adaptive_sampling(self):
+        """Synchronize adaptive motion sampling weights across GPU processes."""
+        sync_adaptive_sampling_all_gpus_freq = self.env.config.get(
+            "sync_adaptive_sampling_all_gpus_freq", 200
+        )
+        sync_across_gpus = (
+            sync_adaptive_sampling_all_gpus_freq > 0
+            and (self.state.global_step + 1) % sync_adaptive_sampling_all_gpus_freq == 0
+        )
+        if hasattr(self.env, "sync_and_compute_adaptive_sampling"):
+            self.env.sync_and_compute_adaptive_sampling(
+                self.accelerator, sync_across_gpus=sync_across_gpus
+            )
+
+    def append_to_log_dict(self, log_dict):
+        """Hook for subclasses to inject additional entries into the per-iteration log dict."""
+        pass
+
+    def _eval_mode(self):
+        """Switch models and environment to evaluation mode."""
+        self.model.eval()
+        model = self.accelerator.unwrap_model(self.model)
+        model.set_mode("eval")
+        model.transform_eval()
+        self.env.set_is_evaluating(is_evaluating=True, log_info=False)
+
+    def _train_rollout_mode(self):
+        """Switch to rollout collection mode: model in eval, transforms off, env in train."""
+        self.model.eval()
+        model = self.accelerator.unwrap_model(self.model)
+        model.set_mode("train_rollout")
+        model.transform_eval()
+        if self.train_with_evaluating_env:
+            self.env.set_is_evaluating(is_evaluating=True, log_info=False)
+        else:
+            self.env.set_is_evaluating(is_evaluating=False, log_info=False)
+
+    def _train_mode(self):
+        """Switch to gradient-update mode: model in train, transforms on."""
+        self.model.train()
+        model = self.accelerator.unwrap_model(self.model)
+        model.set_mode("train")
+        model.transform_train()
+        if self.train_with_evaluating_env:
+            self.env.set_is_evaluating(is_evaluating=True, log_info=False)
+        else:
+            self.env.set_is_evaluating(is_evaluating=False, log_info=False)
+
+    def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
+        """
+        Log `logs` on the various objects watching training.
+
+        Subclass and override this method to inject custom behavior.
+
+        Args:
+            logs (`dict[str, float]`):
+                The values to log.
+            start_time (`Optional[float]`):
+                The start of training.
+        """  # noqa: D212
+        if self.state.epoch is not None:
+            logs["epoch"] = self.state.epoch
+        if self.args.include_num_input_tokens_seen:
+            logs["num_input_tokens_seen"] = self.state.num_input_tokens_seen
+            if start_time is not None:
+                speed_metrics(  # noqa: F405
+                    "train", start_time, num_tokens=self.state.num_input_tokens_seen
+                )
+
+        # Sanitize all caller logs at this boundary: rank 0 stores only detached CPU/Python values.
+        if self.state.is_world_process_zero:
+            output = {}
+            for key, value in logs.items():
+                if isinstance(value, torch.Tensor):
+                    value = value.detach().cpu().item()
+                elif isinstance(value, np.ndarray):
+                    value = float(value)
+                output[key] = value
+            output["step"] = self.state.global_step
+            self.state.log_history.append(output)
+
+        self.control = self.callback_handler.on_log(self.args, self.state, self.control, logs)
+
+    def _gradient_clipping(self):
+        """Clip gradients and detect NaN/Inf, skipping the update if found.
+
+        Returns:
+            The global gradient norm after clipping, or None if NaN/Inf
+            gradients were detected (signaling the caller to skip the
+            optimizer step).
+        """
+        args = self.args
+        model = self.model
+
+        # Check for NaN/Inf in gradients
+        for name, param in model.named_parameters():
+            if param.grad is not None and (
+                torch.isnan(param.grad).any() or torch.isinf(param.grad).any()
+            ):
+                print(  # noqa: T201
+                    f"[Rank {self.accelerator.process_index}] NaN/Inf grad in {name}, norm={param.grad.norm():.3e}"
+                )
+                self.optimizer.zero_grad()
+                return None
+        grad_norm = None
+        if args.max_grad_norm is not None and args.max_grad_norm > 0:
+            # deepspeed does its own clipping
+
+            if is_sagemaker_mp_enabled() and args.fp16:  # noqa: F405
+                _grad_norm = self.optimizer.clip_master_grads(args.max_grad_norm)
+            elif self.use_apex:
+                # Revert to normal clipping otherwise, handling Apex or full precision
+                _grad_norm = nn.utils.clip_grad_norm_(
+                    amp.master_params(self.optimizer),  # noqa: F405
+                    args.max_grad_norm,
+                )
+            else:
+                _grad_norm = self.accelerator.clip_grad_norm_(
+                    model.parameters(),
+                    args.max_grad_norm,
+                )
+
+            if (
+                transformers_utils.is_accelerate_available()
+                and self.accelerator.distributed_type == accelerate.DistributedType.DEEPSPEED
+            ):
+                grad_norm = model.get_global_grad_norm()
+                # In some cases the grad norm may not return a float
+                if hasattr(grad_norm, "item"):
+                    grad_norm = grad_norm.item()
+            else:
+                grad_norm = _grad_norm
+
+        return grad_norm
+
+    def _compute_returns(self, values, last_values, policy_state_dict):
+        """Compute the returns and advantages for the given policy state.
+        This function calculates the returns and advantages for each step in the
+        environment based on the provided observations and policy state. It uses
+        Generalized Advantage Estimation (GAE) to compute the advantages, which
+        helps in reducing the variance of the policy gradient estimates.
+        Args:
+            values (torch.Tensor): The values for each step.
+            last_values (torch.Tensor): The last values for the last step.
+            policy_state_dict (dict): A dictionary containing the policy state
+                          information, including 'values', 'dones',
+                          and 'rewards'.
+        Returns:
+            tuple: A tuple containing:
+            - returns (torch.Tensor): The computed returns for each step.
+            - advantages (torch.Tensor): The normalized advantages for each step.
+        """  # noqa: D205, D410, D411
+        device = self.accelerator.device
+        advantage = 0
+
+        dones = policy_state_dict["dones"]
+        rewards = policy_state_dict["rewards"]
+
+        dones = dones.to(device)
+        rewards = rewards.to(device)
+
+        returns = torch.zeros_like(values)
+
+        num_steps = returns.shape[0]
+
+        for step in reversed(range(num_steps)):
+            if step == num_steps - 1:  # noqa: SIM108
+                next_values = last_values
+            else:
+                next_values = values[step + 1]
+            next_is_not_terminal = 1.0 - dones[step].float()
+            delta = rewards[step] + next_is_not_terminal * self.gamma * next_values - values[step]
+            advantage = delta + next_is_not_terminal * self.gamma * self.lam * advantage
+            returns[step] = advantage + values[step]
+
+        # Compute and normalize the advantages
+        advantages = returns - values
+        if self.sync_advantage_normalization:
+            # gather advantages from all processes before normalization
+            advantages = self.accelerator.gather(advantages)
+            advantages = (advantages - advantages.mean(dim=(0, 1), keepdim=True)) / (
+                advantages.std(dim=(0, 1), keepdim=True) + 1e-8
+            )
+            # ungather advantages
+            advantages = advantages.reshape(
+                self.accelerator.num_processes, -1, *advantages.shape[1:]
+            )[self.accelerator.process_index].to(device)
+        else:
+            advantages = (advantages - advantages.mean(dim=(0, 1), keepdim=True)) / (
+                advantages.std(dim=(0, 1), keepdim=True) + 1e-8
+            )
+        return returns, advantages
+
+    def _adjust_learning_rate_based_on_kl(self, kl_mean, optimizer):
+        """Adjust the learning rate based on the KL divergence.
+
+        This function implements a learning rate schedule that adjusts the learning rate
+        based on the KL divergence between the current policy and the old policy.
+        If the KL divergence is too high, the learning rate is decreased.
+        If the KL divergence is too low, the learning rate is increased.
+
+        Args:
+            kl_mean (float): The mean KL divergence across all processes.
+            optimizer (torch.optim.Optimizer): The optimizer to update.
+        """
+        if self.desired_kl is None:
+            return
+
+        if kl_mean > self.desired_kl * 2.0:
+            new_lr = max(self.adaptive_lr_min, self.args.learning_rate / 1.5)
+        elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0:
+            new_lr = min(self.adaptive_lr_max, self.args.learning_rate * 1.5)
+        else:
+            new_lr = self.args.learning_rate
+        self.args.learning_rate = new_lr
+
+        for param_group in optimizer.param_groups:
+            param_group["lr"] = self.args.learning_rate
+
+    def load_checkpoint(self, checkpoint_path, resume=False):  # noqa: D417
+        """Load a checkpoint to restore model weights and optionally full training state.
+
+        Args:
+            checkpoint_path: Path to the ``.pt`` checkpoint file.
+            resume: If True, also restore optimizer state, LR scheduler,
+                environment state, and trainer counters for seamless resumption.
+
+        Returns:
+            The loaded checkpoint dict.
+        """
+        print(f"Loading checkpoint from {checkpoint_path}")  # noqa: T201
+        checkpoint = torch.load(
+            checkpoint_path, map_location=self.accelerator.device, weights_only=False
+        )
+
+        # Load model state
+        model = self.accelerator.unwrap_model(self.model)
+        if "actor_model_state_dict" in checkpoint:
+            model.policy.load_state_dict(checkpoint["actor_model_state_dict"])
+        elif "policy_state_dict" in checkpoint:
+            model.policy.load_state_dict(checkpoint["policy_state_dict"], strict=False)
+        if "value_state_dict" in checkpoint and model.value_model is not None:
+            model.value_model.load_state_dict(checkpoint["value_state_dict"])
+
+        if resume:
+            # Load optimizer state
+            if (
+                "optimizer_state_dict" in checkpoint
+                and checkpoint["optimizer_state_dict"] is not None
+            ):
+                self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
+
+                # Update learning rate if available
+                if "args" in checkpoint and hasattr(checkpoint["args"], "learning_rate"):
+                    self.args.learning_rate = checkpoint["args"].learning_rate
+                    for param_group in self.optimizer.param_groups:
+                        param_group["lr"] = self.args.learning_rate
+
+            # Load learning rate scheduler state
+            if (
+                "lr_scheduler_state_dict" in checkpoint
+                and checkpoint["lr_scheduler_state_dict"] is not None
+            ):
+                self.lr_scheduler.load_state_dict(checkpoint["lr_scheduler_state_dict"])
+
+            if "env_state_dict" in checkpoint:
+                self.env.load_env_state_dict(checkpoint["env_state_dict"])
+
+            if "state" in checkpoint:
+                for key, value in checkpoint["state"].__dict__.items():
+                    if key in ["cur_reward_sum", "cur_episode_length"]:
+                        cur_value = getattr(self, key)
+                        if cur_value.shape != value.shape:
+                            continue
+                        setattr(self, key, value)
+                    if key not in [
+                        "stateful_callbacks",
+                        "is_local_process_zero",
+                        "is_world_process_zero",
+                        "log_history",
+                    ]:
+                        setattr(self.state, key, value)
+
+        print(f"Loaded checkpoint from step {checkpoint['state'].global_step}")  # noqa: T201
+        return checkpoint
+
+    def eval(self):
+        """Run an infinite deterministic evaluation loop with the current policy.
+
+        Resets the environment, then repeatedly queries the policy for mean
+        actions (no sampling) and steps the environment.  Intended for
+        interactive visualization; exits only when interrupted externally.
+        """
+        self._eval_mode()
+        self.env.set_is_evaluating()
+        self.model.policy.eval_mode()
+        obs_dict = self.env.reset_all()
+        for obs_key in obs_dict.keys():  # noqa: SIM118
+            obs_dict[obs_key] = obs_dict[obs_key].to(self.accelerator.device)
+
+        self.callback_handler.on_step_end(self.args, self.state, self.control)
+
+        with torch.no_grad():  # noqa: SIM117
+            with models_utils.unwrap_model_for_generation(
+                self.model,
+                self.accelerator,
+                gather_deepspeed3_params=self.args.ds3_gather_for_generation,
+            ) as model:
+                while True:
+                    device = self.accelerator.device
+                    policy_model = model.policy
+                    value_model = model.value_model  # noqa: F841
+                    policy_model.init_rollout()
+
+                    policy_state_dict = {}  # noqa: F841
+                    actor_state = {}
+                    actions = policy_model.rollout(obs_dict=obs_dict)  # noqa: F841
+                    action_mean = policy_model.action_mean.detach()
+
+                    actor_state["actions"] = action_mean
+                    results = self.env.step(actor_state)
+                    obs_dict, rewards, dones, infos = (
+                        results[0],
+                        results[1],
+                        results[2],
+                        results[3],
+                    )  # noqa: F841
+
+                    for obs_key in obs_dict.keys():  # noqa: SIM118
+                        obs_dict[obs_key] = obs_dict[obs_key].to(device)
+
+    @torch.no_grad()
+    def get_example_obs(self):
+        """Reset the environment and return a CPU observation dict for inspection or ONNX tracing.
+
+        Returns:
+            Dict mapping observation keys to CPU tensors ``(num_envs, obs_dim)``.
+        """
+        obs_dict = self.env.reset_all()
+        for obs_key in obs_dict.keys():  # noqa: SIM118
+            print(obs_key, sorted(self.env.config.obs.obs_dict[obs_key]))  # noqa: T201
+        # move to cpu
+        for k in obs_dict:
+            obs_dict[k] = obs_dict[k].cpu()
+        return obs_dict
+
+    @property
+    def inference_model(self):
+        return {"actor": self.model.policy, "critic": self.model.value_model}
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/trainer/ppo_trainer_aux_loss.py b/GR00T-WholeBodyControl/gear_sonic/trl/trainer/ppo_trainer_aux_loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..ea54219633c70aad9ecb18b22eced3e0ef5f9527
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/trainer/ppo_trainer_aux_loss.py
@@ -0,0 +1,306 @@
+import torch
+
+from gear_sonic.trl.trainer.ppo_trainer import TRLPPOTrainer
+
+
+class TRLAuxLossPPOTrainer(TRLPPOTrainer):
+    """PPO trainer extended with per-step auxiliary losses.
+
+    Subclasses :class:`TRLPPOTrainer` to add support for auxiliary losses
+    that the policy's forward pass returns alongside the standard PPO
+    objective.  Typical use-case is SONIC / universal token training where
+    the ``UniversalTokenModule`` emits reconstruction and latent-alignment
+    losses together with the action mean.
+
+    The total loss is::
+
+        loss = ppo_loss + aux_loss_scale * sum(coef_i * aux_loss_i)
+
+    Auxiliary losses and their per-loss coefficients are expected in the
+    ``policy_results`` dict under the keys ``"aux_losses"`` and
+    ``"aux_loss_coef"`` respectively.
+
+    Config keys (read from ``self.config``):
+
+    * ``aux_loss_scale`` (float, default 1.0) – global scale applied to the
+      weighted sum of auxiliary losses.
+    * ``compute_aux_loss`` (bool, default ``True``) – disable to skip all
+      auxiliary loss computation (useful for ablations).
+    """
+
+    _tag_names = ["trl", "aux_loss_ppo"]
+
+    def _init_config(self):
+        """Extend base config initialisation with auxiliary loss settings.
+
+        Reads ``aux_loss_scale`` and ``compute_aux_loss`` from
+        ``self.config`` and stores them as instance attributes.
+        """
+        super()._init_config()
+
+        # Auxiliary loss configuration
+        # aux_loss_scale: overall scalar to scale the total auxiliary loss
+        self.aux_loss_scale = self.config.get("aux_loss_scale", 1.0)
+        self.compute_aux_loss = self.config.get("compute_aux_loss", True)
+
+    def _register_stats_buffer(self):
+        """Allocate per-step statistics tensors for auxiliary losses.
+
+        Calls the parent method first to register base PPO stats, then
+        allocates the following additional buffers when
+        ``self.compute_aux_loss`` is ``True``:
+
+        * ``self.aux_loss_stats`` – empty dict; individual loss tensors are
+          added lazily on first occurrence (see :meth:`_update_stats_buffer`).
+        * ``self.total_aux_loss_unscaled_stats`` – shape
+          ``(num_ppo_epochs, num_mini_batches, num_micro_batches)``, stores
+          the coefficient-weighted sum before the global scale.
+        * ``self.total_aux_loss_stats`` – same shape, stores the fully scaled
+          total auxiliary loss.
+        """
+        super()._register_stats_buffer()
+
+        if self.compute_aux_loss:
+            args = self.args
+            device = self.accelerator.device
+
+            stats_shape = (args.num_ppo_epochs, args.num_mini_batches, args.num_micro_batches)
+            # Store stats as dictionaries to support multiple auxiliary losses
+            self.aux_loss_stats = {}
+            self.total_aux_loss_unscaled_stats = torch.zeros(stats_shape, device=device)
+            self.total_aux_loss_stats = torch.zeros(stats_shape, device=device)
+
+    def _extract_aux_losses_from_forward_results(self, forward_results):
+        """Pull auxiliary losses and their coefficients from the policy output dict.
+
+        Args:
+            forward_results: Dict returned by the policy's forward pass.  The
+                following optional keys are consumed:
+
+                * ``"aux_losses"`` – dict mapping loss name to scalar tensor.
+                * ``"aux_loss_coef"`` – dict mapping loss name to float
+                  coefficient; defaults to ``None`` when absent.
+
+        Returns:
+            Tuple of ``(aux_losses_dict, aux_loss_coef)`` where
+            ``aux_losses_dict`` maps loss name to tensor and
+            ``aux_loss_coef`` maps loss name to float (or ``None`` when the
+            key was not present in ``forward_results``).
+        """
+        aux_losses_dict = {}
+        aux_loss_coef = None
+
+        if "aux_losses" in forward_results:
+            aux_losses_dict = forward_results["aux_losses"]
+
+        # Extract coefficients if provided in forward_results
+        if "aux_loss_coef" in forward_results:
+            aux_loss_coef = forward_results["aux_loss_coef"]
+
+        return aux_losses_dict, aux_loss_coef
+
+    def _compute_aux_loss(self, policy_results, mb_rollout_data):
+        """Compute the weighted auxiliary loss for one mini-batch.
+
+        Extracts individual auxiliary losses and their coefficients from
+        ``policy_results``, computes the per-loss weighted sum, and applies
+        the global ``aux_loss_scale``.
+
+        Args:
+            policy_results: Dict from the policy forward pass (the value
+                stored at ``forward_results["policy_results"]``).  Expected
+                to contain ``"aux_losses"`` and optionally ``"aux_loss_coef"``.
+            mb_rollout_data: Mini-batch rollout dict (not used directly but
+                available for subclass overrides).
+
+        Returns:
+            Dict with keys:
+
+            * ``"aux_losses_dict"`` – raw loss tensors keyed by name.
+            * ``"aux_loss_coef"`` – per-loss coefficient dict.
+            * ``"total_aux_loss_unscaled"`` – weighted sum before
+              ``aux_loss_scale``, shape ``()``.
+            * ``"total_aux_loss"`` – final scaled auxiliary loss,
+              shape ``()``.
+        """
+        device = self.accelerator.device
+
+        # Extract auxiliary losses and coefficients from forward results
+        aux_losses_dict, aux_loss_coef = self._extract_aux_losses_from_forward_results(
+            policy_results
+        )
+
+        if not aux_losses_dict:
+            # No auxiliary losses found
+            return {
+                "aux_losses_dict": {},
+                "aux_loss_coef": {},
+                "total_aux_loss_unscaled": torch.tensor(0.0, device=device),
+                "total_aux_loss": torch.tensor(0.0, device=device),
+            }
+
+        # Compute weighted sum of auxiliary losses
+        total_aux_loss_unscaled = torch.tensor(0.0, device=device)
+        for loss_name, loss_value in aux_losses_dict.items():
+            coef = aux_loss_coef.get(loss_name, 0.0)
+            total_aux_loss_unscaled += coef * loss_value
+
+        # Apply overall scale
+        total_aux_loss = total_aux_loss_unscaled * self.aux_loss_scale
+
+        return {
+            "aux_losses_dict": aux_losses_dict,
+            "aux_loss_coef": aux_loss_coef,
+            "total_aux_loss_unscaled": total_aux_loss_unscaled,
+            "total_aux_loss": total_aux_loss,
+        }
+
+    def _compute_loss(self, forward_results, mb_rollout_data):
+        """Compute total training loss as PPO loss plus auxiliary loss.
+
+        Calls the parent ``_compute_loss`` for the standard PPO objective
+        (and optionally imitation BC loss), then adds the auxiliary loss
+        computed from ``forward_results["policy_results"]`` when
+        ``self.compute_aux_loss`` is ``True``.
+
+        Args:
+            forward_results: Dict returned by the full forward pass.  Must
+                contain ``"policy_results"`` (the policy module's output dict)
+                plus whatever the parent method expects.
+            mb_rollout_data: Mini-batch of rollout transitions used to
+                compute advantages, returns, and old log-probs.
+
+        Returns:
+            Dict with at minimum:
+
+            * ``"loss"`` – scalar total loss tensor (PPO + aux).
+            * ``"aux_loss_dict"`` – result dict from
+              :meth:`_compute_aux_loss` (only present when
+              ``compute_aux_loss`` is ``True``).
+            * All keys returned by the parent ``_compute_loss``.
+        """
+        # Compute PPO loss (includes ppo_loss and optionally imgaug_bc_loss)
+        loss_dict = super()._compute_loss(forward_results, mb_rollout_data)
+
+        # Compute and add auxiliary loss if enabled
+        if self.compute_aux_loss:
+            aux_loss_result = self._compute_aux_loss(
+                forward_results["policy_results"], mb_rollout_data
+            )
+
+            loss_dict["loss"] += aux_loss_result["total_aux_loss"]
+
+            # Add auxiliary loss dict to return dict
+            loss_dict["aux_loss_dict"] = aux_loss_result
+
+        return loss_dict
+
+    def _update_stats_buffer(
+        self,
+        ppo_epoch_idx,
+        minibatch_idx,
+        microbatch_idx,
+        loss_dict,
+        forward_results,
+        mb_rollout_data,
+    ):
+        """Record per-step loss values into pre-allocated statistics buffers.
+
+        Delegates to the parent method for base PPO stats, then writes
+        individual auxiliary loss values and aggregate totals into
+        ``self.aux_loss_stats``, ``self.total_aux_loss_unscaled_stats``, and
+        ``self.total_aux_loss_stats``.  Individual loss buffers are lazily
+        initialised on first encounter.
+
+        Args:
+            ppo_epoch_idx: Index of the current PPO epoch (0-based).
+            minibatch_idx: Index of the current mini-batch within the epoch.
+            microbatch_idx: Index of the current micro-batch within the
+                mini-batch (used for gradient accumulation).
+            loss_dict: Output of :meth:`_compute_loss` for this step.
+                Expected to contain ``"aux_loss_dict"`` when
+                ``compute_aux_loss`` is ``True``.
+            forward_results: Full forward-pass result dict (passed through to
+                the parent method).
+            mb_rollout_data: Mini-batch rollout dict (passed through to the
+                parent method).
+        """
+        # Update PPO stats
+        super()._update_stats_buffer(
+            ppo_epoch_idx,
+            minibatch_idx,
+            microbatch_idx,
+            loss_dict,
+            forward_results,
+            mb_rollout_data,
+        )
+
+        # Update auxiliary loss stats if enabled
+        if self.compute_aux_loss and "aux_loss_dict" in loss_dict:
+            aux_loss_result = loss_dict["aux_loss_dict"]
+            aux_losses_dict = aux_loss_result["aux_losses_dict"]
+            total_aux_loss_unscaled = aux_loss_result["total_aux_loss_unscaled"]
+            total_aux_loss = aux_loss_result["total_aux_loss"]
+
+            # Update stats for each individual auxiliary loss
+            for loss_name, loss_value in aux_losses_dict.items():
+                if loss_name not in self.aux_loss_stats:
+                    # Lazily initialize stats buffer for this loss
+                    args = self.args
+                    device = self.accelerator.device
+                    stats_shape = (
+                        args.num_ppo_epochs,
+                        args.num_mini_batches,
+                        args.num_micro_batches,
+                    )
+                    self.aux_loss_stats[loss_name] = torch.zeros(stats_shape, device=device)
+
+                self.aux_loss_stats[loss_name][
+                    ppo_epoch_idx, minibatch_idx, microbatch_idx
+                ] = loss_value
+
+            # Update total auxiliary loss stats
+            self.total_aux_loss_unscaled_stats[ppo_epoch_idx, minibatch_idx, microbatch_idx] = (
+                total_aux_loss_unscaled
+            )
+            self.total_aux_loss_stats[ppo_epoch_idx, minibatch_idx, microbatch_idx] = total_aux_loss
+
+    def _get_train_metrics(self):
+        """Collect training metrics including auxiliary loss averages.
+
+        Calls the parent method for standard PPO metrics, then appends:
+
+        * ``"loss/aux_{name}_avg"`` – mean of each individual auxiliary loss
+          across all PPO epochs, mini-batches, and micro-batches.
+        * ``"loss/total_aux_loss_unscaled_avg"`` – mean of the
+          coefficient-weighted sum before global scaling.
+        * ``"loss/total_aux_loss_avg"`` – mean of the fully scaled total
+          auxiliary loss.
+        * ``"aux_loss_scale"`` – the configured global scale factor.
+
+        Returns:
+            Dict of metric name → scalar value for the completed training
+            iteration, suitable for logging to W&B or TensorBoard.
+        """
+        metrics = super()._get_train_metrics()
+
+        # Add auxiliary loss metrics if enabled
+        if self.compute_aux_loss:
+            # Add metrics for each individual auxiliary loss
+            for loss_name, loss_stats in self.aux_loss_stats.items():
+                metrics[f"loss/aux_{loss_name}_avg"] = (
+                    self.accelerator.gather_for_metrics(loss_stats).mean().item()
+                )
+
+            # Add total auxiliary loss metrics
+            metrics["loss/total_aux_loss_unscaled_avg"] = (
+                self.accelerator.gather_for_metrics(self.total_aux_loss_unscaled_stats)
+                .mean()
+                .item()
+            )
+            metrics["loss/total_aux_loss_avg"] = (
+                self.accelerator.gather_for_metrics(self.total_aux_loss_stats).mean().item()
+            )
+            metrics["aux_loss_scale"] = self.aux_loss_scale
+
+        return metrics
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/__init__.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/__init__.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b1a2238384c427f763fc89af519132c95a4b2b45
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/__init__.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/__init__.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5b14fc70f580969604b30c7187f0e08274f0842f
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/__init__.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/common.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/common.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9654a5ee5bb595829a7126cf9f71017c277ba69f
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/common.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/common.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/common.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2326084647213e8a0c358837a776c564fe3f1093
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/common.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/kornia_transform.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/kornia_transform.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0309b39aac0bcbc14c8e30c5a699e8abba759cd6
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/kornia_transform.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/kornia_transform.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/kornia_transform.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..789e593d9a4e3fbfc73c8a255b56ea30af006e94
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/kornia_transform.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/order_converter.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/order_converter.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0320d1ece84fe91455ec44fc451bc7e33d75ca0b
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/order_converter.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/order_converter.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/order_converter.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..effe1212a529b8b847e6d535090bcff2f3eee260
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/order_converter.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/rl.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/rl.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5056632dc48fa823ca9d214372a4e44f52e95f87
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/rl.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/rl.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/rl.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5ed7a972605f1c905036b018caf35a3ea092f3df
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/rl.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/scheduler.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/scheduler.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..649a1e6043959fecbdc609ccec1f327ea6f28cdd
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/scheduler.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/scheduler.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/scheduler.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1ec52ed868bea2aa135b12ff5ec6ce15278fd10f
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/scheduler.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/torch_transform.cpython-311.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/torch_transform.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..022ae195f4a0323c424f7cc75abb64c03686951e
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/torch_transform.cpython-311.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/torch_transform.cpython-312.pyc b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/torch_transform.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ab933683714300b955173d2e720e23b961d31a49
Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/trl/utils/__pycache__/torch_transform.cpython-312.pyc differ
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/common.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/common.py
new file mode 100644
index 0000000000000000000000000000000000000000..b792b01455d2a2c8dffe1f44ff2454832b9b7b7a
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/common.py
@@ -0,0 +1,163 @@
+"""Miscellaneous training utilities: W&B helpers, dynamic imports, OmegaConf tools, and timers."""
+
+import wandb
+import importlib
+import os
+import time
+from omegaconf import OmegaConf, DictConfig, ListConfig
+
+
+def wandb_run_exists():
+    return isinstance(wandb.run, wandb.sdk.wandb_run.Run)
+
+
+def import_type_from_str(s):
+    module_name, type_name = s.rsplit(".", 1)
+    module = importlib.import_module(module_name)
+    type_to_import = getattr(module, type_name)
+    return type_to_import
+
+
+def recursive_set_struct(cfg, struct_value: bool):
+    OmegaConf.set_struct(cfg, struct_value)
+    if isinstance(cfg, DictConfig):
+        for key in cfg.keys():
+            try:
+                value = cfg[key]
+                if isinstance(value, (DictConfig, ListConfig)):
+                    recursive_set_struct(value, struct_value)
+            except Exception as e:
+                # print(e)
+                pass
+    elif isinstance(cfg, ListConfig):
+        for item in cfg:
+            if isinstance(item, (DictConfig, ListConfig)):
+                recursive_set_struct(item, struct_value)
+
+
+def materialize_lazy_params(policy, env):
+    """Materialize lazy parameters (nn.LazyLinear, nn.LazyConv2d) with a dummy forward pass.
+
+    Must be called before DDP wrapping, since accelerator.prepare() requires all params initialized.
+    Uses env.reset() with default flatten_dict_obs=True to get flat tensors (not sub-dicts).
+    """
+    import torch
+    import torch.nn as nn
+
+    if any(isinstance(m, (nn.LazyLinear, nn.LazyConv2d)) for m in policy.modules()):
+        dummy_obs = env.reset()
+        with torch.no_grad():
+            policy.act(dummy_obs)
+
+
+def get_filtered_state_dict(state_dict, state_dict_key):
+    """
+    Filter state_dict keys that start with the given prefix and remove the prefix.
+
+    Args:
+        state_dict: Dictionary of state dict keys and values
+        state_dict_key: Prefix string to filter by
+
+    Returns:
+        Filtered dictionary with prefix removed from keys
+    """
+    filtered_dict = {}
+    for key, value in state_dict.items():
+        if key.startswith(state_dict_key):
+            # Remove the prefix from the key
+            new_key = key[len(state_dict_key) :].lstrip(".")
+            filtered_dict[new_key] = value
+    return filtered_dict
+
+
+def custom_instantiate(d, _resolve=True, _recursive=False, **add_kwargs):
+    """
+    Recursively instantiate nested configs with _target_ fields.
+    """
+
+    def _recursive_instantiate(obj):
+        # If it's a dict and has a _target_, instantiate it
+        if isinstance(obj, dict) and "_target_" in obj:
+            if obj.get("_recursive_", None) == True:
+                assert False, "recursive is not supported"
+            obj = obj.copy()
+            obj.pop("_recursive_", None)
+            obj.pop("_convert_", None)
+            obj.pop("_partial_", None)
+            _type = import_type_from_str(obj.pop("_target_"))
+            # Recursively instantiate all dict/list values
+            for k, v in list(obj.items()):
+                if isinstance(v, (dict, DictConfig)):
+                    obj[k] = _recursive_instantiate(v)
+                elif isinstance(v, (list, ListConfig)):
+                    obj[k] = [_recursive_instantiate(i) for i in v]
+            return _type(**obj)
+        # If it's a dict, recursively instantiate its values
+        elif isinstance(obj, dict):
+            return {k: _recursive_instantiate(v) for k, v in obj.items()}
+        # If it's a list, recursively instantiate its items
+        elif isinstance(obj, list):
+            return [_recursive_instantiate(i) for i in obj]
+        else:
+            return obj
+
+    # Top-level: allow add_kwargs to override
+    d = d.copy()
+    if isinstance(d, DictConfig):
+        if _resolve:
+            d = OmegaConf.to_container(d, resolve=_resolve)
+        else:
+            recursive_set_struct(d, False)
+    if d.get("_recursive_", None) == True:
+        assert False, "recursive is not supported"
+    d.pop("_recursive_", None)
+    d.pop("_convert_", None)
+    d.pop("_partial_", None)
+    _type = import_type_from_str(d.pop("_target_"))
+    if _recursive:
+        # Recursively instantiate all dict/list values
+        for k, v in list(d.items()):
+            if isinstance(v, (dict, DictConfig)):
+                d[k] = _recursive_instantiate(v)
+            elif isinstance(v, (list, ListConfig)):
+                d[k] = [_recursive_instantiate(i) for i in v]
+    return _type(**d, **add_kwargs)
+
+
+# Global variable for timing indentation level
+timer_indent_level = 0
+
+
+# Context manager for timing
+class Timer:
+    def __init__(self, name="", instance_enabled=True):
+        self.name = name
+        self.start_time = None
+        self.enabled = instance_enabled and os.environ.get("TIMER_ENABLED", "0") == "1"
+        if "LOCAL_RANK" in os.environ:
+            self.rank = int(os.environ["LOCAL_RANK"])
+        else:
+            self.rank = 0
+        self.show_rank = os.environ.get("TIMER_SHOW_RANK", "0") == "1"
+        self.rank_zero_only = os.environ.get("TIMER_RANK_ZERO_ONLY", "0") == "1"
+
+    def __enter__(self):
+        if (not self.enabled) or (self.rank_zero_only and self.rank != 0):
+            return self
+        global timer_indent_level
+        self.start_time = time.perf_counter()
+        self.current_indent = timer_indent_level  # Capture current indent level
+        timer_indent_level += 1  # Increment global indent level for next call
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if exc_type:
+            return False  # Re-raise the exception
+        if (not self.enabled) or (self.rank_zero_only and self.rank != 0):
+            return self
+        global timer_indent_level
+        elapsed_time = time.perf_counter() - self.start_time
+        indent = "    " * self.current_indent  # 4 spaces per indent level
+        rank_str = f"[rank{self.rank}] " if self.show_rank else ""
+        print(f"{indent}{rank_str}[{self.name}] time: {elapsed_time:.4f} seconds")
+        timer_indent_level -= 1  # Decrement global indent level after finishing
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/data.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/data.py
new file mode 100644
index 0000000000000000000000000000000000000000..f7f10c529096ba651d7b93e93a87f1e8897fcb46
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/data.py
@@ -0,0 +1,16 @@
+"""Lightweight dataset helpers (dummy datasets, episode attention masks)."""
+
+from datasets import Dataset
+
+
+def create_dummy_dataset(num_samples: int = 100) -> Dataset:
+    """Create a dummy dataset with the specified number of samples.
+
+    Args:
+        num_samples (int): Number of samples to create in the dataset.
+
+    Returns:
+        Dataset: A HuggingFace Dataset containing dummy prompts.
+    """
+    dummy_data = {"prompt": [f"Sample prompt {i}" for i in range(num_samples)]}
+    return Dataset.from_dict(dummy_data)
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/kornia_transform.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/kornia_transform.py
new file mode 100644
index 0000000000000000000000000000000000000000..307321813b3331440358ccee958235b79c4e90d9
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/kornia_transform.py
@@ -0,0 +1,1079 @@
+"""Kornia-derived 2D/3D geometric transforms with optional JIT compilation.
+
+Angle, coordinate, quaternion, rotation-matrix, and axis-angle conversions
+used throughout GEAR-SONIC training and inference. Many functions are
+conditionally wrapped with torch.jit.script (controlled by the
+USE_JIT_TORCH_TRANSFORM env var) for performance.
+"""
+
+import enum
+import os
+import warnings
+from typing import Tuple
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+# Check environment variable to enable/disable torch.jit.script
+USE_JIT_TORCH_TRANSFORM = os.getenv("USE_JIT_TORCH_TRANSFORM", "1").lower() in ("1", "true", "yes")
+
+
+def conditional_jit_script(func):
+    """Conditionally apply torch.jit.script based on USE_JIT_TORCH_TRANSFORM env var"""
+    if USE_JIT_TORCH_TRANSFORM:
+        return torch.jit.script(func)
+    return func
+
+
+__all__ = [
+    # functional api
+    "rad2deg",
+    "deg2rad",
+    "pol2cart",
+    "cart2pol",
+    "convert_points_from_homogeneous",
+    "convert_points_to_homogeneous",
+    "convert_affinematrix_to_homography",
+    "convert_affinematrix_to_homography3d",
+    "angle_axis_to_rotation_matrix",
+    "angle_axis_to_quaternion",
+    "rotation_matrix_to_angle_axis",
+    "rotation_matrix_to_quaternion",
+    "quaternion_to_angle_axis",
+    "quaternion_to_rotation_matrix",
+    "quaternion_log_to_exp",
+    "quaternion_exp_to_log",
+    "denormalize_pixel_coordinates",
+    "normalize_pixel_coordinates",
+    "normalize_quaternion",
+    "denormalize_pixel_coordinates3d",
+    "normalize_pixel_coordinates3d",
+]
+
+
+class QuaternionCoeffOrder(enum.Enum):
+    XYZW = "xyzw"
+    WXYZ = "wxyz"
+
+
+@conditional_jit_script
+def torch_safe_atan2(y, x, eps: float = 1e-6):
+    y = y.clone()
+    if len(y.shape) == 0:
+        if y.abs() < eps and x.abs() < eps:
+            y += eps
+    else:
+        y[(y.abs() < eps) & (x.abs() < eps)] += eps
+    return torch.atan2(y, x)
+
+
+@conditional_jit_script
+def rad2deg(tensor: torch.Tensor) -> torch.Tensor:
+    r"""Function that converts angles from radians to degrees.
+
+    Args:
+        tensor: Tensor of arbitrary shape.
+
+    Returns:
+        Tensor with same shape as input.
+
+    Example:
+        >>> input = torch.tensor(3.1415926535) * torch.rand(1, 3, 3)
+        >>> output = rad2deg(input)
+    """
+
+    pi = np.pi
+    if not isinstance(tensor, torch.Tensor):
+        raise TypeError("Input type is not a torch.Tensor. Got {}".format(type(tensor)))
+
+    return 180.0 * tensor / pi
+
+
+@conditional_jit_script
+def deg2rad(tensor: torch.Tensor) -> torch.Tensor:
+    r"""Function that converts angles from degrees to radians.
+
+    Args:
+        tensor: Tensor of arbitrary shape.
+
+    Returns:
+        tensor with same shape as input.
+
+    Examples:
+        >>> input = 360. * torch.rand(1, 3, 3)
+        >>> output = deg2rad(input)
+    """
+
+    pi = np.pi
+    if not isinstance(tensor, torch.Tensor):
+        raise TypeError("Input type is not a torch.Tensor. Got {}".format(type(tensor)))
+
+    return tensor * pi / 180.0
+
+
+@conditional_jit_script
+def pol2cart(rho: torch.Tensor, phi: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    r"""Function that converts polar coordinates to cartesian coordinates.
+
+    Args:
+        rho: Tensor of arbitrary shape.
+        phi: Tensor of same arbitrary shape.
+
+    Returns:
+        Tensor with same shape as input.
+
+    Example:
+        >>> rho = torch.rand(1, 3, 3)
+        >>> phi = torch.rand(1, 3, 3)
+        >>> x, y = pol2cart(rho, phi)
+    """
+    if not (isinstance(rho, torch.Tensor) & isinstance(phi, torch.Tensor)):
+        raise TypeError("Input type is not a torch.Tensor. Got {}, {}".format(type(rho), type(phi)))
+
+    x = rho * torch.cos(phi)
+    y = rho * torch.sin(phi)
+    return x, y
+
+
+@conditional_jit_script
+def cart2pol(
+    x: torch.Tensor, y: torch.Tensor, eps: float = 1.0e-8
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Function that converts cartesian coordinates to polar coordinates.
+
+    Args:
+        rho: Tensor of arbitrary shape.
+        phi: Tensor of same arbitrary shape.
+        eps: To avoid division by zero.
+
+    Returns:
+        Tensor with same shape as input.
+
+    Example:
+        >>> x = torch.rand(1, 3, 3)
+        >>> y = torch.rand(1, 3, 3)
+        >>> rho, phi = cart2pol(x, y)
+    """
+    if not (isinstance(x, torch.Tensor) & isinstance(y, torch.Tensor)):
+        raise TypeError("Input type is not a torch.Tensor. Got {}, {}".format(type(x), type(y)))
+
+    rho = torch.sqrt((x**2 + y**2).clamp_min(eps))
+    phi = torch_safe_atan2(y, x)
+    return rho, phi
+
+
+@conditional_jit_script
+def convert_points_from_homogeneous(points: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
+    r"""Function that converts points from homogeneous to Euclidean space.
+
+    Args:
+        points: the points to be transformed.
+        eps: to avoid division by zero.
+
+    Returns:
+        the points in Euclidean space.
+
+    Examples:
+        >>> input = torch.rand(2, 4, 3)  # BxNx3
+        >>> output = convert_points_from_homogeneous(input)  # BxNx2
+    """
+    if not isinstance(points, torch.Tensor):
+        raise TypeError("Input type is not a torch.Tensor. Got {}".format(type(points)))
+
+    if len(points.shape) < 2:
+        raise ValueError("Input must be at least a 2D tensor. Got {}".format(points.shape))
+
+    # we check for points at max_val
+    z_vec: torch.Tensor = points[..., -1:]
+
+    # set the results of division by zeror/near-zero to 1.0
+    # follow the convention of opencv:
+    # https://github.com/opencv/opencv/pull/14411/files
+    mask: torch.Tensor = torch.abs(z_vec) > eps
+    scale = torch.where(mask, 1.0 / (z_vec + eps), torch.ones_like(z_vec))
+
+    return scale * points[..., :-1]
+
+
+@conditional_jit_script
+def convert_points_to_homogeneous(points: torch.Tensor) -> torch.Tensor:
+    r"""Function that converts points from Euclidean to homogeneous space.
+
+    Args:
+        points: the points to be transformed.
+
+    Returns:
+        the points in homogeneous coordinates.
+
+    Examples:
+        >>> input = torch.rand(2, 4, 3)  # BxNx3
+        >>> output = convert_points_to_homogeneous(input)  # BxNx4
+    """
+    if not isinstance(points, torch.Tensor):
+        raise TypeError("Input type is not a torch.Tensor. Got {}".format(type(points)))
+    if len(points.shape) < 2:
+        raise ValueError("Input must be at least a 2D tensor. Got {}".format(points.shape))
+
+    return torch.nn.functional.pad(points, [0, 1], "constant", 1.0)
+
+
+@conditional_jit_script
+def _convert_affinematrix_to_homography_impl(A: torch.Tensor) -> torch.Tensor:
+    H: torch.Tensor = torch.nn.functional.pad(A, [0, 0, 0, 1], "constant", value=0.0)
+    H[..., -1, -1] += 1.0
+    return H
+
+
+@conditional_jit_script
+def convert_affinematrix_to_homography(A: torch.Tensor) -> torch.Tensor:
+    r"""Function that converts batch of affine matrices.
+
+    Args:
+        A: the affine matrix with shape :math:`(B,2,3)`.
+
+    Returns:
+         the homography matrix with shape of :math:`(B,3,3)`.
+
+    Examples:
+        >>> input = torch.rand(2, 2, 3)  # Bx2x3
+        >>> output = convert_affinematrix_to_homography(input)  # Bx3x3
+    """
+    if not isinstance(A, torch.Tensor):
+        raise TypeError("Input type is not a torch.Tensor. Got {}".format(type(A)))
+    if not (len(A.shape) == 3 and A.shape[-2:] == (2, 3)):
+        raise ValueError("Input matrix must be a Bx2x3 tensor. Got {}".format(A.shape))
+    return _convert_affinematrix_to_homography_impl(A)
+
+
+@conditional_jit_script
+def convert_affinematrix_to_homography3d(A: torch.Tensor) -> torch.Tensor:
+    r"""Function that converts batch of 3d affine matrices.
+
+    Args:
+        A: the affine matrix with shape :math:`(B,3,4)`.
+
+    Returns:
+         the homography matrix with shape of :math:`(B,4,4)`.
+
+    Examples:
+        >>> input = torch.rand(2, 3, 4)  # Bx3x4
+        >>> output = convert_affinematrix_to_homography3d(input)  # Bx4x4
+    """
+    if not isinstance(A, torch.Tensor):
+        raise TypeError("Input type is not a torch.Tensor. Got {}".format(type(A)))
+    if not (len(A.shape) == 3 and A.shape[-2:] == (3, 4)):
+        raise ValueError("Input matrix must be a Bx3x4 tensor. Got {}".format(A.shape))
+    return _convert_affinematrix_to_homography_impl(A)
+
+
+@conditional_jit_script
+def _compute_rotation_matrix(angle_axis, theta2, eps: float = 1e-6):
+    # We want to be careful to only evaluate the square root if the
+    # norm of the angle_axis vector is greater than zero. Otherwise
+    # we get a division by zero.
+    k_one = 1.0
+    theta = torch.sqrt(theta2.clamp_min(eps))
+    wxyz = angle_axis / (theta + eps)
+    wx, wy, wz = torch.chunk(wxyz, 3, dim=1)
+    cos_theta = torch.cos(theta)
+    sin_theta = torch.sin(theta)
+
+    r00 = cos_theta + wx * wx * (k_one - cos_theta)
+    r10 = wz * sin_theta + wx * wy * (k_one - cos_theta)
+    r20 = -wy * sin_theta + wx * wz * (k_one - cos_theta)
+    r01 = wx * wy * (k_one - cos_theta) - wz * sin_theta
+    r11 = cos_theta + wy * wy * (k_one - cos_theta)
+    r21 = wx * sin_theta + wy * wz * (k_one - cos_theta)
+    r02 = wy * sin_theta + wx * wz * (k_one - cos_theta)
+    r12 = -wx * sin_theta + wy * wz * (k_one - cos_theta)
+    r22 = cos_theta + wz * wz * (k_one - cos_theta)
+    rotation_matrix = torch.cat([r00, r01, r02, r10, r11, r12, r20, r21, r22], dim=1)
+    return rotation_matrix.view(-1, 3, 3)
+
+
+@conditional_jit_script
+def _compute_rotation_matrix_taylor(angle_axis):
+    rx, ry, rz = torch.chunk(angle_axis, 3, dim=1)
+    k_one = torch.ones_like(rx)
+    rotation_matrix = torch.cat([k_one, -rz, ry, rz, k_one, -rx, -ry, rx, k_one], dim=1)
+    return rotation_matrix.view(-1, 3, 3)
+
+
+@conditional_jit_script
+def angle_axis_to_rotation_matrix(angle_axis: torch.Tensor) -> torch.Tensor:
+    r"""Convert 3d vector of axis-angle rotation to 3x3 rotation matrix.
+
+    Args:
+        angle_axis: tensor of 3d vector of axis-angle rotations.
+
+    Returns:
+        tensor of 3x3 rotation matrices.
+
+    Shape:
+        - Input: :math:`(N, 3)`
+        - Output: :math:`(N, 3, 3)`
+
+    Example:
+        >>> input = torch.rand(1, 3)  # Nx3
+        >>> output = angle_axis_to_rotation_matrix(input)  # Nx3x3
+    """
+    if not isinstance(angle_axis, torch.Tensor):
+        raise TypeError("Input type is not a torch.Tensor. Got {}".format(type(angle_axis)))
+
+    if not angle_axis.shape[-1] == 3:
+        raise ValueError("Input size must be a (*, 3) tensor. Got {}".format(angle_axis.shape))
+
+    orig_shape = angle_axis.shape
+    angle_axis = angle_axis.reshape(-1, 3)
+
+    # stolen from ceres/rotation.h
+
+    _angle_axis = torch.unsqueeze(angle_axis, dim=1)
+    theta2 = torch.matmul(_angle_axis, _angle_axis.transpose(1, 2))
+    theta2 = torch.squeeze(theta2, dim=1)
+
+    # compute rotation matrices
+    rotation_matrix_normal = _compute_rotation_matrix(angle_axis, theta2)
+    rotation_matrix_taylor = _compute_rotation_matrix_taylor(angle_axis)
+
+    # create mask to handle both cases
+    eps = 1e-6
+    mask = (theta2 > eps).view(-1, 1, 1).to(theta2.device)
+    mask_pos = (mask).type_as(theta2)
+    mask_neg = (mask == torch.tensor(False)).type_as(theta2)  # noqa
+
+    # create output pose matrix
+    batch_size = angle_axis.shape[0]
+    rotation_matrix = torch.eye(3).to(angle_axis.device).type_as(angle_axis)
+    rotation_matrix = rotation_matrix.view(1, 3, 3).repeat(batch_size, 1, 1)
+    # fill output matrix with masked values
+    rotation_matrix[..., :3, :3] = (
+        mask_pos * rotation_matrix_normal + mask_neg * rotation_matrix_taylor
+    )
+
+    rotation_matrix = rotation_matrix.view(orig_shape[:-1] + (3, 3))
+    return rotation_matrix  # Nx3x3
+
+
+# @conditional_jit_script
+def safe_zero_division(
+    numerator: torch.Tensor, denominator: torch.Tensor, eps: float = 1.0e-6
+) -> torch.Tensor:
+    denominator = denominator.clone()
+    if len(denominator.shape) == 0:
+        if denominator.abs() < eps:
+            denominator += eps
+    else:
+        # denominator[denominator.abs() < eps] += eps
+        denominator = torch.where(denominator.abs() < eps, denominator + eps, denominator)
+    return numerator / denominator
+
+
+# @conditional_jit_script
+def rotation_matrix_to_quaternion(
+    rotation_matrix: torch.Tensor,
+    eps: float = 1.0e-6,
+    order: QuaternionCoeffOrder = QuaternionCoeffOrder.WXYZ,
+) -> torch.Tensor:
+    r"""Convert 3x3 rotation matrix to 4d quaternion vector.
+
+    The quaternion vector has components in (w, x, y, z) or (x, y, z, w) format.
+
+    .. note::
+        The (x, y, z, w) order is going to be deprecated in favor of efficiency.
+
+    Args:
+        rotation_matrix: the rotation matrix to convert.
+        eps: small value to avoid zero division.
+        order: quaternion coefficient order. Note: 'xyzw' will be deprecated in favor of 'wxyz'.
+
+    Return:
+        the rotation in quaternion.
+
+    Shape:
+        - Input: :math:`(*, 3, 3)`
+        - Output: :math:`(*, 4)`
+
+    Example:
+        >>> input = torch.rand(4, 3, 3)  # Nx3x3
+        >>> output = rotation_matrix_to_quaternion(input, eps=torch.finfo(input.dtype).eps,
+        ...                                        order=QuaternionCoeffOrder.WXYZ)  # Nx4
+    """
+    if not isinstance(rotation_matrix, torch.Tensor):
+        raise TypeError(f"Input type is not a torch.Tensor. Got {type(rotation_matrix)}")
+
+    if not rotation_matrix.shape[-2:] == (3, 3):
+        raise ValueError(f"Input size must be a (*, 3, 3) tensor. Got {rotation_matrix.shape}")
+
+    # if not torch.jit.is_scripting():
+    #    if order.name not in QuaternionCoeffOrder.__members__.keys():
+    #        raise ValueError(
+    #            f"order must be one of {QuaternionCoeffOrder.__members__.keys()}"
+    #        )
+
+    if order == QuaternionCoeffOrder.XYZW:
+        warnings.warn(
+            "`XYZW` quaternion coefficient order is deprecated and"
+            " will be removed after > 0.6. "
+            "Please use `QuaternionCoeffOrder.WXYZ` instead."
+        )
+
+    m00, m01, m02 = (
+        rotation_matrix[..., 0, 0],
+        rotation_matrix[..., 0, 1],
+        rotation_matrix[..., 0, 2],
+    )
+    m10, m11, m12 = (
+        rotation_matrix[..., 1, 0],
+        rotation_matrix[..., 1, 1],
+        rotation_matrix[..., 1, 2],
+    )
+    m20, m21, m22 = (
+        rotation_matrix[..., 2, 0],
+        rotation_matrix[..., 2, 1],
+        rotation_matrix[..., 2, 2],
+    )
+
+    trace: torch.Tensor = m00 + m11 + m22
+
+    sq = torch.sqrt((trace + 1.0).clamp_min(eps)) * 2.0  # sq = 4 * qw.
+    qw = 0.25 * sq
+    qx = safe_zero_division(m21 - m12, sq)
+    qy = safe_zero_division(m02 - m20, sq)
+    qz = safe_zero_division(m10 - m01, sq)
+    if order == QuaternionCoeffOrder.XYZW:
+        trace_positive_cond = torch.stack((qx, qy, qz, qw), dim=-1)
+    trace_positive_cond = torch.stack((qw, qx, qy, qz), dim=-1)
+
+    sq = torch.sqrt((1.0 + m00 - m11 - m22).clamp_min(eps)) * 2.0  # sq = 4 * qx.
+    qw = safe_zero_division(m21 - m12, sq)
+    qx = 0.25 * sq
+    qy = safe_zero_division(m01 + m10, sq)
+    qz = safe_zero_division(m02 + m20, sq)
+    if order == QuaternionCoeffOrder.XYZW:
+        cond_1 = torch.stack((qx, qy, qz, qw), dim=-1)
+    cond_1 = torch.stack((qw, qx, qy, qz), dim=-1)
+
+    sq = torch.sqrt((1.0 + m11 - m00 - m22).clamp_min(eps)) * 2.0  # sq = 4 * qy.
+    qw = safe_zero_division(m02 - m20, sq)
+    qx = safe_zero_division(m01 + m10, sq)
+    qy = 0.25 * sq
+    qz = safe_zero_division(m12 + m21, sq)
+    if order == QuaternionCoeffOrder.XYZW:
+        cond_2 = torch.stack((qx, qy, qz, qw), dim=-1)
+    cond_2 = torch.stack((qw, qx, qy, qz), dim=-1)
+
+    sq = torch.sqrt((1.0 + m22 - m00 - m11).clamp_min(eps)) * 2.0  # sq = 4 * qz.
+    qw = safe_zero_division(m10 - m01, sq)
+    qx = safe_zero_division(m02 + m20, sq)
+    qy = safe_zero_division(m12 + m21, sq)
+    qz = 0.25 * sq
+    if order == QuaternionCoeffOrder.XYZW:
+        cond_3 = torch.stack((qx, qy, qz, qw), dim=-1)
+    cond_3 = torch.stack((qw, qx, qy, qz), dim=-1)
+
+    where_2 = torch.where((m11 > m22).unsqueeze(-1), cond_2, cond_3)
+    where_1 = torch.where(((m00 > m11) & (m00 > m22)).unsqueeze(-1), cond_1, where_2)
+
+    quaternion: torch.Tensor = torch.where(
+        (trace > 0.0).unsqueeze(-1), trace_positive_cond, where_1
+    )
+    return quaternion
+
+
+# @conditional_jit_script
+def normalize_quaternion(quaternion: torch.Tensor, eps: float = 1.0e-12) -> torch.Tensor:
+    r"""Normalizes a quaternion.
+
+    The quaternion should be in (x, y, z, w) format.
+
+    Args:
+        quaternion: a tensor containing a quaternion to be normalized.
+          The tensor can be of shape :math:`(*, 4)`.
+        eps: small value to avoid division by zero.
+
+    Return:
+        the normalized quaternion of shape :math:`(*, 4)`.
+
+    Example:
+        >>> quaternion = torch.tensor((1., 0., 1., 0.))
+        >>> normalize_quaternion(quaternion)
+        tensor([0.7071, 0.0000, 0.7071, 0.0000])
+    """
+    if not isinstance(quaternion, torch.Tensor):
+        raise TypeError("Input type is not a torch.Tensor. Got {}".format(type(quaternion)))
+
+    if not quaternion.shape[-1] == 4:
+        raise ValueError("Input must be a tensor of shape (*, 4). Got {}".format(quaternion.shape))
+    return F.normalize(quaternion, p=2.0, dim=-1, eps=eps)
+
+
+# based on:
+# https://github.com/matthew-brett/transforms3d/blob/8965c48401d9e8e66b6a8c37c65f2fc200a076fa/transforms3d/quaternions.py#L101
+# https://github.com/tensorflow/graphics/blob/master/tensorflow_graphics/geometry/transformation/rotation_matrix_3d.py#L247
+
+
+# @conditional_jit_script
+def quaternion_to_rotation_matrix(
+    quaternion: torch.Tensor, order: QuaternionCoeffOrder = QuaternionCoeffOrder.WXYZ
+) -> torch.Tensor:
+    r"""Converts a quaternion to a rotation matrix.
+
+    The quaternion should be in (x, y, z, w) or (w, x, y, z) format.
+
+    Args:
+        quaternion: a tensor containing a quaternion to be converted.
+          The tensor can be of shape :math:`(*, 4)`.
+        order: quaternion coefficient order. Note: 'xyzw' will be deprecated in favor of 'wxyz'.
+
+    Return:
+        the rotation matrix of shape :math:`(*, 3, 3)`.
+
+    Example:
+        >>> quaternion = torch.tensor((0., 0., 0., 1.))
+        >>> quaternion_to_rotation_matrix(quaternion, order=QuaternionCoeffOrder.WXYZ)
+        tensor([[-1.,  0.,  0.],
+                [ 0., -1.,  0.],
+                [ 0.,  0.,  1.]])
+    """
+    if not isinstance(quaternion, torch.Tensor):
+        raise TypeError(f"Input type is not a torch.Tensor. Got {type(quaternion)}")
+
+    if not quaternion.shape[-1] == 4:
+        raise ValueError(f"Input must be a tensor of shape (*, 4). Got {quaternion.shape}")
+
+    # if not torch.jit.is_scripting():
+    #    if order.name not in QuaternionCoeffOrder.__members__.keys():
+    #        raise ValueError(
+    #            f"order must be one of {QuaternionCoeffOrder.__members__.keys()}"
+    #        )
+
+    if order == QuaternionCoeffOrder.XYZW:
+        warnings.warn(
+            "`XYZW` quaternion coefficient order is deprecated and"
+            " will be removed after > 0.6. "
+            "Please use `QuaternionCoeffOrder.WXYZ` instead."
+        )
+
+    # normalize the input quaternion
+    quaternion_norm: torch.Tensor = normalize_quaternion(quaternion)
+
+    # unpack the normalized quaternion components
+    if order == QuaternionCoeffOrder.XYZW:
+        x, y, z, w = (
+            quaternion_norm[..., 0],
+            quaternion_norm[..., 1],
+            quaternion_norm[..., 2],
+            quaternion_norm[..., 3],
+        )
+    else:
+        w, x, y, z = (
+            quaternion_norm[..., 0],
+            quaternion_norm[..., 1],
+            quaternion_norm[..., 2],
+            quaternion_norm[..., 3],
+        )
+
+    # compute the actual conversion
+    tx: torch.Tensor = 2.0 * x
+    ty: torch.Tensor = 2.0 * y
+    tz: torch.Tensor = 2.0 * z
+    twx: torch.Tensor = tx * w
+    twy: torch.Tensor = ty * w
+    twz: torch.Tensor = tz * w
+    txx: torch.Tensor = tx * x
+    txy: torch.Tensor = ty * x
+    txz: torch.Tensor = tz * x
+    tyy: torch.Tensor = ty * y
+    tyz: torch.Tensor = tz * y
+    tzz: torch.Tensor = tz * z
+    one: torch.Tensor = torch.tensor(1.0)
+
+    matrix: torch.Tensor = torch.stack(
+        (
+            one - (tyy + tzz),
+            txy - twz,
+            txz + twy,
+            txy + twz,
+            one - (txx + tzz),
+            tyz - twx,
+            txz - twy,
+            tyz + twx,
+            one - (txx + tyy),
+        ),
+        dim=-1,
+    ).view(quaternion.shape[:-1] + (3, 3))
+
+    # if len(quaternion.shape) == 1:
+    #     matrix = torch.squeeze(matrix, dim=0)
+    return matrix
+
+
+@conditional_jit_script
+def quaternion_to_angle_axis(
+    quaternion: torch.Tensor,
+    eps: float = 1.0e-6,
+    order: QuaternionCoeffOrder = QuaternionCoeffOrder.WXYZ,
+) -> torch.Tensor:
+    """Convert quaternion vector to angle axis of rotation.
+
+    The quaternion should be in (x, y, z, w) or (w, x, y, z) format.
+
+    Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h
+
+    Args:
+        quaternion: tensor with quaternions.
+        order: quaternion coefficient order. Note: 'xyzw' will be deprecated in favor of 'wxyz'.
+
+    Return:
+        tensor with angle axis of rotation.
+
+    Shape:
+        - Input: :math:`(*, 4)` where `*` means, any number of dimensions
+        - Output: :math:`(*, 3)`
+
+    Example:
+        >>> quaternion = torch.rand(2, 4)  # Nx4
+        >>> angle_axis = quaternion_to_angle_axis(quaternion)  # Nx3
+    """
+
+    if not quaternion.shape[-1] == 4:
+        raise ValueError(f"Input must be a tensor of shape Nx4 or 4. Got {quaternion.shape}")
+
+    if not torch.jit.is_scripting():
+        if order.name not in QuaternionCoeffOrder.__members__.keys():
+            raise ValueError(f"order must be one of {QuaternionCoeffOrder.__members__.keys()}")
+
+    if order == QuaternionCoeffOrder.XYZW:
+        warnings.warn(
+            "`XYZW` quaternion coefficient order is deprecated and"
+            " will be removed after > 0.6. "
+            "Please use `QuaternionCoeffOrder.WXYZ` instead."
+        )
+    # unpack input and compute conversion
+    q1: torch.Tensor = torch.tensor([])
+    q2: torch.Tensor = torch.tensor([])
+    q3: torch.Tensor = torch.tensor([])
+    cos_theta: torch.Tensor = torch.tensor([])
+
+    if order == QuaternionCoeffOrder.XYZW:
+        q1 = quaternion[..., 0]
+        q2 = quaternion[..., 1]
+        q3 = quaternion[..., 2]
+        cos_theta = quaternion[..., 3]
+    else:
+        cos_theta = quaternion[..., 0]
+        q1 = quaternion[..., 1]
+        q2 = quaternion[..., 2]
+        q3 = quaternion[..., 3]
+
+    sin_squared_theta: torch.Tensor = q1 * q1 + q2 * q2 + q3 * q3
+
+    sin_theta: torch.Tensor = torch.sqrt((sin_squared_theta).clamp_min(eps))
+    two_theta: torch.Tensor = 2.0 * torch.where(
+        cos_theta < 0.0,
+        torch_safe_atan2(-sin_theta, -cos_theta),
+        torch_safe_atan2(sin_theta, cos_theta),
+    )
+
+    k_pos: torch.Tensor = safe_zero_division(two_theta, sin_theta, eps)
+    k_neg: torch.Tensor = 2.0 * torch.ones_like(sin_theta)
+    k: torch.Tensor = torch.where(sin_squared_theta > 0.0, k_pos, k_neg)
+
+    angle_axis: torch.Tensor = torch.zeros_like(quaternion)[..., :3]
+    angle_axis[..., 0] += q1 * k
+    angle_axis[..., 1] += q2 * k
+    angle_axis[..., 2] += q3 * k
+    return angle_axis
+
+
+@conditional_jit_script
+def rotation_matrix_to_angle_axis(rotation_matrix: torch.Tensor) -> torch.Tensor:
+    r"""Convert 3x3 rotation matrix to Rodrigues vector.
+
+    Args:
+        rotation_matrix: rotation matrix.
+
+    Returns:
+        Rodrigues vector transformation.
+
+    Shape:
+        - Input: :math:`(N, 3, 3)`
+        - Output: :math:`(N, 3)`
+
+    Example:
+        >>> input = torch.rand(2, 3, 3)  # Nx3x3
+        >>> output = rotation_matrix_to_angle_axis(input)  # Nx3
+    """
+    if not isinstance(rotation_matrix, torch.Tensor):
+        raise TypeError(f"Input type is not a torch.Tensor. Got {type(rotation_matrix)}")
+
+    if not rotation_matrix.shape[-2:] == (3, 3):
+        raise ValueError(f"Input size must be a (*, 3, 3) tensor. Got {rotation_matrix.shape}")
+    quaternion: torch.Tensor = rotation_matrix_to_quaternion(
+        rotation_matrix, order=QuaternionCoeffOrder.WXYZ
+    )
+    return quaternion_to_angle_axis(quaternion, order=QuaternionCoeffOrder.WXYZ)
+
+
+@conditional_jit_script
+def quaternion_log_to_exp(
+    quaternion: torch.Tensor,
+    eps: float = 1.0e-6,
+    order: QuaternionCoeffOrder = QuaternionCoeffOrder.WXYZ,
+) -> torch.Tensor:
+    r"""Applies exponential map to log quaternion.
+
+    The quaternion should be in (x, y, z, w) or (w, x, y, z) format.
+
+    Args:
+        quaternion: a tensor containing a quaternion to be converted.
+          The tensor can be of shape :math:`(*, 3)`.
+        order: quaternion coefficient order. Note: 'xyzw' will be deprecated in favor of 'wxyz'.
+
+    Return:
+        the quaternion exponential map of shape :math:`(*, 4)`.
+
+    Example:
+        >>> quaternion = torch.tensor((0., 0., 0.))
+        >>> quaternion_log_to_exp(quaternion, eps=torch.finfo(quaternion.dtype).eps,
+        ...                       order=QuaternionCoeffOrder.WXYZ)
+        tensor([1., 0., 0., 0.])
+    """
+    if not isinstance(quaternion, torch.Tensor):
+        raise TypeError(f"Input type is not a torch.Tensor. Got {type(quaternion)}")
+
+    if not quaternion.shape[-1] == 3:
+        raise ValueError(f"Input must be a tensor of shape (*, 3). Got {quaternion.shape}")
+
+    if not torch.jit.is_scripting():
+        if order.name not in QuaternionCoeffOrder.__members__.keys():
+            raise ValueError(f"order must be one of {QuaternionCoeffOrder.__members__.keys()}")
+
+    if order == QuaternionCoeffOrder.XYZW:
+        warnings.warn(
+            "`XYZW` quaternion coefficient order is deprecated and"
+            " will be removed after > 0.6. "
+            "Please use `QuaternionCoeffOrder.WXYZ` instead."
+        )
+
+    # compute quaternion norm
+    norm_q: torch.Tensor = torch.norm(quaternion, p=2, dim=-1, keepdim=True).clamp(min=eps)
+
+    # compute scalar and vector
+    quaternion_vector: torch.Tensor = quaternion * torch.sin(norm_q) / norm_q
+    quaternion_scalar: torch.Tensor = torch.cos(norm_q)
+
+    # compose quaternion and return
+    quaternion_exp: torch.Tensor = torch.tensor([])
+    if order == QuaternionCoeffOrder.XYZW:
+        quaternion_exp = torch.cat((quaternion_vector, quaternion_scalar), dim=-1)
+    else:
+        quaternion_exp = torch.cat((quaternion_scalar, quaternion_vector), dim=-1)
+
+    return quaternion_exp
+
+
+@conditional_jit_script
+def quaternion_exp_to_log(
+    quaternion: torch.Tensor,
+    eps: float = 1.0e-6,
+    order: QuaternionCoeffOrder = QuaternionCoeffOrder.WXYZ,
+) -> torch.Tensor:
+    r"""Applies the log map to a quaternion.
+
+    The quaternion should be in (x, y, z, w) format.
+
+    Args:
+        quaternion: a tensor containing a quaternion to be converted.
+          The tensor can be of shape :math:`(*, 4)`.
+        eps: A small number for clamping.
+        order: quaternion coefficient order. Note: 'xyzw' will be deprecated in favor of 'wxyz'.
+
+    Return:
+        the quaternion log map of shape :math:`(*, 3)`.
+
+    Example:
+        >>> quaternion = torch.tensor((1., 0., 0., 0.))
+        >>> quaternion_exp_to_log(quaternion, eps=torch.finfo(quaternion.dtype).eps,
+        ...                       order=QuaternionCoeffOrder.WXYZ)
+        tensor([0., 0., 0.])
+    """
+    if not isinstance(quaternion, torch.Tensor):
+        raise TypeError(f"Input type is not a torch.Tensor. Got {type(quaternion)}")
+
+    if not quaternion.shape[-1] == 4:
+        raise ValueError(f"Input must be a tensor of shape (*, 4). Got {quaternion.shape}")
+
+    if not torch.jit.is_scripting():
+        if order.name not in QuaternionCoeffOrder.__members__.keys():
+            raise ValueError(f"order must be one of {QuaternionCoeffOrder.__members__.keys()}")
+
+    if order == QuaternionCoeffOrder.XYZW:
+        warnings.warn(
+            "`XYZW` quaternion coefficient order is deprecated and"
+            " will be removed after > 0.6. "
+            "Please use `QuaternionCoeffOrder.WXYZ` instead."
+        )
+
+    # unpack quaternion vector and scalar
+    quaternion_vector: torch.Tensor = torch.tensor([])
+    quaternion_scalar: torch.Tensor = torch.tensor([])
+
+    if order == QuaternionCoeffOrder.XYZW:
+        quaternion_vector = quaternion[..., 0:3]
+        quaternion_scalar = quaternion[..., 3:4]
+    else:
+        quaternion_scalar = quaternion[..., 0:1]
+        quaternion_vector = quaternion[..., 1:4]
+
+    # compute quaternion norm
+    norm_q: torch.Tensor = torch.norm(quaternion_vector, p=2, dim=-1, keepdim=True).clamp(min=eps)
+
+    # apply log map
+    quaternion_log: torch.Tensor = (
+        quaternion_vector
+        * torch.acos(torch.clamp(quaternion_scalar, min=-1.0 + eps, max=1.0 - eps))
+        / norm_q
+    )
+
+    return quaternion_log
+
+
+# based on:
+# https://github.com/facebookresearch/QuaterNet/blob/master/common/quaternion.py#L138
+
+
+@conditional_jit_script
+def angle_axis_to_quaternion(
+    angle_axis: torch.Tensor,
+    eps: float = 1.0e-6,
+    order: QuaternionCoeffOrder = QuaternionCoeffOrder.WXYZ,
+) -> torch.Tensor:
+    r"""Convert an angle axis to a quaternion.
+
+    The quaternion vector has components in (x, y, z, w) or (w, x, y, z) format.
+
+    Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h
+
+    Args:
+        angle_axis: tensor with angle axis.
+        order: quaternion coefficient order. Note: 'xyzw' will be deprecated in favor of 'wxyz'.
+
+    Return:
+        tensor with quaternion.
+
+    Shape:
+        - Input: :math:`(*, 3)` where `*` means, any number of dimensions
+        - Output: :math:`(*, 4)`
+
+    Example:
+        >>> angle_axis = torch.rand(2, 3)  # Nx3
+        >>> quaternion = angle_axis_to_quaternion(angle_axis, order=QuaternionCoeffOrder.WXYZ)  # Nx4
+    """
+
+    if not angle_axis.shape[-1] == 3:
+        raise ValueError(f"Input must be a tensor of shape Nx3 or 3. Got {angle_axis.shape}")
+
+    if not torch.jit.is_scripting():
+        if order.name not in QuaternionCoeffOrder.__members__.keys():
+            raise ValueError(f"order must be one of {QuaternionCoeffOrder.__members__.keys()}")
+
+    if order == QuaternionCoeffOrder.XYZW:
+        warnings.warn(
+            "`XYZW` quaternion coefficient order is deprecated and"
+            " will be removed after > 0.6. "
+            "Please use `QuaternionCoeffOrder.WXYZ` instead."
+        )
+
+    # unpack input and compute conversion
+    a0: torch.Tensor = angle_axis[..., 0:1]
+    a1: torch.Tensor = angle_axis[..., 1:2]
+    a2: torch.Tensor = angle_axis[..., 2:3]
+    theta_squared: torch.Tensor = a0 * a0 + a1 * a1 + a2 * a2
+
+    theta: torch.Tensor = torch.sqrt((theta_squared).clamp_min(eps))
+    half_theta: torch.Tensor = theta * 0.5
+
+    mask: torch.Tensor = theta_squared > 0.0
+    ones: torch.Tensor = torch.ones_like(half_theta)
+
+    k_neg: torch.Tensor = 0.5 * ones
+    k_pos: torch.Tensor = safe_zero_division(torch.sin(half_theta), theta, eps)
+    k: torch.Tensor = torch.where(mask, k_pos, k_neg)
+    w: torch.Tensor = torch.where(mask, torch.cos(half_theta), ones)
+
+    quaternion: torch.Tensor = torch.zeros(
+        size=angle_axis.shape[:-1] + (4,),
+        dtype=angle_axis.dtype,
+        device=angle_axis.device,
+    )
+    if order == QuaternionCoeffOrder.XYZW:
+        quaternion[..., 0:1] = a0 * k
+        quaternion[..., 1:2] = a1 * k
+        quaternion[..., 2:3] = a2 * k
+        quaternion[..., 3:4] = w
+    else:
+        quaternion[..., 1:2] = a0 * k
+        quaternion[..., 2:3] = a1 * k
+        quaternion[..., 3:4] = a2 * k
+        quaternion[..., 0:1] = w
+    return quaternion
+
+
+# based on:
+# https://github.com/ClementPinard/SfmLearner-Pytorch/blob/master/inverse_warp.py#L65-L71
+
+
+@conditional_jit_script
+def normalize_pixel_coordinates(
+    pixel_coordinates: torch.Tensor, height: int, width: int, eps: float = 1e-8
+) -> torch.Tensor:
+    r"""Normalize pixel coordinates between -1 and 1.
+
+    Normalized, -1 if on extreme left, 1 if on extreme right (x = w-1).
+
+    Args:
+        pixel_coordinates: the grid with pixel coordinates. Shape can be :math:`(*, 2)`.
+        width: the maximum width in the x-axis.
+        height: the maximum height in the y-axis.
+        eps: safe division by zero.
+
+    Return:
+        the normalized pixel coordinates.
+    """
+    if pixel_coordinates.shape[-1] != 2:
+        raise ValueError(
+            "Input pixel_coordinates must be of shape (*, 2). Got {}".format(
+                pixel_coordinates.shape
+            )
+        )
+    # compute normalization factor
+    hw: torch.Tensor = torch.stack(
+        [
+            torch.tensor(width, device=pixel_coordinates.device, dtype=pixel_coordinates.dtype),
+            torch.tensor(height, device=pixel_coordinates.device, dtype=pixel_coordinates.dtype),
+        ]
+    )
+
+    factor: torch.Tensor = torch.tensor(
+        2.0, device=pixel_coordinates.device, dtype=pixel_coordinates.dtype
+    ) / (hw - 1).clamp(eps)
+
+    return factor * pixel_coordinates - 1
+
+
+@conditional_jit_script
+def denormalize_pixel_coordinates(
+    pixel_coordinates: torch.Tensor, height: int, width: int, eps: float = 1e-8
+) -> torch.Tensor:
+    r"""Denormalize pixel coordinates.
+
+    The input is assumed to be -1 if on extreme left, 1 if on extreme right (x = w-1).
+
+    Args:
+        pixel_coordinates: the normalized grid coordinates. Shape can be :math:`(*, 2)`.
+        width: the maximum width in the x-axis.
+        height: the maximum height in the y-axis.
+        eps: safe division by zero.
+
+    Return:
+        the denormalized pixel coordinates.
+    """
+    if pixel_coordinates.shape[-1] != 2:
+        raise ValueError(
+            "Input pixel_coordinates must be of shape (*, 2). Got {}".format(
+                pixel_coordinates.shape
+            )
+        )
+    # compute normalization factor
+    hw: torch.Tensor = (
+        torch.stack([torch.tensor(width), torch.tensor(height)])
+        .to(pixel_coordinates.device)
+        .to(pixel_coordinates.dtype)
+    )
+
+    factor: torch.Tensor = torch.tensor(2.0) / (hw - 1).clamp(eps)
+
+    return torch.tensor(1.0) / factor * (pixel_coordinates + 1)
+
+
+@conditional_jit_script
+def normalize_pixel_coordinates3d(
+    pixel_coordinates: torch.Tensor,
+    depth: int,
+    height: int,
+    width: int,
+    eps: float = 1e-8,
+) -> torch.Tensor:
+    r"""Normalize pixel coordinates between -1 and 1.
+
+    Normalized, -1 if on extreme left, 1 if on extreme right (x = w-1).
+
+    Args:
+        pixel_coordinates: the grid with pixel coordinates. Shape can be :math:`(*, 3)`.
+        depth: the maximum depth in the z-axis.
+        height: the maximum height in the y-axis.
+        width: the maximum width in the x-axis.
+        eps: safe division by zero.
+
+    Return:
+        the normalized pixel coordinates.
+    """
+    if pixel_coordinates.shape[-1] != 3:
+        raise ValueError(
+            "Input pixel_coordinates must be of shape (*, 3). Got {}".format(
+                pixel_coordinates.shape
+            )
+        )
+    # compute normalization factor
+    dhw: torch.Tensor = (
+        torch.stack([torch.tensor(depth), torch.tensor(width), torch.tensor(height)])
+        .to(pixel_coordinates.device)
+        .to(pixel_coordinates.dtype)
+    )
+
+    factor: torch.Tensor = torch.tensor(2.0) / (dhw - 1).clamp(eps)
+
+    return factor * pixel_coordinates - 1
+
+
+@conditional_jit_script
+def denormalize_pixel_coordinates3d(
+    pixel_coordinates: torch.Tensor,
+    depth: int,
+    height: int,
+    width: int,
+    eps: float = 1e-8,
+) -> torch.Tensor:
+    r"""Denormalize pixel coordinates.
+
+    The input is assumed to be -1 if on extreme left, 1 if on extreme right (x = w-1).
+
+    Args:
+        pixel_coordinates: the normalized grid coordinates. Shape can be :math:`(*, 3)`.
+        depth: the maximum depth in the x-axis.
+        height: the maximum height in the y-axis.
+        width: the maximum width in the x-axis.
+        eps: safe division by zero.
+
+    Return:
+        the denormalized pixel coordinates.
+    """
+    if pixel_coordinates.shape[-1] != 3:
+        raise ValueError(
+            "Input pixel_coordinates must be of shape (*, 3). Got {}".format(
+                pixel_coordinates.shape
+            )
+        )
+    # compute normalization factor
+    dhw: torch.Tensor = (
+        torch.stack([torch.tensor(depth), torch.tensor(width), torch.tensor(height)])
+        .to(pixel_coordinates.device)
+        .to(pixel_coordinates.dtype)
+    )
+
+    factor: torch.Tensor = torch.tensor(2.0) / (dhw - 1).clamp(eps)
+
+    return torch.tensor(1.0) / factor * (pixel_coordinates + 1)
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/math.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/math.py
new file mode 100644
index 0000000000000000000000000000000000000000..815508f2a53f4ef300fe55faf7163887b5ddceb5
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/math.py
@@ -0,0 +1,180 @@
+"""Interpolation and frame-rate rescaling for pose sequences.
+
+Provides linear interpolation (via scipy), quaternion slerp, and
+functions to up/down-sample joint pose trajectories to a target frame rate.
+"""
+
+import torch
+import numpy as np
+from scipy.interpolate import interp1d
+from .kornia_transform import angle_axis_to_quaternion, quaternion_to_angle_axis
+
+
+def interp_tensor_with_scipy(x, new_len=None, scale=None, dim=-1):
+    orig_len = x.shape[dim]
+    if new_len is None:
+        new_len = int(orig_len * scale)
+    T = orig_len
+    f = interp1d(
+        np.linspace(0, T, orig_len),
+        x.cpu().numpy(),
+        axis=dim,
+        assume_sorted=True,
+        fill_value="extrapolate",
+    )
+    x_interp = torch.from_numpy(f(np.linspace(0, T, new_len))).type_as(x)
+    return x_interp
+
+
+def slerp(q0, q1, t):
+    # type: (torch.Tensor, torch.Tensor, torch.Tensor) -> torch.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[:, None]) * half_theta) / sin_half_theta
+    ratioB = torch.sin(t[:, None] * 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
+
+
+def _slerp_batch(a: torch.Tensor, b: torch.Tensor, blend: torch.Tensor) -> torch.Tensor:
+    """Spherical linear interpolation between two quaternions."""
+    slerped_quats = torch.zeros_like(a)
+    slerped_quats = slerp(a, b, blend)
+    return slerped_quats
+
+
+def interpolate_quaternions(
+    pose_quat: torch.Tensor, source_fps: float, target_fps: float
+) -> torch.Tensor:
+    """
+    Interpolate quaternions from source_fps to target_fps.
+
+    Args:
+        pose_quat: Input quaternions, shape (1, T, 4)
+        source_fps: Source frame rate
+        target_fps: Target frame rate
+
+    Returns:
+        Interpolated quaternions
+    """
+    device = pose_quat.device
+    in_shape = pose_quat.shape
+    assert in_shape[0] == 1, "Only support single sequence for now"
+
+    T = in_shape[1]
+    duration = (T - 1) * (1 / source_fps)
+    times = torch.arange(0, duration + 1e-6, 1 / target_fps, dtype=torch.float32, device=device)
+    times = times[times <= duration]
+
+    # Compute frame indices and blend factors
+    frame_indices = times * source_fps
+    index_0 = torch.floor(frame_indices).long()
+    index_1 = torch.min(index_0 + 1, torch.tensor(T - 1, device=device))
+    blend = frame_indices - index_0.float()
+
+    pose_quat_interp = _slerp_batch(pose_quat[0, index_0], pose_quat[0, index_1], blend)
+    pose_quat_interp = pose_quat_interp.unsqueeze(0)
+
+    return pose_quat_interp
+
+
+def interpolate_pose(
+    pose_aa: torch.Tensor,
+    source_fps: float,
+    target_fps: float,
+    device: str = "cpu",
+    interpolation_type: str = "slerp",
+    rot_type: str = "aa",
+) -> torch.Tensor:
+    """
+    Interpolate pose_aa from source_fps to target_fps using specified interpolation method.
+
+    Args:
+        pose_aa: Input pose in angle-axis format, shape (T, N*3) where T is number of frames and N is number of joints
+        source_fps: Source frame rate
+        target_fps: Target frame rate
+        device: Device to run computations on
+        interpolation_type: Type of interpolation to use ("linear" or "slerp")
+
+    Returns:
+        Interpolated pose_aa with new frame rate, shape (T_new, N*3)
+    """
+    # pose_aa: (T, N*3)
+    orig_shape = pose_aa.shape[1:]
+    if pose_aa.ndim != 2:
+        pose_aa = pose_aa.reshape(pose_aa.shape[0], -1)
+    T, D = pose_aa.shape
+
+    if interpolation_type == "linear":
+        # Direct linear interpolation on angle-axis representation
+        duration = (T - 1) * (1 / source_fps)
+        times = torch.arange(0, duration + 1e-6, 1 / target_fps, dtype=torch.float32, device=device)
+        times = times[times <= duration]
+
+        # Compute frame indices and blend factors for linear interpolation
+        frame_indices = times * source_fps
+        index_0 = torch.floor(frame_indices).long()
+        index_1 = torch.min(index_0 + 1, torch.tensor(T - 1, device=device))
+        blend = frame_indices - index_0.float()
+
+        # Linear interpolation on the entire 2D tensor
+        pose_aa_interp = (1 - blend.unsqueeze(1)) * pose_aa[index_0] + blend.unsqueeze(1) * pose_aa[
+            index_1
+        ]
+        pose_aa_interp = pose_aa_interp.view(pose_aa_interp.shape[0], *orig_shape)
+        if pose_aa.dtype == torch.int64:
+            pose_aa_interp = pose_aa_interp.round()
+        pose_aa_interp = pose_aa_interp.type_as(pose_aa)
+        return pose_aa_interp
+
+    elif interpolation_type == "slerp":
+        dim = 3 if rot_type == "aa" else 4
+        N = D // dim
+        pose_aa_reshaped = pose_aa.view(T, N, dim)
+        # Original spherical linear interpolation on quaternions
+        pose_aa_interp_list = []
+        for i in range(N):
+            # Convert angle-axis to quaternion for this joint
+            if rot_type == "aa":
+                pose_quat = angle_axis_to_quaternion(pose_aa_reshaped[:, i])  # (T, 4)
+            else:
+                pose_quat = pose_aa_reshaped[:, i]
+            pose_quat_batch = pose_quat.unsqueeze(0)  # (1, T, 4)
+            pose_quat_interp = interpolate_quaternions(pose_quat_batch, source_fps, target_fps)
+            pose_quat_interp = pose_quat_interp[0]  # (T_new, 4)
+            if rot_type == "aa":
+                pose_aa_interp = quaternion_to_angle_axis(pose_quat_interp)  # (T_new, 3)
+            else:
+                pose_aa_interp = pose_quat_interp
+            pose_aa_interp_list.append(pose_aa_interp)
+
+        # Concatenate all joints: (T_new, N, 3) -> (T_new, N*3)
+        pose_aa_interp = torch.stack(pose_aa_interp_list, dim=1)  # (T_new, N, 3)
+        pose_aa_interp = pose_aa_interp.view(pose_aa_interp.shape[0], -1)  # (T_new, N*3)
+        pose_aa_interp = pose_aa_interp.view(pose_aa_interp.shape[0], *orig_shape).to(pose_aa)
+        return pose_aa_interp
+
+    else:
+        raise ValueError(
+            f"Unsupported interpolation_type: {interpolation_type}. Must be 'linear' or 'slerp'."
+        )
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/order_converter.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/order_converter.py
new file mode 100644
index 0000000000000000000000000000000000000000..078f2fc3bb794403f5ea58e72b506e9bcba6fdc2
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/order_converter.py
@@ -0,0 +1,248 @@
+"""
+IsaacLab <-> MuJoCo ordering conversion utilities.
+
+This module provides utilities for converting both DOF and body ordering
+between IsaacLab and MuJoCo conventions for humanoid robots.
+
+qpos format: [root_trans(3), root_quat(4), dof_angles(N)]
+"""
+
+from abc import ABC
+
+import numpy as np
+import torch
+
+
+class IsaacLabMuJoCoConverter(ABC):
+    """Abstract base class for DOF/body order conversion between IsaacLab and MuJoCo.
+
+    Subclasses must define DOF_MAPPINGS for their specific robot.
+    """
+
+    ROOT_QPOS_OFFSET = 7  # root_trans(3) + root_quat(4)
+    DOF_MAPPINGS: dict = {}  # Must be overridden by subclasses
+    VALID_DOF_ORDERS = ("mujoco", "isaaclab")
+
+    def convert(self, data: torch.Tensor, from_order: str, to_order: str) -> torch.Tensor:
+        """Convert DOF order between conventions.
+
+        Auto-detects input format:
+            - qpos:           [..., 7 + num_dof]           -> reorder DOF portion, keep root
+            - dof angles:     [..., num_dof]                -> reorder directly
+            - body transforms [..., num_dof + 1, D]         -> keep root body, reorder DOF bodies
+            - body transforms [..., num_dof, D]             -> reorder directly
+            - rotation mats   [..., num_dof (+ 1), 3, 3]   -> same as above on dim -3
+
+        Note: Body reordering for per-body transforms uses DOF_MAPPINGS (not
+        BODY_MAPPINGS) because each DOF link has exactly one parent body in G1,
+        so DOF order and DOF-body order are consistent by construction.  The
+        separate BODY_MAPPINGS include the root (pelvis) at index 0 and are used
+        only when the full 30-body ordering is needed.
+        """
+        if from_order == to_order:
+            return data
+
+        mapping = self.DOF_MAPPINGS[(from_order, to_order)]
+        last_dim = data.shape[-1]
+
+        if last_dim == self.ROOT_QPOS_OFFSET + self.num_dof:
+            # qpos: [..., 7 + num_dof]
+            return torch.cat(
+                [
+                    data[..., : self.ROOT_QPOS_OFFSET],
+                    data[..., self.ROOT_QPOS_OFFSET :][..., mapping],
+                ],
+                dim=-1,
+            )
+        elif last_dim == self.num_dof:
+            # Raw DOF angles: [..., num_dof]
+            return data[..., mapping]
+        else:
+            # Per-body transforms: body dim is -3 for [..., J, 3, 3], else -2 for [..., J, D]
+            body_dim = (
+                -3 if (data.ndim >= 3 and data.shape[-1] == 3 and data.shape[-2] == 3) else -2
+            )
+            num_bodies = data.shape[body_dim]
+
+            if num_bodies == self.num_dof + 1:
+                # Includes root at index 0 — keep root, reorder DOF bodies
+                body_order = [0] + [m + 1 for m in mapping]
+            elif num_bodies == self.num_dof:
+                body_order = mapping
+            else:
+                raise ValueError(
+                    f"Cannot detect format: last_dim={last_dim}, body_dim size={num_bodies}, "
+                    f"expected num_dof+1={self.num_dof + 1} or num_dof={self.num_dof}"
+                )
+
+            if body_dim == -2:
+                return data[..., body_order, :]
+            else:  # body_dim == -3
+                return data[..., body_order, :, :]
+
+    def to_mujoco(self, data: torch.Tensor) -> torch.Tensor:
+        """Convert from IsaacLab to MuJoCo DOF order."""
+        return self.convert(data, from_order="isaaclab", to_order="mujoco")
+
+    def to_isaaclab(self, data: torch.Tensor) -> torch.Tensor:
+        """Convert from MuJoCo to IsaacLab DOF order."""
+        return self.convert(data, from_order="mujoco", to_order="isaaclab")
+
+    @property
+    def num_dof(self) -> int:
+        """Number of actuated DOFs (excluding root)."""
+        return len(self.DOF_MAPPINGS[(self.VALID_DOF_ORDERS[0], self.VALID_DOF_ORDERS[1])])
+
+
+class G1Converter(IsaacLabMuJoCoConverter):
+    """G1 ordering converter.
+
+    Imports G1 body/DOF/joint mappings from gear_sonic.envs.manager_env.robots.g1.
+    """
+
+    def __init__(self):
+        # Lazy import to avoid circular dependency:
+        # order_converter -> g1 -> mdp/__init__ -> commands -> order_converter
+        from gear_sonic.envs.manager_env.robots.g1 import (
+            G1_ISAACLAB_JOINTS,
+            G1_ISAACLAB_TO_MUJOCO_BODY,
+            G1_ISAACLAB_TO_MUJOCO_DOF,
+            G1_MUJOCO_TO_ISAACLAB_BODY,
+            G1_MUJOCO_TO_ISAACLAB_DOF,
+        )
+
+        self.JOINT_NAMES = G1_ISAACLAB_JOINTS
+        self.DOF_MAPPINGS = {
+            ("isaaclab", "mujoco"): G1_ISAACLAB_TO_MUJOCO_DOF,
+            ("mujoco", "isaaclab"): G1_MUJOCO_TO_ISAACLAB_DOF,
+        }
+        self.BODY_MAPPINGS = {
+            ("isaaclab", "mujoco"): G1_ISAACLAB_TO_MUJOCO_BODY,
+            ("mujoco", "isaaclab"): G1_MUJOCO_TO_ISAACLAB_BODY,
+        }
+
+    # Body subset names for MPJPE metrics (used by reconstruction_trainer)
+    VR_3POINTS_BODY_NAMES = ["torso_link", "left_wrist_yaw_link", "right_wrist_yaw_link"]
+    FOOT_BODY_NAMES = ["left_ankle_roll_link", "right_ankle_roll_link"]
+
+    @property
+    def vr_3points_mujoco_indices(self):
+        """VR 3-point body indices in full (30-body) MuJoCo body order.
+
+        These index into the full body array after isaaclab_to_mujoco_body
+        reordering, NOT the 14-body motion.yaml body_names subset.
+        """
+        mj_names = [self.JOINT_NAMES[i] for i in self.isaaclab_to_mujoco_body]
+        return [mj_names.index(n) for n in self.VR_3POINTS_BODY_NAMES]
+
+    @property
+    def foot_mujoco_indices(self):
+        """Foot body indices in full (30-body) MuJoCo body order.
+
+        These index into the full body array after isaaclab_to_mujoco_body
+        reordering, NOT the 14-body motion.yaml body_names subset.
+        """
+        mj_names = [self.JOINT_NAMES[i] for i in self.isaaclab_to_mujoco_body]
+        return [mj_names.index(n) for n in self.FOOT_BODY_NAMES]
+
+    @property
+    def isaaclab_to_mujoco_dof(self):
+        """DOF reorder indices: IsaacLab -> MuJoCo."""
+        return self.DOF_MAPPINGS[("isaaclab", "mujoco")]
+
+    @property
+    def mujoco_to_isaaclab_dof(self):
+        """DOF reorder indices: MuJoCo -> IsaacLab."""
+        return self.DOF_MAPPINGS[("mujoco", "isaaclab")]
+
+    @property
+    def isaaclab_to_mujoco_body(self):
+        """Body reorder indices: IsaacLab -> MuJoCo."""
+        return self.BODY_MAPPINGS[("isaaclab", "mujoco")]
+
+    @property
+    def mujoco_to_isaaclab_body(self):
+        """Body reorder indices: MuJoCo -> IsaacLab."""
+        return self.BODY_MAPPINGS[("mujoco", "isaaclab")]
+
+    def get_isaaclab_to_mujoco_mapping(self):
+        """Return the full mapping dict for body/DOF reordering."""
+        return {
+            "isaaclab_joints": self.JOINT_NAMES,
+            "isaaclab_to_mujoco_dof": self.isaaclab_to_mujoco_dof,
+            "mujoco_to_isaaclab_dof": self.mujoco_to_isaaclab_dof,
+            "isaaclab_to_mujoco_body": self.isaaclab_to_mujoco_body,
+            "mujoco_to_isaaclab_body": self.mujoco_to_isaaclab_body,
+        }
+
+
+class H2Converter(IsaacLabMuJoCoConverter):
+    """H2 robot joint/body order converter between IsaacLab and MuJoCo conventions."""
+
+    def __init__(self):
+        from gear_sonic.envs.manager_env.robots.h2 import (
+            H2_ISAACLAB_JOINTS,
+            H2_ISAACLAB_TO_MUJOCO_BODY,
+            H2_ISAACLAB_TO_MUJOCO_DOF,
+            H2_MUJOCO_TO_ISAACLAB_BODY,
+            H2_MUJOCO_TO_ISAACLAB_DOF,
+        )
+
+        self.JOINT_NAMES = H2_ISAACLAB_JOINTS
+        self.DOF_MAPPINGS = {
+            ("isaaclab", "mujoco"): H2_ISAACLAB_TO_MUJOCO_DOF,
+            ("mujoco", "isaaclab"): H2_MUJOCO_TO_ISAACLAB_DOF,
+        }
+        self.BODY_MAPPINGS = {
+            ("isaaclab", "mujoco"): H2_ISAACLAB_TO_MUJOCO_BODY,
+            ("mujoco", "isaaclab"): H2_MUJOCO_TO_ISAACLAB_BODY,
+        }
+
+    VR_3POINTS_BODY_NAMES = ["torso_link", "left_wrist_pitch_link", "right_wrist_pitch_link"]
+    FOOT_BODY_NAMES = ["left_ankle_roll_link", "right_ankle_roll_link"]
+
+
+def load_qpos_from_csv(csv_path: str) -> torch.Tensor:
+    """Load qpos [T, D] from CSV."""
+    import pandas as pd
+
+    return torch.from_numpy(pd.read_csv(csv_path).values.astype(np.float32))
+
+
+def save_qpos_to_csv(qpos: torch.Tensor, csv_path: str):
+    """Save qpos to CSV."""
+    import pandas as pd
+
+    data = qpos[0].cpu().numpy() if qpos.dim() == 3 else qpos.cpu().numpy()
+    pd.DataFrame(data).to_csv(csv_path, index=False)
+
+
+if __name__ == "__main__":
+    from isaacsim import SimulationApp
+
+    _sim_app = SimulationApp({"headless": True})
+
+    # Test DOF conversion round-trip
+    converter = G1Converter()
+
+    # Create random qpos (G1 has 29 DOFs)
+    T, num_dof = 30, 29
+    qpos = torch.cat(
+        [
+            torch.tensor([[0.0, 0.0, 1.0]]).expand(T, 3),
+            torch.tensor([[1.0, 0.0, 0.0, 0.0]]).expand(T, 4),
+            torch.randn(T, num_dof) * 0.3,
+        ],
+        dim=-1,
+    )
+
+    # Test round-trip conversion
+    qpos_mujoco = converter.to_mujoco(qpos)
+    qpos_back = converter.to_isaaclab(qpos_mujoco)
+    print(f"Round-trip error: {(qpos - qpos_back).abs().max():.2e}")
+
+    # Test explicit convert API
+    qpos_mujoco2 = converter.convert(qpos, from_order="isaaclab", to_order="mujoco")
+    print(f"to_mujoco vs convert match: {torch.allclose(qpos_mujoco, qpos_mujoco2)}")
+
+    _sim_app.close()
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/rl.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/rl.py
new file mode 100644
index 0000000000000000000000000000000000000000..3eee2916ac036c88425976f8e4b93ce4b875591f
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/rl.py
@@ -0,0 +1,29 @@
+"""RL-specific helpers: episode attention masks and (legacy) schedule utilities."""
+
+import torch
+
+
+def compute_episode_attnmask(dones):
+    """
+    Compute an attention mask that prevents the model from attending to observations from different episodes.
+
+    Args:
+        dones (torch.Tensor): A tensor of shape (num_envs, num_steps) indicating when each environment episode ends.
+                                A value of 1.0 indicates the end of an episode.
+
+    Returns:
+        torch.Tensor: An attention mask of shape (num_envs, num_steps, num_steps) where True values indicate
+                        positions that should be masked (i.e., the model should not attend to these positions).
+    """
+    # Create cumulative sum of dones to identify different episodes
+    episode_starts = torch.roll(dones, 1, dims=1)
+    episode_starts[:, 0] = True  # First step is always start of an episode
+    episode_ids = torch.cumsum(episode_starts, dim=1)  # (num_envs, num_steps)
+
+    # Expand episode_ids for broadcasting
+    episode_ids_i = episode_ids.unsqueeze(2)  # (num_envs, num_steps, 1)
+    episode_ids_j = episode_ids.unsqueeze(1)  # (num_envs, 1, num_steps)
+
+    # Create mask where True indicates positions from different episodes
+    attnmask = episode_ids_i != episode_ids_j
+    return attnmask
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/rotation_conversion.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/rotation_conversion.py
new file mode 100644
index 0000000000000000000000000000000000000000..eec519cd5df8c2a5f550970a95ca1250152e69de
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/rotation_conversion.py
@@ -0,0 +1,618 @@
+# This file contains code derived from PyTorch3D (via https://github.com/Mathux/ACTOR),
+# originally Copyright (c) Facebook, Inc. and its affiliates, licensed under BSD.
+# See https://github.com/facebookresearch/pytorch3d/blob/main/LICENSE for the original license.
+
+"""PyTorch 3D rotation conversions (quaternion, rotation matrix, axis-angle, 6D).
+
+Provides batched, differentiable conversions between quaternion (wxyz),
+rotation matrix, axis-angle, and 6D rotation representations, plus
+swing-twist decomposition and yaw extraction.
+
+The transformation matrices returned from the functions in this file assume
+the points on which the transformation will be applied are column vectors.
+i.e. the R matrix is structured as
+
+    R = [
+            [Rxx, Rxy, Rxz],
+            [Ryx, Ryy, Ryz],
+            [Rzx, Rzy, Rzz],
+        ]  # (3, 3)
+
+This matrix can be applied to column vectors by post multiplication
+by the points e.g.
+
+    points = [[0], [1], [2]]  # (3 x 1) xyz coordinates of a point
+    transformed_points = R * points
+
+To apply the same matrix to points which are row vectors, the R matrix
+can be transposed and pre multiplied by the points:
+
+e.g.
+    points = [[0, 1, 2]]  # (1 x 3) xyz coordinates of a point
+    transformed_points = points * R.transpose(1, 0)
+"""
+
+import functools
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+
+try:
+    from pytorch3d.transforms.rotation_conversions import matrix_to_axis_angle
+except ImportError:
+    matrix_to_axis_angle = None
+
+
+def quaternion_to_matrix(quaternions):
+    """
+    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))
+
+
+def _copysign(a, b):
+    """
+    Return a tensor where each element has the absolute value taken from the,
+    corresponding element of a, with sign taken from the corresponding
+    element of b. This is like the standard copysign floating-point operation,
+    but is not careful about negative 0 and NaN.
+
+    Args:
+        a: source tensor.
+        b: tensor whose signs will be used, of the same shape as a.
+
+    Returns:
+        Tensor of the same shape as a with the signs of b.
+    """
+    signs_differ = (a < 0) != (b < 0)
+    return torch.where(signs_differ, -a, a)
+
+
+def _sqrt_positive_part(x):
+    """
+    Returns torch.sqrt(torch.max(0, x))
+    but with a zero subgradient where x is 0.
+    """
+    return torch.where(x > 0, torch.sqrt(torch.clamp_min(x, 1e-5)), 0)
+
+
+def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor:
+    """
+    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(
+        [
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1),
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1),
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1),
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            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)
+    out = quat_candidates[F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, :].reshape(
+        batch_dim + (4,)
+    )
+    return standardize_quaternion(out)
+
+
+def _axis_angle_rotation(axis: str, angle):
+    """
+    Return the rotation matrices for one of the rotations about an axis
+    of which Euler angles describe, for each value of the angle given.
+
+    Args:
+        axis: Axis label "X" or "Y or "Z".
+        angle: any shape tensor of Euler angles in radians
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+
+    cos = torch.cos(angle)
+    sin = torch.sin(angle)
+    one = torch.ones_like(angle)
+    zero = torch.zeros_like(angle)
+
+    if axis == "X":
+        R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos)
+    if axis == "Y":
+        R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos)
+    if axis == "Z":
+        R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one)
+
+    return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3))
+
+
+def euler_angles_to_matrix(euler_angles, convention: str):
+    """
+    Convert rotations given as Euler angles in radians to rotation matrices.
+
+    Args:
+        euler_angles: Euler angles in radians as tensor of shape (..., 3).
+        convention: Convention string of three uppercase letters from
+            {"X", "Y", and "Z"}.
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3:
+        raise ValueError("Invalid input euler angles.")
+    if len(convention) != 3:
+        raise ValueError("Convention must have 3 letters.")
+    if convention[1] in (convention[0], convention[2]):
+        raise ValueError(f"Invalid convention {convention}.")
+    for letter in convention:
+        if letter not in ("X", "Y", "Z"):
+            raise ValueError(f"Invalid letter {letter} in convention string.")
+    matrices = map(_axis_angle_rotation, convention, torch.unbind(euler_angles, -1))
+    return functools.reduce(torch.matmul, matrices)
+
+
+def _angle_from_tan(axis: str, other_axis: str, data, horizontal: bool, tait_bryan: bool):
+    """
+    Extract the first or third Euler angle from the two members of
+    the matrix which are positive constant times its sine and cosine.
+
+    Args:
+        axis: Axis label "X" or "Y or "Z" for the angle we are finding.
+        other_axis: Axis label "X" or "Y or "Z" for the middle axis in the
+            convention.
+        data: Rotation matrices as tensor of shape (..., 3, 3).
+        horizontal: Whether we are looking for the angle for the third axis,
+            which means the relevant entries are in the same row of the
+            rotation matrix. If not, they are in the same column.
+        tait_bryan: Whether the first and third axes in the convention differ.
+
+    Returns:
+        Euler Angles in radians for each matrix in dataset as a tensor
+        of shape (...).
+    """
+
+    i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis]
+    if horizontal:
+        i2, i1 = i1, i2
+    even = (axis + other_axis) in ["XY", "YZ", "ZX"]
+    if horizontal == even:
+        return torch.atan2(data[..., i1], data[..., i2])
+    if tait_bryan:
+        return torch.atan2(-data[..., i2], data[..., i1])
+    return torch.atan2(data[..., i2], -data[..., i1])
+
+
+def _index_from_letter(letter: str):
+    if letter == "X":
+        return 0
+    if letter == "Y":
+        return 1
+    if letter == "Z":
+        return 2
+
+
+def matrix_to_euler_angles(matrix, convention: str):
+    """
+    Convert rotations given as rotation matrices to Euler angles in radians.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+        convention: Convention string of three uppercase letters.
+
+    Returns:
+        Euler angles in radians as tensor of shape (..., 3).
+    """
+    if len(convention) != 3:
+        raise ValueError("Convention must have 3 letters.")
+    if convention[1] in (convention[0], convention[2]):
+        raise ValueError(f"Invalid convention {convention}.")
+    for letter in convention:
+        if letter not in ("X", "Y", "Z"):
+            raise ValueError(f"Invalid letter {letter} in convention string.")
+    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
+        raise ValueError(f"Invalid rotation matrix  shape f{matrix.shape}.")
+    i0 = _index_from_letter(convention[0])
+    i2 = _index_from_letter(convention[2])
+    tait_bryan = i0 != i2
+    if tait_bryan:
+        central_angle = torch.asin(matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0))
+    else:
+        central_angle = torch.acos(matrix[..., i0, i0])
+
+    o = (
+        _angle_from_tan(convention[0], convention[1], matrix[..., i2], False, tait_bryan),
+        central_angle,
+        _angle_from_tan(convention[2], convention[1], matrix[..., i0, :], True, tait_bryan),
+    )
+    return torch.stack(o, -1)
+
+
+def random_quaternions(
+    n: int, dtype: Optional[torch.dtype] = None, device=None, requires_grad=False
+):
+    """
+    Generate random quaternions representing rotations,
+    i.e. versors with nonnegative real part.
+
+    Args:
+        n: Number of quaternions in a batch to return.
+        dtype: Type to return.
+        device: Desired device of returned tensor. Default:
+            uses the current device for the default tensor type.
+        requires_grad: Whether the resulting tensor should have the gradient
+            flag set.
+
+    Returns:
+        Quaternions as tensor of shape (N, 4).
+    """
+    o = torch.randn((n, 4), dtype=dtype, device=device, requires_grad=requires_grad)
+    s = (o * o).sum(1)
+    o = o / _copysign(torch.sqrt(s), o[:, 0])[:, None]
+    return o
+
+
+def random_rotations(n: int, dtype: Optional[torch.dtype] = None, device=None, requires_grad=False):
+    """
+    Generate random rotations as 3x3 rotation matrices.
+
+    Args:
+        n: Number of rotation matrices in a batch to return.
+        dtype: Type to return.
+        device: Device of returned tensor. Default: if None,
+            uses the current device for the default tensor type.
+        requires_grad: Whether the resulting tensor should have the gradient
+            flag set.
+
+    Returns:
+        Rotation matrices as tensor of shape (n, 3, 3).
+    """
+    quaternions = random_quaternions(n, dtype=dtype, device=device, requires_grad=requires_grad)
+    return quaternion_to_matrix(quaternions)
+
+
+def random_rotation(dtype: Optional[torch.dtype] = None, device=None, requires_grad=False):
+    """
+    Generate a single random 3x3 rotation matrix.
+
+    Args:
+        dtype: Type to return
+        device: Device of returned tensor. Default: if None,
+            uses the current device for the default tensor type
+        requires_grad: Whether the resulting tensor should have the gradient
+            flag set
+
+    Returns:
+        Rotation matrix as tensor of shape (3, 3).
+    """
+    return random_rotations(1, dtype, device, requires_grad)[0]
+
+
+def standardize_quaternion(quaternions):
+    """
+    Convert a unit quaternion to a standard form: one in which the real
+    part is non negative.
+
+    Args:
+        quaternions: Quaternions with real part first,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Standardized quaternions as tensor of shape (..., 4).
+    """
+    return torch.where(quaternions[..., 0:1] < 0, -quaternions, quaternions)
+
+
+def quaternion_raw_multiply(a, b):
+    """
+    Multiply two quaternions.
+    Usual torch rules for broadcasting apply.
+
+    Args:
+        a: Quaternions as tensor of shape (..., 4), real part first.
+        b: Quaternions as tensor of shape (..., 4), real part first.
+
+    Returns:
+        The product of a and b, a tensor of quaternions shape (..., 4).
+    """
+    aw, ax, ay, az = torch.unbind(a, -1)
+    bw, bx, by, bz = torch.unbind(b, -1)
+    ow = aw * bw - ax * bx - ay * by - az * bz
+    ox = aw * bx + ax * bw + ay * bz - az * by
+    oy = aw * by - ax * bz + ay * bw + az * bx
+    oz = aw * bz + ax * by - ay * bx + az * bw
+    return torch.stack((ow, ox, oy, oz), -1)
+
+
+def quaternion_multiply(a, b):
+    """
+    Multiply two quaternions representing rotations, returning the quaternion
+    representing their composition, i.e. the versor with nonnegative real part.
+    Usual torch rules for broadcasting apply.
+
+    Args:
+        a: Quaternions as tensor of shape (..., 4), real part first.
+        b: Quaternions as tensor of shape (..., 4), real part first.
+
+    Returns:
+        The product of a and b, a tensor of quaternions of shape (..., 4).
+    """
+    ab = quaternion_raw_multiply(a, b)
+    return standardize_quaternion(ab)
+
+
+def quaternion_invert(quaternion):
+    """
+    Given a quaternion representing rotation, get the quaternion representing
+    its inverse.
+
+    Args:
+        quaternion: Quaternions as tensor of shape (..., 4), with real part
+            first, which must be versors (unit quaternions).
+
+    Returns:
+        The inverse, a tensor of quaternions of shape (..., 4).
+    """
+
+    return quaternion * quaternion.new_tensor([1, -1, -1, -1])
+
+
+def quaternion_apply(quaternion, point):
+    """
+    Apply the rotation given by a quaternion to a 3D point.
+    Usual torch rules for broadcasting apply.
+
+    Args:
+        quaternion: Tensor of quaternions, real part first, of shape (..., 4).
+        point: Tensor of 3D points of shape (..., 3).
+
+    Returns:
+        Tensor of rotated points of shape (..., 3).
+    """
+    if point.size(-1) != 3:
+        raise ValueError(f"Points are not in 3D, f{point.shape}.")
+    real_parts = point.new_zeros(point.shape[:-1] + (1,))
+    point_as_quaternion = torch.cat((real_parts, point), -1)
+    out = quaternion_raw_multiply(
+        quaternion_raw_multiply(quaternion, point_as_quaternion),
+        quaternion_invert(quaternion),
+    )
+    return out[..., 1:]
+
+
+def axis_angle_to_matrix(axis_angle):
+    """
+    Convert rotations given as axis/angle to rotation matrices.
+
+    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:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle))
+
+
+def matrix_to_axis_angle(matrix):
+    """
+    Convert rotations given as rotation matrices to axis/angle.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+
+    Returns:
+        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.
+    """
+    return quaternion_to_axis_angle(matrix_to_quaternion(matrix))
+
+
+def axis_angle_to_quaternion(axis_angle):
+    """
+    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 = 0.5 * angles
+    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
+    )
+    """
+    sin_half_angles_over_angles = torch.where(
+        angles.abs() < eps, 0.5 - (angles * angles) / 48, torch.sin(half_angles) / angles
+    )
+    quaternions = torch.cat(
+        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
+    )
+    return quaternions
+
+
+def quaternion_to_axis_angle(quaternions):
+    """
+    Convert rotations given as quaternions to axis/angle.
+
+    Args:
+        quaternions: quaternions with real part first,
+            as tensor of shape (..., 4).
+
+    Returns:
+        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.
+    """
+    norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True)
+    half_angles = torch.atan2(norms, quaternions[..., :1])
+    angles = 2 * half_angles
+    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
+    )
+    """
+    sin_half_angles_over_angles = torch.where(
+        angles.abs() < eps, 0.5 - (angles * angles) / 48, torch.sin(half_angles) / angles
+    )
+    return quaternions[..., 1:] / sin_half_angles_over_angles
+
+
+def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor:
+    """
+    Converts 6D rotation representation by Zhou et al. [1] to rotation matrix
+    using Gram--Schmidt orthogonalisation per Section B of [1].
+    Args:
+        d6: 6D rotation representation, of size (*, 6)
+
+    Returns:
+        batch of rotation matrices of size (*, 3, 3)
+
+    [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
+    On the Continuity of Rotation Representations in Neural Networks.
+    IEEE Conference on Computer Vision and Pattern Recognition, 2019.
+    Retrieved from http://arxiv.org/abs/1812.07035
+    """
+
+    a1, a2 = d6[..., :3], d6[..., 3:]
+    b1 = F.normalize(a1, dim=-1)
+    b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1
+    b2 = F.normalize(b2, dim=-1)
+    b3 = torch.cross(b1, b2, dim=-1)
+    return torch.stack((b1, b2, b3), dim=-2)
+
+
+def matrix_to_rotation_6d(matrix: torch.Tensor) -> torch.Tensor:
+    """
+    Converts rotation matrices to 6D rotation representation by Zhou et al. [1]
+    by dropping the last row. Note that 6D representation is not unique.
+    Args:
+        matrix: batch of rotation matrices of size (*, 3, 3)
+
+    Returns:
+        6D rotation representation, of size (*, 6)
+
+    [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
+    On the Continuity of Rotation Representations in Neural Networks.
+    IEEE Conference on Computer Vision and Pattern Recognition, 2019.
+    Retrieved from http://arxiv.org/abs/1812.07035
+    """
+    return matrix[..., :2, :].clone().reshape(*matrix.size()[:-2], 6)
+
+
+def quaternion_multiply_np(a, b):
+    aw, ax, ay, az = np.split(a, 4, axis=1)
+    bw, bx, by, bz = np.split(b, 4, axis=1)
+    ow = aw * bw - ax * bx - ay * by - az * bz
+    ox = aw * bx + ax * bw + ay * bz - az * by
+    oy = aw * by - ax * bz + ay * bw + az * bx
+    oz = aw * bz + ax * by - ay * bx + az * bw
+    return np.concatenate([ow, ox, oy, oz], axis=1)
+
+
+def decompose_rotation_aa(rotation_aa, v2):
+    angle = np.linalg.norm(rotation_aa, axis=1)[:, None]
+    w = np.cos(angle / 2)
+    v = np.sin(angle / 2) * rotation_aa / angle
+    q = np.concatenate([w, v], axis=1)
+
+    v_twist = np.dot(v, v2)[:, None] * v2
+    q_twist = np.concatenate([w, v_twist], axis=1)
+    q_twist = q_twist / np.linalg.norm(q_twist, axis=1)[:, None]
+
+    q_twist_inv = q_twist * np.array([1, -1, -1, -1])
+    q_swing = quaternion_multiply_np(q_twist_inv, q)
+
+    return q_twist, q_swing
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/scheduler.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/scheduler.py
new file mode 100644
index 0000000000000000000000000000000000000000..2bef76dda25ba82b7ea792b8fddb6675c0ed71ab
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/scheduler.py
@@ -0,0 +1,469 @@
+"""Utilities for scheduled parameter updates and learning-rate scheduling.
+
+Includes object-path navigation for dynamically accessing/mutating nested
+config attributes, a WarmupCosineScheduler for LR with linear warm-up and
+cosine decay, and helpers for managing parameter change schedules.
+"""
+
+import numpy
+import torch
+import math
+import re
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import _LRScheduler
+from omegaconf.dictconfig import DictConfig
+
+
+def _navigate_object_path(obj, path, split_char="@"):
+    """
+    Navigate through a complex object path that may include:
+    - Attribute access: obj.attr
+    - Function calls: obj.method('param')
+    - Dictionary/array access: obj['key'][0]
+    - Mixed combinations: obj.method('param')['key'][0].attr
+    """
+    current_obj = obj
+
+    # Split the path by the split_char and process each segment
+    segments = path.split(split_char)
+
+    for segment in segments:
+        current_obj = _process_path_segment(current_obj, segment)
+
+    return current_obj
+
+
+def _process_path_segment(obj, segment):
+    """
+    Process a single path segment that may contain:
+    - Simple attribute: attr
+    - Function call: method('param')
+    - Bracket access: ['key'][0]
+    - Combined: method('param')['key']
+    """
+    current_obj = obj
+
+    # Parse the segment to identify different access patterns
+    i = 0
+    while i < len(segment):
+        if segment[i] == "[":
+            # Handle bracket access
+            bracket_end = _find_matching_bracket(segment, i)
+            bracket_content = segment[i + 1 : bracket_end]
+
+            # Evaluate the bracket content
+            if bracket_content.startswith("'") and bracket_content.endswith("'"):
+                # String key
+                key = bracket_content[1:-1]
+                current_obj = current_obj[key]
+            elif bracket_content.startswith('"') and bracket_content.endswith('"'):
+                # String key with double quotes
+                key = bracket_content[1:-1]
+                current_obj = current_obj[key]
+            elif bracket_content.lstrip("-").isdigit():
+                # Numeric index
+                index = int(bracket_content)
+                current_obj = current_obj[index]
+            else:
+                # Try to evaluate as expression (for complex keys)
+                try:
+                    key = eval(bracket_content)
+                    current_obj = current_obj[key]
+                except:
+                    # Fallback to string key
+                    current_obj = current_obj[bracket_content]
+
+            i = bracket_end + 1
+
+        else:
+            # Handle attribute access or function call
+            attr_start = i
+            # Find the end of the identifier (attribute or method name)
+            while i < len(segment) and (segment[i].isalnum() or segment[i] == "_"):
+                i += 1
+
+            if attr_start < i:
+                attr_name = segment[attr_start:i]
+
+                # Check if this is followed by parentheses (function call)
+                if i < len(segment) and segment[i] == "(":
+                    # This is a function call
+                    paren_end = _find_matching_paren(segment, i)
+                    args_str = segment[i + 1 : paren_end]
+
+                    # Parse and evaluate arguments
+                    args = _parse_function_args(args_str)
+
+                    # Call the method
+                    method = getattr(current_obj, attr_name)
+                    current_obj = method(*args)
+
+                    i = paren_end + 1
+                else:
+                    # This is a simple attribute access
+                    if attr_name.lstrip("-").isdigit():
+                        # Numeric index for direct access
+                        current_obj = current_obj[int(attr_name)]
+                    else:
+                        # Attribute access
+                        current_obj = getattr(current_obj, attr_name)
+            else:
+                # Skip non-alphanumeric characters that aren't brackets or parentheses
+                i += 1
+
+    return current_obj
+
+
+def _find_matching_bracket(s, start):
+    """Find the matching closing bracket for an opening bracket at position start."""
+    count = 1
+    i = start + 1
+    while i < len(s) and count > 0:
+        if s[i] == "[":
+            count += 1
+        elif s[i] == "]":
+            count -= 1
+        i += 1
+    return i - 1
+
+
+def _find_matching_paren(s, start):
+    """Find the matching closing parenthesis for an opening parenthesis at position start."""
+    count = 1
+    i = start + 1
+    while i < len(s) and count > 0:
+        if s[i] == "(":
+            count += 1
+        elif s[i] == ")":
+            count -= 1
+        i += 1
+    return i - 1
+
+
+def _parse_function_args(args_str):
+    """Parse function arguments from a string."""
+    if not args_str.strip():
+        return []
+
+    args = []
+    current_arg = ""
+    paren_count = 0
+    bracket_count = 0
+    in_quotes = False
+    quote_char = None
+
+    for char in args_str:
+        if char in ['"', "'"] and not in_quotes:
+            in_quotes = True
+            quote_char = char
+            current_arg += char
+        elif char == quote_char and in_quotes:
+            in_quotes = False
+            quote_char = None
+            current_arg += char
+        elif not in_quotes:
+            if char == "(":
+                paren_count += 1
+                current_arg += char
+            elif char == ")":
+                paren_count -= 1
+                current_arg += char
+            elif char == "[":
+                bracket_count += 1
+                current_arg += char
+            elif char == "]":
+                bracket_count -= 1
+                current_arg += char
+            elif char == "," and paren_count == 0 and bracket_count == 0:
+                args.append(_evaluate_arg(current_arg.strip()))
+                current_arg = ""
+            else:
+                current_arg += char
+        else:
+            current_arg += char
+
+    if current_arg.strip():
+        args.append(_evaluate_arg(current_arg.strip()))
+
+    return args
+
+
+def _evaluate_arg(arg_str):
+    """Evaluate a function argument string to its proper type."""
+    arg_str = arg_str.strip()
+
+    # String literals
+    if (arg_str.startswith("'") and arg_str.endswith("'")) or (
+        arg_str.startswith('"') and arg_str.endswith('"')
+    ):
+        return arg_str[1:-1]
+
+    # Numeric literals
+    if arg_str.lstrip("-").replace(".", "").isdigit():
+        if "." in arg_str:
+            return float(arg_str)
+        else:
+            return int(arg_str)
+
+    # Boolean literals
+    if arg_str.lower() == "true":
+        return True
+    elif arg_str.lower() == "false":
+        return False
+    elif arg_str.lower() == "none":
+        return None
+
+    # For complex expressions, try eval (be careful in production)
+    try:
+        return eval(arg_str)
+    except:
+        # Fallback to string
+        return arg_str
+
+
+def _get_final_target(obj, target_attr):
+    """Get the final target object for reading, handling complex paths."""
+    if _is_complex_path(target_attr):
+        return _process_path_segment(obj, target_attr)
+    else:
+        # Simple attribute or numeric index
+        if target_attr.lstrip("-").isdigit():
+            return obj[int(target_attr)]
+        else:
+            return getattr(obj, target_attr)
+
+
+def _set_final_target(obj, target_attr, value):
+    """Set the final target value, handling complex paths."""
+    if _is_complex_path(target_attr):
+        # For complex paths, we need to navigate to the parent and set the final element
+        _set_complex_path_value(obj, target_attr, value)
+    else:
+        # Simple attribute or numeric index
+        if target_attr.lstrip("-").isdigit():
+            obj[int(target_attr)] = value
+        else:
+            setattr(obj, target_attr, value)
+
+
+def _is_complex_path(path):
+    """Check if a path contains complex access patterns (brackets or parentheses)."""
+    return "[" in path or "(" in path
+
+
+def _set_complex_path_value(obj, path, value):
+    """Set a value using a complex path by navigating to the parent and setting the final element."""
+    # Parse the path to find the parent path and final accessor
+    parent_obj = obj
+
+    # Find the last bracket or the final attribute
+    last_bracket = path.rfind("[")
+    last_paren = path.rfind("(")
+
+    if last_bracket > last_paren:
+        # Last accessor is a bracket
+        bracket_end = _find_matching_bracket(path, last_bracket)
+        parent_path = path[:last_bracket]
+        bracket_content = path[last_bracket + 1 : bracket_end]
+
+        if parent_path:
+            parent_obj = _process_path_segment(obj, parent_path)
+
+        # Set the value using bracket access
+        if bracket_content.startswith("'") and bracket_content.endswith("'"):
+            key = bracket_content[1:-1]
+            parent_obj[key] = value
+        elif bracket_content.startswith('"') and bracket_content.endswith('"'):
+            key = bracket_content[1:-1]
+            parent_obj[key] = value
+        elif bracket_content.lstrip("-").isdigit():
+            index = int(bracket_content)
+            parent_obj[index] = value
+        else:
+            try:
+                key = eval(bracket_content)
+                parent_obj[key] = value
+            except:
+                parent_obj[bracket_content] = value
+    else:
+        # No brackets, treat as simple attribute
+        if path.lstrip("-").isdigit():
+            obj[int(path)] = value
+        else:
+            setattr(obj, path, value)
+
+
+def update_scheduled_params(obj, scheduler_dict, step, split_char="@"):
+    scheduled_params_dict = {}
+    for target, cfg in scheduler_dict.items():
+        sch_type = cfg["type"]
+        val_type = cfg.get("val_type", "float")
+        target_attr = target
+        target_obj = obj
+        if split_char in target:
+            target_obj_str, target_attr = target.rsplit(split_char, 1)
+            target_obj = _navigate_object_path(obj, target_obj_str, split_char)
+        if sch_type == "linear":
+            i = len(cfg["seg_vals"]) - 1
+            while step < cfg["seg_steps"][i]:
+                i -= 1
+            if i == len(cfg["seg_vals"]) - 1:
+                val = cfg["seg_vals"][i]
+            else:
+                t = (step - cfg["seg_steps"][i]) / (cfg["seg_steps"][i + 1] - cfg["seg_steps"][i])
+                t = max(0.0, min(1.0, t))
+                val = (1.0 - t) * cfg["seg_vals"][i] + t * cfg["seg_vals"][i + 1]
+        elif sch_type == "segment":
+            i = len(cfg["seg_vals"]) - 1
+            while step < cfg["seg_steps"][i]:
+                i -= 1
+            val = cfg["seg_vals"][i]
+
+        val = eval(val_type)(val)
+
+        if type(val) is DictConfig or type(val) is dict:
+            # Handle complex path for dict/config access
+            tmp_obj = _get_final_target(target_obj, target_attr)
+
+            if cfg.get("overwrite_dict", False):
+                _set_final_target(target_obj, target_attr, val)
+            else:
+                for k, v in val.items():
+                    if type(tmp_obj) is dict:
+                        tmp_obj[k] = v
+                    else:
+                        setattr(tmp_obj, k, v)
+        else:
+            # Handle complex path for direct value assignment
+            _set_final_target(target_obj, target_attr, val)
+
+        scheduled_params_dict[target] = val
+
+        if "trigger_func" in cfg and step == cfg["seg_steps"][i]:
+            target_func = cfg["trigger_func"]
+            print(f"Triggering function: {target_func}")
+            if split_char in target_func:
+                target_obj_str, target_func_name = target_func.rsplit(split_char, 1)
+                target_obj = _navigate_object_path(obj, target_obj_str, split_char)
+            else:
+                target_obj = obj
+                target_func_name = target_func
+            getattr(target_obj, target_func_name)()
+
+    return scheduled_params_dict
+
+
+class WarmupCosineScheduler(_LRScheduler):
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        num_warmup_steps: int,
+        num_training_steps: int,
+        final_lr: float = 0.0,
+        last_epoch: int = -1,
+    ):
+        self.num_warmup_steps = num_warmup_steps
+        self.num_training_steps = num_training_steps
+        self.final_lr = final_lr
+        super(WarmupCosineScheduler, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        current_step = self.last_epoch
+        if current_step < self.num_warmup_steps:
+            return [
+                base_lr * float(current_step) / float(max(1, self.num_warmup_steps))
+                for base_lr in self.base_lrs
+            ]
+        else:
+            progress = float(current_step - self.num_warmup_steps) / float(
+                max(1, self.num_training_steps - self.num_warmup_steps)
+            )
+            cosine_decay = 0.5 * (1.0 + math.cos(math.pi * min(progress, 1.0)))
+            return [
+                self.final_lr + (base_lr - self.final_lr) * cosine_decay
+                for base_lr in self.base_lrs
+            ]
+
+
+if __name__ == "__main__":
+    # Test the complex path navigation
+    class MockEventManager:
+        def __init__(self):
+            self.configs = {
+                "push_robot": {"params": {"velocity_range": {"x": [1.0, 2.0], "y": [0.5, 1.5]}}}
+            }
+
+        def get_term_cfg(self, term_name):
+            return self.configs[term_name]
+
+    class MockEnv:
+        def __init__(self):
+            self.event_manager = MockEventManager()
+
+    class MockSimulator:
+        def __init__(self):
+            self.env = MockEnv()
+
+    # Test complex path navigation
+    mock_obj = MockSimulator()
+
+    # Test the path: env@event_manager@get_term_cfg('push_robot')@params@velocity_range@x@0
+    test_path = "env@event_manager@get_term_cfg('push_robot')['params']['velocity_range']['x'][0]"
+
+    # Create a simple scheduler config to test
+    scheduler_config = {
+        test_path: {"type": "linear", "seg_steps": [0, 100], "seg_vals": [5.0, 10.0]}
+    }
+
+    # Test the function
+    print("Testing complex path navigation...")
+    print(
+        f"Original value: {mock_obj.env.event_manager.get_term_cfg('push_robot')['params']['velocity_range']['x'][0]}"
+    )
+
+    result = update_scheduled_params(mock_obj, scheduler_config, 50)
+    print(
+        f"Updated value: {mock_obj.env.event_manager.get_term_cfg('push_robot')['params']['velocity_range']['x'][0]}"
+    )
+    print(f"Scheduler result: {result}")
+
+    # Test with step that triggers second segment
+    result2 = update_scheduled_params(mock_obj, scheduler_config, 150)
+    print(
+        f"Updated value (step 150): {mock_obj.env.event_manager.get_term_cfg('push_robot')['params']['velocity_range']['x'][0]}"
+    )
+    print(f"Scheduler result: {result2}")
+
+    print("\nOriginal learning rate scheduler test:")
+
+    class YourModel(torch.nn.Module):
+        def __init__(self):
+            super(YourModel, self).__init__()
+            self.fc = torch.nn.Linear(10, 1)
+
+        def forward(self, x):
+            return self.fc(x)
+
+    model = YourModel()
+    optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)
+
+    num_warmup_steps = 1000
+    num_training_steps = 10000
+    final_lr = 0.0001
+
+    scheduler = WarmupCosineScheduler(optimizer, num_warmup_steps, num_training_steps, final_lr)
+
+    lrs = []
+    for step in range(num_training_steps):
+        scheduler.step()
+        lrs.append(scheduler.get_lr()[0])
+
+    # Plotting the learning rate vs training steps
+    import matplotlib.pyplot as plt
+
+    plt.plot(range(num_training_steps), lrs)
+    plt.xlabel("Training Steps")
+    plt.ylabel("Learning Rate")
+    plt.title("Learning Rate vs Training Steps")
+    # plt.show()
+    plt.savefig("out/lr_vs_steps.png")
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/README.md b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..ec1770250b6f1bca8fff3bf0c06ab95e580a37d7
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/README.md
@@ -0,0 +1,3 @@
+# README
+
+Contents of this folder are modified from HuMoR repository.
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/__init__.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..26f4858ef74a17283cbd0224d27876c35538a6b2
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/__init__.py
@@ -0,0 +1,5 @@
+"""Body model wrappers for SMPL, SMPLH, and SMPLX."""
+
+from .body_model import BodyModel
+from .body_model_smplh import BodyModelSMPLH
+from .body_model_smplx import BodyModelSMPLX
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/body_model.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/body_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..cd53bf58a96c908d9b882125abb9cb338301d1be
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/body_model.py
@@ -0,0 +1,136 @@
+"""Generic BodyModel wrapper around SMPL/SMPLH/SMPLX with optional vertex selection."""
+
+import numpy as np
+import torch
+import torch.nn as nn
+from smplx import SMPL, SMPLH, SMPLX
+from smplx.utils import Struct
+from smplx.vertex_ids import vertex_ids
+
+
+class BodyModel(nn.Module):
+    """
+    Wrapper around SMPLX body model class.
+    modified by Zehong Shen
+    """
+
+    def __init__(self, bm_path, num_betas=16, use_vtx_selector=False, model_type="smplh"):
+        super().__init__()
+        """
+        Creates the body model object at the given path.
+
+        :param bm_path: path to the body model pkl file
+        :param model_type: one of [smpl, smplh, smplx]
+        :param use_vtx_selector: if true, returns additional vertices as joints that correspond to OpenPose joints
+        """
+        self.use_vtx_selector = use_vtx_selector
+        cur_vertex_ids = None
+        if self.use_vtx_selector:
+            cur_vertex_ids = vertex_ids[model_type]
+        data_struct = None
+        if ".npz" in bm_path:
+            # smplx does not support .npz by default, so have to load in manually
+            smpl_dict = np.load(bm_path, encoding="latin1")
+            data_struct = Struct(**smpl_dict)
+            # print(smpl_dict.files)
+            if model_type == "smplh":
+                data_struct.hands_componentsl = np.zeros((0))
+                data_struct.hands_componentsr = np.zeros((0))
+                data_struct.hands_meanl = np.zeros((15 * 3))
+                data_struct.hands_meanr = np.zeros((15 * 3))
+                V, D, B = data_struct.shapedirs.shape
+                data_struct.shapedirs = np.concatenate(
+                    [data_struct.shapedirs, np.zeros((V, D, SMPL.SHAPE_SPACE_DIM - B))],
+                    axis=-1,
+                )  # super hacky way to let smplh use 16-size beta
+        kwargs = {
+            "model_type": model_type,
+            "data_struct": data_struct,
+            "num_betas": num_betas,
+            "vertex_ids": cur_vertex_ids,
+            "use_pca": False,
+            "flat_hand_mean": True,
+            # - enable variable batchsize, since we don't need module variable - #
+            "create_body_pose": False,
+            "create_betas": False,
+            "create_global_orient": False,
+            "create_transl": False,
+            "create_left_hand_pose": False,
+            "create_right_hand_pose": False,
+        }
+        assert model_type in ["smpl", "smplh", "smplx"]
+        if model_type == "smpl":
+            self.bm = SMPL(bm_path, **kwargs)
+            self.num_joints = SMPL.NUM_JOINTS
+        elif model_type == "smplh":
+            self.bm = SMPLH(bm_path, **kwargs)
+            self.num_joints = SMPLH.NUM_JOINTS
+        elif model_type == "smplx":
+            self.bm = SMPLX(bm_path, **kwargs)
+            self.num_joints = SMPLX.NUM_JOINTS
+
+        self.model_type = model_type
+
+    def forward(
+        self,
+        root_orient=None,
+        pose_body=None,
+        pose_hand=None,
+        pose_jaw=None,
+        pose_eye=None,
+        betas=None,
+        trans=None,
+        dmpls=None,
+        expression=None,
+        return_dict=False,
+        **kwargs,
+    ):
+        """
+        Note dmpls are not supported.
+        """
+        assert dmpls is None
+        B = pose_body.shape[0]
+        if pose_hand is None:
+            pose_hand = torch.zeros((B, 2 * SMPLH.NUM_HAND_JOINTS * 3), device=pose_body.device)
+        if len(betas.shape) == 1:
+            betas = betas.reshape((1, -1)).expand(B, -1)
+
+        out_obj = self.bm(
+            betas=betas,
+            global_orient=root_orient,
+            body_pose=pose_body,
+            left_hand_pose=pose_hand[:, : (SMPLH.NUM_HAND_JOINTS * 3)],
+            right_hand_pose=pose_hand[:, (SMPLH.NUM_HAND_JOINTS * 3) :],
+            transl=trans,
+            expression=expression,
+            jaw_pose=pose_jaw,
+            leye_pose=None if pose_eye is None else pose_eye[:, :3],
+            reye_pose=None if pose_eye is None else pose_eye[:, 3:],
+            return_full_pose=True,
+            **kwargs,
+        )
+
+        out = {
+            "v": out_obj.vertices,
+            "f": self.bm.faces_tensor,
+            "Jtr": out_obj.joints,
+        }
+
+        if not self.use_vtx_selector:
+            # don't need extra joints
+            out["Jtr"] = out["Jtr"][:, : self.num_joints + 1]  # add one for the root
+
+        if not return_dict:
+            out = Struct(**out)
+
+        return out
+
+    def forward_motion(self, **kwargs):
+        B, W, _ = kwargs["pose_body"].shape
+        kwargs = {k: v.reshape(B * W, v.shape[-1]) for k, v in kwargs.items()}
+
+        smpl_opt = self.forward(**kwargs)
+        smpl_opt.v = smpl_opt.v.reshape(B, W, -1, 3)
+        smpl_opt.Jtr = smpl_opt.Jtr.reshape(B, W, -1, 3)
+
+        return smpl_opt
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/body_model_smplh.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/body_model_smplh.py
new file mode 100644
index 0000000000000000000000000000000000000000..85dab2e8089182a22b6a6e43983dc3b1bb9ec163
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/body_model_smplh.py
@@ -0,0 +1,110 @@
+"""SMPLH body model wrapper for batch inference with explicit pose parameters."""
+
+import smplx
+import torch
+import torch.nn as nn
+
+kwargs_disable_member_var = {
+    "create_body_pose": False,
+    "create_betas": False,
+    "create_global_orient": False,
+    "create_transl": False,
+    "create_left_hand_pose": False,
+    "create_right_hand_pose": False,
+}
+
+
+class BodyModelSMPLH(nn.Module):
+    """Support Batch inference"""
+
+    def __init__(self, model_path, **kwargs):
+        super().__init__()
+        # enable flexible batchsize, handle missing variable at forward()
+        kwargs.update(kwargs_disable_member_var)
+        self.bm = smplx.create(model_path=model_path, **kwargs)
+        self.faces = self.bm.faces
+        self.is_smpl = kwargs.get("model_type", "smpl") == "smpl"
+        if not self.is_smpl:
+            self.hand_pose_dim = (
+                self.bm.num_pca_comps if self.bm.use_pca else 3 * self.bm.NUM_HAND_JOINTS
+            )
+
+        # For fast computing of skeleton under beta
+        shapedirs = self.bm.shapedirs  # (V, 3, 10)
+        J_regressor = self.bm.J_regressor[:22, :]  # (22, V)
+        v_template = self.bm.v_template  # (V, 3)
+        J_template = J_regressor @ v_template  # (22, 3)
+        J_shapedirs = torch.einsum("jv, vcd -> jcd", J_regressor, shapedirs)  # (22, 3, 10)
+        self.register_buffer("J_template", J_template, False)
+        self.register_buffer("J_shapedirs", J_shapedirs, False)
+
+    def forward(
+        self,
+        betas=None,
+        global_orient=None,
+        transl=None,
+        body_pose=None,
+        left_hand_pose=None,
+        right_hand_pose=None,
+        **kwargs,
+    ):
+        device, dtype = self.bm.shapedirs.device, self.bm.shapedirs.dtype
+
+        model_vars = [
+            betas,
+            global_orient,
+            body_pose,
+            transl,
+            left_hand_pose,
+            right_hand_pose,
+        ]
+        batch_size = 1
+        for var in model_vars:
+            if var is None:
+                continue
+            batch_size = max(batch_size, len(var))
+
+        if global_orient is None:
+            global_orient = torch.zeros([batch_size, 3], dtype=dtype, device=device)
+        if body_pose is None:
+            body_pose = (
+                torch.zeros(3 * self.bm.NUM_BODY_JOINTS, device=device, dtype=dtype)[None]
+                .expand(batch_size, -1)
+                .contiguous()
+            )
+        if not self.is_smpl:
+            if left_hand_pose is None:
+                left_hand_pose = (
+                    torch.zeros(self.hand_pose_dim, device=device, dtype=dtype)[None]
+                    .expand(batch_size, -1)
+                    .contiguous()
+                )
+            if right_hand_pose is None:
+                right_hand_pose = (
+                    torch.zeros(self.hand_pose_dim, device=device, dtype=dtype)[None]
+                    .expand(batch_size, -1)
+                    .contiguous()
+                )
+        if betas is None:
+            betas = torch.zeros([batch_size, self.bm.num_betas], dtype=dtype, device=device)
+        if transl is None:
+            transl = torch.zeros([batch_size, 3], dtype=dtype, device=device)
+
+        bm_out = self.bm(
+            betas=betas,
+            global_orient=global_orient,
+            body_pose=body_pose,
+            left_hand_pose=left_hand_pose,
+            right_hand_pose=right_hand_pose,
+            transl=transl,
+            **kwargs,
+        )
+
+        return bm_out
+
+    def get_skeleton(self, betas):
+        """betas: (*, 10) -> skeleton_beta: (*, 22, 3)"""
+        skeleton_beta = self.J_template + torch.einsum(
+            "...d, jcd -> ...jc", betas, self.J_shapedirs
+        )  # (22, 3)
+        return skeleton_beta
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/body_model_smplx.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/body_model_smplx.py
new file mode 100644
index 0000000000000000000000000000000000000000..50faf6d4ac57b2d2a79b8cb6d84df76c42d947ec
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/body_model_smplx.py
@@ -0,0 +1,143 @@
+"""SMPLX body model wrapper for batch inference with face/hand/body pose parameters."""
+
+import smplx
+import torch
+import torch.nn as nn
+
+kwargs_disable_member_var = {
+    "create_body_pose": False,
+    "create_betas": False,
+    "create_global_orient": False,
+    "create_transl": False,
+    "create_left_hand_pose": False,
+    "create_right_hand_pose": False,
+    "create_expression": False,
+    "create_jaw_pose": False,
+    "create_leye_pose": False,
+    "create_reye_pose": False,
+}
+
+
+class BodyModelSMPLX(nn.Module):
+    """Support Batch inference"""
+
+    def __init__(self, model_path, **kwargs):
+        super().__init__()
+        # enable flexible batchsize, handle missing variable at forward()
+        kwargs.update(kwargs_disable_member_var)
+        self.bm = smplx.create(model_path=model_path, **kwargs)
+        self.faces = self.bm.faces
+        self.hand_pose_dim = (
+            self.bm.num_pca_comps if self.bm.use_pca else 3 * self.bm.NUM_HAND_JOINTS
+        )
+
+        # For fast computing of skeleton under beta
+        shapedirs = self.bm.shapedirs  # (V, 3, 10)
+        J_regressor = self.bm.J_regressor[:22, :]  # (22, V)
+        v_template = self.bm.v_template  # (V, 3)
+        J_template = J_regressor @ v_template  # (22, 3)
+        J_shapedirs = torch.einsum("jv, vcd -> jcd", J_regressor, shapedirs)  # (22, 3, 10)
+        self.register_buffer("J_template", J_template, False)
+        self.register_buffer("J_shapedirs", J_shapedirs, False)
+
+    def forward(
+        self,
+        betas=None,
+        global_orient=None,
+        transl=None,
+        body_pose=None,
+        left_hand_pose=None,
+        right_hand_pose=None,
+        expression=None,
+        jaw_pose=None,
+        leye_pose=None,
+        reye_pose=None,
+        **kwargs,
+    ):
+        device, dtype = self.bm.shapedirs.device, self.bm.shapedirs.dtype
+
+        model_vars = [
+            betas,
+            global_orient,
+            body_pose,
+            transl,
+            expression,
+            left_hand_pose,
+            right_hand_pose,
+            jaw_pose,
+            leye_pose,
+            reye_pose,
+        ]
+        batch_size = 1
+        for var in model_vars:
+            if var is None:
+                continue
+            batch_size = max(batch_size, len(var))
+
+        if global_orient is None:
+            global_orient = torch.zeros([batch_size, 3], dtype=dtype, device=device)
+        if body_pose is None:
+            body_pose = (
+                torch.zeros(3 * self.bm.NUM_BODY_JOINTS, device=device, dtype=dtype)[None]
+                .expand(batch_size, -1)
+                .contiguous()
+            )
+        if left_hand_pose is None:
+            left_hand_pose = (
+                torch.zeros(self.hand_pose_dim, device=device, dtype=dtype)[None]
+                .expand(batch_size, -1)
+                .contiguous()
+            )
+        if right_hand_pose is None:
+            right_hand_pose = (
+                torch.zeros(self.hand_pose_dim, device=device, dtype=dtype)[None]
+                .expand(batch_size, -1)
+                .contiguous()
+            )
+        if jaw_pose is None:
+            jaw_pose = torch.zeros([batch_size, 3], dtype=dtype, device=device)
+        if leye_pose is None:
+            leye_pose = torch.zeros([batch_size, 3], dtype=dtype, device=device)
+        if reye_pose is None:
+            reye_pose = torch.zeros([batch_size, 3], dtype=dtype, device=device)
+        if expression is None:
+            expression = torch.zeros(
+                [batch_size, self.bm.num_expression_coeffs], dtype=dtype, device=device
+            )
+        if betas is None:
+            betas = torch.zeros([batch_size, self.bm.num_betas], dtype=dtype, device=device)
+        if transl is None:
+            transl = torch.zeros([batch_size, 3], dtype=dtype, device=device)
+
+        bm_out = self.bm(
+            betas=betas,
+            global_orient=global_orient,
+            body_pose=body_pose,
+            left_hand_pose=left_hand_pose,
+            right_hand_pose=right_hand_pose,
+            transl=transl,
+            expression=expression,
+            jaw_pose=jaw_pose,
+            leye_pose=leye_pose,
+            reye_pose=reye_pose,
+            **kwargs,
+        )
+
+        return bm_out
+
+    def get_skeleton(self, betas):
+        """betas: (*, 10) -> skeleton_beta: (*, 22, 3)"""
+        skeleton_beta = self.J_template + torch.einsum(
+            "...d, jcd -> ...jc", betas, self.J_shapedirs
+        )  # (22, 3)
+        return skeleton_beta
+
+    def forward_bfc(self, **kwargs):
+        """Wrap (B, F, C) to (B*F, C) and unwrap (B*F, C) to (B, F, C)"""
+        for k in kwargs:
+            assert len(kwargs[k].shape) == 3
+        B, F = kwargs["body_pose"].shape[:2]
+        smplx_out = self.forward(**{k: v.reshape(B * F, -1) for k, v in kwargs.items()})
+        smplx_out.vertices = smplx_out.vertices.reshape(B, F, -1, 3)
+        smplx_out.joints = smplx_out.joints.reshape(B, F, -1, 3)
+        return smplx_out
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/min_lbs.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/min_lbs.py
new file mode 100644
index 0000000000000000000000000000000000000000..5e3cc843e1e0078b302e71d1b9e46b9b06b727f8
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/min_lbs.py
@@ -0,0 +1,126 @@
+"""Minimal Linear Blend Skinning (LBS) for sparse sensor-point vertices on SMPLH."""
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from smplx.utils import Struct, to_np, to_tensor
+
+from hmr4d.utils.smplx_utils import forward_kinematics_motion
+from motiondiff.models.mdm.rotation_conversions import axis_angle_to_matrix
+
+
+class MinimalLBS(nn.Module):
+    def __init__(self, sp_ids, bm_dir="models/smplh", num_betas=16, model_type="smplh", **kwargs):
+        super().__init__()
+        self.num_betas = num_betas
+        self.sensor_point_vid = torch.tensor(sp_ids)
+
+        # load struct data on predefined sensor-point
+        self.load_struct_on_sp(f"{bm_dir}/male/model.npz", prefix="male")
+        self.load_struct_on_sp(f"{bm_dir}/female/model.npz", prefix="female")
+
+    def load_struct_on_sp(self, bm_path, prefix="m"):
+        """
+        Load 4 weights from body-model-struct.
+        Keep the sensor points only. Use prefix to label different bm.
+        """
+        num_betas = self.num_betas
+        sp_vid = self.sensor_point_vid
+        # load data
+        data_struct = Struct(**np.load(bm_path, encoding="latin1"))
+
+        # v-template
+        v_template = to_tensor(to_np(data_struct.v_template))  # (V, 3)
+        v_template_sp = v_template[sp_vid]  # (N, 3)
+        self.register_buffer(f"{prefix}_v_template_sp", v_template_sp, False)
+
+        # shapedirs
+        shapedirs = to_tensor(to_np(data_struct.shapedirs[:, :, :num_betas]))  # (V, 3, NB)
+        shapedirs_sp = shapedirs[sp_vid]
+        self.register_buffer(f"{prefix}_shapedirs_sp", shapedirs_sp, False)
+
+        # posedirs
+        posedirs = to_tensor(to_np(data_struct.posedirs))  # (V, 3, 51*9)
+        posedirs_sp = posedirs[sp_vid]
+        posedirs_sp = posedirs_sp.reshape(len(sp_vid) * 3, -1).T  # (51*9, N*3)
+        self.register_buffer(f"{prefix}_posedirs_sp", posedirs_sp, False)
+
+        # lbs_weights
+        lbs_weights = to_tensor(to_np(data_struct.weights))  # (V, J+1)
+        lbs_weights_sp = lbs_weights[sp_vid]
+        self.register_buffer(f"{prefix}_lbs_weights_sp", lbs_weights_sp, False)
+
+    def forward(
+        self,
+        root_orient=None,
+        pose_body=None,
+        trans=None,
+        betas=None,
+        A=None,
+        recompute_A=False,
+        genders=None,
+        joints_zero=None,
+    ):
+        """
+        Args:
+            root_orient, Optional: (B, T, 3)
+            pose_body: (B, T, J*3)
+            trans: (B, T, 3)
+            betas: (B, T, 16)
+            A, Optional: (B, T, J+1, 4, 4)
+            recompute_A: if True, root_orient should be given, otherwise use A
+            genders, List: ['male', 'female', ...]
+            joints_zero: (B, J+1, 3), required when recompute_A is True
+        Returns:
+            sensor_verts: (B, T, N, 3)
+        """
+        B, T = pose_body.shape[:2]
+
+        v_template = torch.stack(
+            [getattr(self, f"{g}_v_template_sp") for g in genders]
+        )  # (B, N, 3)
+        shapedirs = torch.stack(
+            [getattr(self, f"{g}_shapedirs_sp") for g in genders]
+        )  # (B, N, 3, NB)
+        posedirs = torch.stack(
+            [getattr(self, f"{g}_posedirs_sp") for g in genders]
+        )  # (B, 51*9, N*3)
+        lbs_weights = torch.stack(
+            [getattr(self, f"{g}_lbs_weights_sp") for g in genders]
+        )  # (B, N, J+1)
+
+        # ===== LBS, handle T ===== #
+        # 2. Add shape contribution
+        if betas.shape[1] == 1:
+            betas = betas.expand(-1, T, -1)
+        blend_shape = torch.einsum("btl,bmkl->btmk", [betas, shapedirs])
+        v_shaped = v_template[:, None] + blend_shape
+
+        # 3. Add pose blend shapes
+        ident = torch.eye(3).to(pose_body)
+        aa = pose_body.reshape(B, T, -1, 3)
+        R = axis_angle_to_matrix(aa)
+        pose_feature = (R - ident).view(B, T, -1)
+        dim_pf = pose_feature.shape[-1]
+        # (B, T, P) @ (B, P, N*3) -> (B, T, N, 3)
+        pose_offsets = torch.matmul(pose_feature, posedirs[:, :dim_pf]).view(B, T, -1, 3)
+        v_posed = pose_offsets + v_shaped
+
+        # 4. Compute A
+        if recompute_A:
+            _, _, A = forward_kinematics_motion(root_orient, pose_body, trans, joints_zero)
+
+        # 5. Skinning
+        W = lbs_weights
+        # (B, 1, N, J+1)) @ (B, T, J+1, 16)
+        num_joints = A.shape[-3]  # 22
+        Ts = torch.matmul(W[:, None, :, :num_joints], A.view(B, T, num_joints, 16))
+        Ts = Ts.view(B, T, -1, 4, 4)  # (B, T, N, 4, 4)
+        v_posed_homo = F.pad(v_posed, (0, 1), value=1)  # (B, T, N, 4)
+        v_homo = torch.matmul(Ts, torch.unsqueeze(v_posed_homo, dim=-1))
+
+        # 6. translate
+        sensor_verts = v_homo[:, :, :, :3, 0] + trans[:, :, None]
+
+        return sensor_verts
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/rotation_conversions.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/rotation_conversions.py
new file mode 100644
index 0000000000000000000000000000000000000000..2aa539757fb42a19e0da6aaff1628eef1c3d8805
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/rotation_conversions.py
@@ -0,0 +1,538 @@
+# This file contains code derived from PyTorch3D (via https://github.com/Mathux/ACTOR),
+# originally Copyright (c) Facebook, Inc. and its affiliates, licensed under BSD.
+# See https://github.com/facebookresearch/pytorch3d/blob/main/LICENSE for the original license.
+
+"""Rotation conversions for body_model subpackage (quaternion, matrix, axis-angle).
+
+The transformation matrices returned from the functions in this file assume
+the points on which the transformation will be applied are column vectors.
+i.e. the R matrix is structured as
+
+    R = [
+            [Rxx, Rxy, Rxz],
+            [Ryx, Ryy, Ryz],
+            [Rzx, Rzy, Rzz],
+        ]  # (3, 3)
+
+This matrix can be applied to column vectors by post multiplication
+by the points e.g.
+
+    points = [[0], [1], [2]]  # (3 x 1) xyz coordinates of a point
+    transformed_points = R * points
+
+To apply the same matrix to points which are row vectors, the R matrix
+can be transposed and pre multiplied by the points:
+
+e.g.
+    points = [[0, 1, 2]]  # (1 x 3) xyz coordinates of a point
+    transformed_points = points * R.transpose(1, 0)
+"""
+
+import functools
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+
+
+def quaternion_to_matrix(quaternions):
+    """
+    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))
+
+
+def _copysign(a, b):
+    """
+    Return a tensor where each element has the absolute value taken from the,
+    corresponding element of a, with sign taken from the corresponding
+    element of b. This is like the standard copysign floating-point operation,
+    but is not careful about negative 0 and NaN.
+
+    Args:
+        a: source tensor.
+        b: tensor whose signs will be used, of the same shape as a.
+
+    Returns:
+        Tensor of the same shape as a with the signs of b.
+    """
+    signs_differ = (a < 0) != (b < 0)
+    return torch.where(signs_differ, -a, a)
+
+
+def _sqrt_positive_part(x):
+    """
+    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 matrix_to_quaternion(matrix):
+    """
+    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 f{matrix.shape}.")
+    m00 = matrix[..., 0, 0]
+    m11 = matrix[..., 1, 1]
+    m22 = matrix[..., 2, 2]
+    o0 = 0.5 * _sqrt_positive_part(1 + m00 + m11 + m22)
+    x = 0.5 * _sqrt_positive_part(1 + m00 - m11 - m22)
+    y = 0.5 * _sqrt_positive_part(1 - m00 + m11 - m22)
+    z = 0.5 * _sqrt_positive_part(1 - m00 - m11 + m22)
+    o1 = _copysign(x, matrix[..., 2, 1] - matrix[..., 1, 2])
+    o2 = _copysign(y, matrix[..., 0, 2] - matrix[..., 2, 0])
+    o3 = _copysign(z, matrix[..., 1, 0] - matrix[..., 0, 1])
+    return torch.stack((o0, o1, o2, o3), -1)
+
+
+def _axis_angle_rotation(axis: str, angle):
+    """
+    Return the rotation matrices for one of the rotations about an axis
+    of which Euler angles describe, for each value of the angle given.
+
+    Args:
+        axis: Axis label "X" or "Y or "Z".
+        angle: any shape tensor of Euler angles in radians
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+
+    cos = torch.cos(angle)
+    sin = torch.sin(angle)
+    one = torch.ones_like(angle)
+    zero = torch.zeros_like(angle)
+
+    if axis == "X":
+        R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos)
+    if axis == "Y":
+        R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos)
+    if axis == "Z":
+        R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one)
+
+    return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3))
+
+
+def euler_angles_to_matrix(euler_angles, convention: str):
+    """
+    Convert rotations given as Euler angles in radians to rotation matrices.
+
+    Args:
+        euler_angles: Euler angles in radians as tensor of shape (..., 3).
+        convention: Convention string of three uppercase letters from
+            {"X", "Y", and "Z"}.
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3:
+        raise ValueError("Invalid input euler angles.")
+    if len(convention) != 3:
+        raise ValueError("Convention must have 3 letters.")
+    if convention[1] in (convention[0], convention[2]):
+        raise ValueError(f"Invalid convention {convention}.")
+    for letter in convention:
+        if letter not in ("X", "Y", "Z"):
+            raise ValueError(f"Invalid letter {letter} in convention string.")
+    matrices = map(_axis_angle_rotation, convention, torch.unbind(euler_angles, -1))
+    return functools.reduce(torch.matmul, matrices)
+
+
+def _angle_from_tan(axis: str, other_axis: str, data, horizontal: bool, tait_bryan: bool):
+    """
+    Extract the first or third Euler angle from the two members of
+    the matrix which are positive constant times its sine and cosine.
+
+    Args:
+        axis: Axis label "X" or "Y or "Z" for the angle we are finding.
+        other_axis: Axis label "X" or "Y or "Z" for the middle axis in the
+            convention.
+        data: Rotation matrices as tensor of shape (..., 3, 3).
+        horizontal: Whether we are looking for the angle for the third axis,
+            which means the relevant entries are in the same row of the
+            rotation matrix. If not, they are in the same column.
+        tait_bryan: Whether the first and third axes in the convention differ.
+
+    Returns:
+        Euler Angles in radians for each matrix in dataset as a tensor
+        of shape (...).
+    """
+
+    i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis]
+    if horizontal:
+        i2, i1 = i1, i2
+    even = (axis + other_axis) in ["XY", "YZ", "ZX"]
+    if horizontal == even:
+        return torch.atan2(data[..., i1], data[..., i2])
+    if tait_bryan:
+        return torch.atan2(-data[..., i2], data[..., i1])
+    return torch.atan2(data[..., i2], -data[..., i1])
+
+
+def _index_from_letter(letter: str):
+    if letter == "X":
+        return 0
+    if letter == "Y":
+        return 1
+    if letter == "Z":
+        return 2
+
+
+def matrix_to_euler_angles(matrix, convention: str):
+    """
+    Convert rotations given as rotation matrices to Euler angles in radians.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+        convention: Convention string of three uppercase letters.
+
+    Returns:
+        Euler angles in radians as tensor of shape (..., 3).
+    """
+    if len(convention) != 3:
+        raise ValueError("Convention must have 3 letters.")
+    if convention[1] in (convention[0], convention[2]):
+        raise ValueError(f"Invalid convention {convention}.")
+    for letter in convention:
+        if letter not in ("X", "Y", "Z"):
+            raise ValueError(f"Invalid letter {letter} in convention string.")
+    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
+        raise ValueError(f"Invalid rotation matrix  shape f{matrix.shape}.")
+    i0 = _index_from_letter(convention[0])
+    i2 = _index_from_letter(convention[2])
+    tait_bryan = i0 != i2
+    if tait_bryan:
+        central_angle = torch.asin(matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0))
+    else:
+        central_angle = torch.acos(matrix[..., i0, i0])
+
+    o = (
+        _angle_from_tan(convention[0], convention[1], matrix[..., i2], False, tait_bryan),
+        central_angle,
+        _angle_from_tan(convention[2], convention[1], matrix[..., i0, :], True, tait_bryan),
+    )
+    return torch.stack(o, -1)
+
+
+def random_quaternions(
+    n: int, dtype: Optional[torch.dtype] = None, device=None, requires_grad=False
+):
+    """
+    Generate random quaternions representing rotations,
+    i.e. versors with nonnegative real part.
+
+    Args:
+        n: Number of quaternions in a batch to return.
+        dtype: Type to return.
+        device: Desired device of returned tensor. Default:
+            uses the current device for the default tensor type.
+        requires_grad: Whether the resulting tensor should have the gradient
+            flag set.
+
+    Returns:
+        Quaternions as tensor of shape (N, 4).
+    """
+    o = torch.randn((n, 4), dtype=dtype, device=device, requires_grad=requires_grad)
+    s = (o * o).sum(1)
+    o = o / _copysign(torch.sqrt(s), o[:, 0])[:, None]
+    return o
+
+
+def random_rotations(n: int, dtype: Optional[torch.dtype] = None, device=None, requires_grad=False):
+    """
+    Generate random rotations as 3x3 rotation matrices.
+
+    Args:
+        n: Number of rotation matrices in a batch to return.
+        dtype: Type to return.
+        device: Device of returned tensor. Default: if None,
+            uses the current device for the default tensor type.
+        requires_grad: Whether the resulting tensor should have the gradient
+            flag set.
+
+    Returns:
+        Rotation matrices as tensor of shape (n, 3, 3).
+    """
+    quaternions = random_quaternions(n, dtype=dtype, device=device, requires_grad=requires_grad)
+    return quaternion_to_matrix(quaternions)
+
+
+def random_rotation(dtype: Optional[torch.dtype] = None, device=None, requires_grad=False):
+    """
+    Generate a single random 3x3 rotation matrix.
+
+    Args:
+        dtype: Type to return
+        device: Device of returned tensor. Default: if None,
+            uses the current device for the default tensor type
+        requires_grad: Whether the resulting tensor should have the gradient
+            flag set
+
+    Returns:
+        Rotation matrix as tensor of shape (3, 3).
+    """
+    return random_rotations(1, dtype, device, requires_grad)[0]
+
+
+def standardize_quaternion(quaternions):
+    """
+    Convert a unit quaternion to a standard form: one in which the real
+    part is non negative.
+
+    Args:
+        quaternions: Quaternions with real part first,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Standardized quaternions as tensor of shape (..., 4).
+    """
+    return torch.where(quaternions[..., 0:1] < 0, -quaternions, quaternions)
+
+
+def quaternion_raw_multiply(a, b):
+    """
+    Multiply two quaternions.
+    Usual torch rules for broadcasting apply.
+
+    Args:
+        a: Quaternions as tensor of shape (..., 4), real part first.
+        b: Quaternions as tensor of shape (..., 4), real part first.
+
+    Returns:
+        The product of a and b, a tensor of quaternions shape (..., 4).
+    """
+    aw, ax, ay, az = torch.unbind(a, -1)
+    bw, bx, by, bz = torch.unbind(b, -1)
+    ow = aw * bw - ax * bx - ay * by - az * bz
+    ox = aw * bx + ax * bw + ay * bz - az * by
+    oy = aw * by - ax * bz + ay * bw + az * bx
+    oz = aw * bz + ax * by - ay * bx + az * bw
+    return torch.stack((ow, ox, oy, oz), -1)
+
+
+def quaternion_multiply(a, b):
+    """
+    Multiply two quaternions representing rotations, returning the quaternion
+    representing their composition, i.e. the versor with nonnegative real part.
+    Usual torch rules for broadcasting apply.
+
+    Args:
+        a: Quaternions as tensor of shape (..., 4), real part first.
+        b: Quaternions as tensor of shape (..., 4), real part first.
+
+    Returns:
+        The product of a and b, a tensor of quaternions of shape (..., 4).
+    """
+    ab = quaternion_raw_multiply(a, b)
+    return standardize_quaternion(ab)
+
+
+def quaternion_invert(quaternion):
+    """
+    Given a quaternion representing rotation, get the quaternion representing
+    its inverse.
+
+    Args:
+        quaternion: Quaternions as tensor of shape (..., 4), with real part
+            first, which must be versors (unit quaternions).
+
+    Returns:
+        The inverse, a tensor of quaternions of shape (..., 4).
+    """
+
+    return quaternion * quaternion.new_tensor([1, -1, -1, -1])
+
+
+def quaternion_apply(quaternion, point):
+    """
+    Apply the rotation given by a quaternion to a 3D point.
+    Usual torch rules for broadcasting apply.
+
+    Args:
+        quaternion: Tensor of quaternions, real part first, of shape (..., 4).
+        point: Tensor of 3D points of shape (..., 3).
+
+    Returns:
+        Tensor of rotated points of shape (..., 3).
+    """
+    if point.size(-1) != 3:
+        raise ValueError(f"Points are not in 3D, f{point.shape}.")
+    real_parts = point.new_zeros(point.shape[:-1] + (1,))
+    point_as_quaternion = torch.cat((real_parts, point), -1)
+    out = quaternion_raw_multiply(
+        quaternion_raw_multiply(quaternion, point_as_quaternion),
+        quaternion_invert(quaternion),
+    )
+    return out[..., 1:]
+
+
+def axis_angle_to_matrix(axis_angle):
+    """
+    Convert rotations given as axis/angle to rotation matrices.
+
+    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:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle))
+
+
+def matrix_to_axis_angle(matrix):
+    """
+    Convert rotations given as rotation matrices to axis/angle.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+
+    Returns:
+        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.
+    """
+    return quaternion_to_axis_angle(matrix_to_quaternion(matrix))
+
+
+def axis_angle_to_quaternion(axis_angle):
+    """
+    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 = 0.5 * angles
+    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
+
+
+def quaternion_to_axis_angle(quaternions):
+    """
+    Convert rotations given as quaternions to axis/angle.
+
+    Args:
+        quaternions: quaternions with real part first,
+            as tensor of shape (..., 4).
+
+    Returns:
+        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.
+    """
+    norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True)
+    half_angles = torch.atan2(norms, quaternions[..., :1])
+    angles = 2 * half_angles
+    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
+    )
+    return quaternions[..., 1:] / sin_half_angles_over_angles
+
+
+def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor:
+    """
+    Converts 6D rotation representation by Zhou et al. [1] to rotation matrix
+    using Gram--Schmidt orthogonalisation per Section B of [1].
+    Args:
+        d6: 6D rotation representation, of size (*, 6)
+
+    Returns:
+        batch of rotation matrices of size (*, 3, 3)
+
+    [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
+    On the Continuity of Rotation Representations in Neural Networks.
+    IEEE Conference on Computer Vision and Pattern Recognition, 2019.
+    Retrieved from http://arxiv.org/abs/1812.07035
+    """
+
+    a1, a2 = d6[..., :3], d6[..., 3:]
+    b1 = F.normalize(a1, dim=-1)
+    b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1
+    b2 = F.normalize(b2, dim=-1)
+    b3 = torch.cross(b1, b2, dim=-1)
+    return torch.stack((b1, b2, b3), dim=-2)
+
+
+def matrix_to_rotation_6d(matrix: torch.Tensor) -> torch.Tensor:
+    """
+    Converts rotation matrices to 6D rotation representation by Zhou et al. [1]
+    by dropping the last row. Note that 6D representation is not unique.
+    Args:
+        matrix: batch of rotation matrices of size (*, 3, 3)
+
+    Returns:
+        6D rotation representation, of size (*, 6)
+
+    [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
+    On the Continuity of Rotation Representations in Neural Networks.
+    IEEE Conference on Computer Vision and Pattern Recognition, 2019.
+    Retrieved from http://arxiv.org/abs/1812.07035
+    """
+    return matrix[..., :2, :].clone().reshape(*matrix.size()[:-2], 6)
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/smpl_lite.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/smpl_lite.py
new file mode 100644
index 0000000000000000000000000000000000000000..0d1c5e9c3eb5bc6e35a411a80cc26875169fbb9a
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/smpl_lite.py
@@ -0,0 +1,148 @@
+"""Lightweight SMPL body models (SmplLite, SmplxLiteJ24) for efficient FK and skinning."""
+
+import pickle
+from pathlib import Path
+from time import time
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import einsum, rearrange
+from smplx.utils import Struct, to_np, to_tensor
+
+from .rotation_conversions import axis_angle_to_matrix
+
+from .smplx_lite import batch_rigid_transform_v2
+
+
+class SmplLite(nn.Module):
+    def __init__(
+        self,
+        model_path="inputs/checkpoints/body_models/smpl",
+        gender="neutral",
+        num_betas=10,
+    ):
+        super().__init__()
+
+        # Load the model
+        model_path = Path(model_path)
+        if model_path.is_dir():
+            smpl_path = Path(model_path) / f"SMPL_{gender.upper()}.pkl"
+        else:
+            smpl_path = model_path
+        assert smpl_path.exists()
+        with open(smpl_path, "rb") as smpl_file:
+            data_struct = Struct(**pickle.load(smpl_file, encoding="latin1"))
+        self.faces = data_struct.f  # (F, 3)
+
+        self.register_smpl_buffers(data_struct, num_betas)
+        self.register_fast_skeleton_computing_buffers()
+
+    def register_smpl_buffers(self, data_struct, num_betas):
+        # shapedirs, (V, 3, N_betas), V=10475 for SMPLX
+        shapedirs = to_tensor(to_np(data_struct.shapedirs[:, :, :num_betas])).float()
+        self.register_buffer("shapedirs", shapedirs, False)
+
+        # v_template, (V, 3)
+        v_template = to_tensor(to_np(data_struct.v_template)).float()
+        self.register_buffer("v_template", v_template, False)
+
+        # J_regressor, (J, V), J=55 for SMPLX
+        J_regressor = to_tensor(to_np(data_struct.J_regressor)).float()
+        self.register_buffer("J_regressor", J_regressor, False)
+
+        # posedirs, (54*9, V, 3), note that the first global_orient is not included
+        posedirs = to_tensor(to_np(data_struct.posedirs)).float()  # (V, 3, 54*9)
+        posedirs = rearrange(posedirs, "v c n -> n v c")
+        self.register_buffer("posedirs", posedirs, False)
+
+        # lbs_weights, (V, J), J=55
+        lbs_weights = to_tensor(to_np(data_struct.weights)).float()
+        self.register_buffer("lbs_weights", lbs_weights, False)
+
+        # parents, (J), long
+        parents = to_tensor(to_np(data_struct.kintree_table[0])).long()
+        parents[0] = -1
+        self.register_buffer("parents", parents, False)
+
+    def register_fast_skeleton_computing_buffers(self):
+        # For fast computing of skeleton under beta
+        J_template = self.J_regressor @ self.v_template  # (J, 3)
+        J_shapedirs = torch.einsum("jv, vcd -> jcd", self.J_regressor, self.shapedirs)  # (J, 3, 10)
+        self.register_buffer("J_template", J_template, False)
+        self.register_buffer("J_shapedirs", J_shapedirs, False)
+
+    def get_skeleton(self, betas):
+        return self.J_template + einsum(betas, self.J_shapedirs, "... k, j c k -> ... j c")
+
+    def forward(
+        self,
+        body_pose,
+        betas,
+        global_orient,
+        transl,
+    ):
+        """
+        Args:
+            body_pose: (B, L, 63)
+            betas: (B, L, 10)
+            global_orient: (B, L, 3)
+            transl: (B, L, 3)
+        Returns:
+            vertices: (B, L, V, 3)
+        """
+        # 1. Convert [global_orient, body_pose] to rot_mats
+        full_pose = torch.cat([global_orient, body_pose], dim=-1)
+        rot_mats = axis_angle_to_matrix(
+            full_pose.reshape(*full_pose.shape[:-1], full_pose.shape[-1] // 3, 3)
+        )
+
+        # 2. Forward Kinematics
+        J = self.get_skeleton(betas)  # (*, 55, 3)
+        A = batch_rigid_transform_v2(rot_mats, J, self.parents)[1]
+
+        # 3. Canonical v_posed = v_template + shaped_offsets + pose_offsets
+        pose_feature = rot_mats[..., 1:, :, :] - rot_mats.new([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
+        pose_feature = pose_feature.view(*pose_feature.shape[:-3], -1)  # (*, 55*3*3)
+        v_posed = (
+            self.v_template
+            + einsum(betas, self.shapedirs, "... k, v c k -> ... v c")
+            + einsum(pose_feature, self.posedirs, "... k, k v c -> ... v c")
+        )
+        del pose_feature, rot_mats, full_pose
+
+        # 4. Skinning
+        T = einsum(self.lbs_weights, A, "v j, ... j c d -> ... v c d")
+        verts = einsum(T[..., :3, :3], v_posed, "... v c d, ... v d -> ... v c") + T[..., :3, 3]
+
+        # 5. Translation
+        verts = verts + transl[..., None, :]
+        return verts
+
+
+class SmplxLiteJ24(SmplLite):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        # Compute mapping
+        smpl2j24 = self.J_regressor  # (24, 6890)
+
+        jids, smplx_vids = torch.where(smpl2j24 != 0)
+        interestd = torch.zeros([len(smplx_vids), 24])
+        for idx, (jid, smplx_vid) in enumerate(zip(jids, smplx_vids)):
+            interestd[idx, jid] = smpl2j24[jid, smplx_vid]
+        self.register_buffer("interestd", interestd, False)  # (236, 24)
+
+        # Update to vertices of interest
+        self.v_template = self.v_template[smplx_vids].clone()  # (V', 3)
+        self.shapedirs = self.shapedirs[smplx_vids].clone()  # (V', 3, K)
+        self.posedirs = self.posedirs[:, smplx_vids].clone()  # (K, V', 3)
+        self.lbs_weights = self.lbs_weights[smplx_vids].clone()  # (V', J)
+
+    def forward(self, body_pose, betas, global_orient, transl):
+        """Returns: joints (*, J, 3). (B, L) or  (B,) are both supported."""
+        # Use super class's forward to get verts
+        verts = super().forward(body_pose, betas, global_orient, transl)  # (*, 236, 3)
+        joints = einsum(self.interestd, verts, "v j, ... v c -> ... j c")
+        return joints
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/smpl_vert_segmentation.json b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/smpl_vert_segmentation.json
new file mode 100644
index 0000000000000000000000000000000000000000..898f0757a893ea347e05f7c7da96a0ce6a515b0a
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/smpl_vert_segmentation.json
@@ -0,0 +1,7440 @@
+{
+    "rightHand": [
+        5442,
+        5443,
+        5444,
+        5445,
+        5446,
+        5447,
+        5448,
+        5449,
+        5450,
+        5451,
+        5452,
+        5453,
+        5454,
+        5455,
+        5456,
+        5457,
+        5458,
+        5459,
+        5460,
+        5461,
+        5462,
+        5463,
+        5464,
+        5465,
+        5466,
+        5467,
+        5468,
+        5469,
+        5470,
+        5471,
+        5472,
+        5473,
+        5474,
+        5475,
+        5476,
+        5477,
+        5478,
+        5479,
+        5480,
+        5481,
+        5482,
+        5483,
+        5484,
+        5485,
+        5486,
+        5487,
+        5492,
+        5493,
+        5494,
+        5495,
+        5496,
+        5497,
+        5502,
+        5503,
+        5504,
+        5505,
+        5506,
+        5507,
+        5508,
+        5509,
+        5510,
+        5511,
+        5512,
+        5513,
+        5514,
+        5515,
+        5516,
+        5517,
+        5518,
+        5519,
+        5520,
+        5521,
+        5522,
+        5523,
+        5524,
+        5525,
+        5526,
+        5527,
+        5530,
+        5531,
+        5532,
+        5533,
+        5534,
+        5535,
+        5536,
+        5537,
+        5538,
+        5539,
+        5540,
+        5541,
+        5542,
+        5543,
+        5544,
+        5545,
+        5546,
+        5547,
+        5548,
+        5549,
+        5550,
+        5551,
+        5552,
+        5553,
+        5554,
+        5555,
+        5556,
+        5557,
+        5558,
+        5559,
+        5560,
+        5561,
+        5562,
+        5569,
+        5571,
+        5574,
+        5575,
+        5576,
+        5577,
+        5578,
+        5579,
+        5580,
+        5581,
+        5582,
+        5583,
+        5588,
+        5589,
+        5592,
+        5593,
+        5594,
+        5595,
+        5596,
+        5597,
+        5598,
+        5599,
+        5600,
+        5601,
+        5602,
+        5603,
+        5604,
+        5605,
+        5610,
+        5611,
+        5612,
+        5613,
+        5614,
+        5621,
+        5622,
+        5625,
+        5631,
+        5632,
+        5633,
+        5634,
+        5635,
+        5636,
+        5637,
+        5638,
+        5639,
+        5640,
+        5641,
+        5643,
+        5644,
+        5645,
+        5646,
+        5649,
+        5650,
+        5652,
+        5653,
+        5654,
+        5655,
+        5656,
+        5657,
+        5658,
+        5659,
+        5660,
+        5661,
+        5662,
+        5663,
+        5664,
+        5667,
+        5670,
+        5671,
+        5672,
+        5673,
+        5674,
+        5675,
+        5682,
+        5683,
+        5684,
+        5685,
+        5686,
+        5687,
+        5688,
+        5689,
+        5690,
+        5692,
+        5695,
+        5697,
+        5698,
+        5699,
+        5700,
+        5701,
+        5707,
+        5708,
+        5709,
+        5710,
+        5711,
+        5712,
+        5713,
+        5714,
+        5715,
+        5716,
+        5717,
+        5718,
+        5719,
+        5720,
+        5721,
+        5723,
+        5724,
+        5725,
+        5726,
+        5727,
+        5728,
+        5729,
+        5730,
+        5731,
+        5732,
+        5735,
+        5736,
+        5737,
+        5738,
+        5739,
+        5740,
+        5745,
+        5746,
+        5748,
+        5749,
+        5750,
+        5751,
+        5752,
+        6056,
+        6057,
+        6066,
+        6067,
+        6158,
+        6159,
+        6160,
+        6161,
+        6162,
+        6163,
+        6164,
+        6165,
+        6166,
+        6167,
+        6168,
+        6169,
+        6170,
+        6171,
+        6172,
+        6173,
+        6174,
+        6175,
+        6176,
+        6177,
+        6178,
+        6179,
+        6180,
+        6181,
+        6182,
+        6183,
+        6184,
+        6185,
+        6186,
+        6187,
+        6188,
+        6189,
+        6190,
+        6191,
+        6192,
+        6193,
+        6194,
+        6195,
+        6196,
+        6197,
+        6198,
+        6199,
+        6200,
+        6201,
+        6202,
+        6203,
+        6204,
+        6205,
+        6206,
+        6207,
+        6208,
+        6209,
+        6210,
+        6211,
+        6212,
+        6213,
+        6214,
+        6215,
+        6216,
+        6217,
+        6218,
+        6219,
+        6220,
+        6221,
+        6222,
+        6223,
+        6224,
+        6225,
+        6226,
+        6227,
+        6228,
+        6229,
+        6230,
+        6231,
+        6232,
+        6233,
+        6234,
+        6235,
+        6236,
+        6237,
+        6238,
+        6239
+    ],
+    "rightUpLeg": [
+        4320,
+        4321,
+        4323,
+        4324,
+        4333,
+        4334,
+        4335,
+        4336,
+        4337,
+        4338,
+        4339,
+        4340,
+        4356,
+        4357,
+        4358,
+        4359,
+        4360,
+        4361,
+        4362,
+        4363,
+        4364,
+        4365,
+        4366,
+        4367,
+        4383,
+        4384,
+        4385,
+        4386,
+        4387,
+        4388,
+        4389,
+        4390,
+        4391,
+        4392,
+        4393,
+        4394,
+        4395,
+        4396,
+        4397,
+        4398,
+        4399,
+        4400,
+        4401,
+        4419,
+        4420,
+        4421,
+        4422,
+        4430,
+        4431,
+        4432,
+        4433,
+        4434,
+        4435,
+        4436,
+        4437,
+        4438,
+        4439,
+        4440,
+        4441,
+        4442,
+        4443,
+        4444,
+        4445,
+        4446,
+        4447,
+        4448,
+        4449,
+        4450,
+        4451,
+        4452,
+        4453,
+        4454,
+        4455,
+        4456,
+        4457,
+        4458,
+        4459,
+        4460,
+        4461,
+        4462,
+        4463,
+        4464,
+        4465,
+        4466,
+        4467,
+        4468,
+        4469,
+        4470,
+        4471,
+        4472,
+        4473,
+        4474,
+        4475,
+        4476,
+        4477,
+        4478,
+        4479,
+        4480,
+        4481,
+        4482,
+        4483,
+        4484,
+        4485,
+        4486,
+        4487,
+        4488,
+        4489,
+        4490,
+        4491,
+        4492,
+        4493,
+        4494,
+        4495,
+        4496,
+        4497,
+        4498,
+        4499,
+        4500,
+        4501,
+        4502,
+        4503,
+        4504,
+        4505,
+        4506,
+        4507,
+        4508,
+        4509,
+        4510,
+        4511,
+        4512,
+        4513,
+        4514,
+        4515,
+        4516,
+        4517,
+        4518,
+        4519,
+        4520,
+        4521,
+        4522,
+        4523,
+        4524,
+        4525,
+        4526,
+        4527,
+        4528,
+        4529,
+        4530,
+        4531,
+        4532,
+        4623,
+        4624,
+        4625,
+        4626,
+        4627,
+        4628,
+        4629,
+        4630,
+        4631,
+        4632,
+        4633,
+        4634,
+        4645,
+        4646,
+        4647,
+        4648,
+        4649,
+        4650,
+        4651,
+        4652,
+        4653,
+        4654,
+        4655,
+        4656,
+        4657,
+        4658,
+        4659,
+        4660,
+        4670,
+        4671,
+        4672,
+        4673,
+        4704,
+        4705,
+        4706,
+        4707,
+        4708,
+        4709,
+        4710,
+        4711,
+        4712,
+        4713,
+        4745,
+        4746,
+        4757,
+        4758,
+        4759,
+        4760,
+        4801,
+        4802,
+        4829,
+        4834,
+        4835,
+        4836,
+        4837,
+        4838,
+        4839,
+        4840,
+        4841,
+        4924,
+        4925,
+        4926,
+        4928,
+        4929,
+        4930,
+        4931,
+        4932,
+        4933,
+        4934,
+        4935,
+        4936,
+        4948,
+        4949,
+        4950,
+        4951,
+        4952,
+        4970,
+        4971,
+        4972,
+        4973,
+        4983,
+        4984,
+        4985,
+        4986,
+        4987,
+        4988,
+        4989,
+        4990,
+        4991,
+        4992,
+        4993,
+        5004,
+        5005,
+        6546,
+        6547,
+        6548,
+        6549,
+        6552,
+        6553,
+        6554,
+        6555,
+        6556,
+        6873,
+        6877
+    ],
+    "leftArm": [
+        626,
+        627,
+        628,
+        629,
+        634,
+        635,
+        680,
+        681,
+        716,
+        717,
+        718,
+        719,
+        769,
+        770,
+        771,
+        772,
+        773,
+        774,
+        775,
+        776,
+        777,
+        778,
+        779,
+        780,
+        784,
+        785,
+        786,
+        787,
+        788,
+        789,
+        790,
+        791,
+        792,
+        793,
+        1231,
+        1232,
+        1233,
+        1234,
+        1258,
+        1259,
+        1260,
+        1261,
+        1271,
+        1281,
+        1282,
+        1310,
+        1311,
+        1314,
+        1315,
+        1340,
+        1341,
+        1342,
+        1343,
+        1355,
+        1356,
+        1357,
+        1358,
+        1376,
+        1377,
+        1378,
+        1379,
+        1380,
+        1381,
+        1382,
+        1383,
+        1384,
+        1385,
+        1386,
+        1387,
+        1388,
+        1389,
+        1390,
+        1391,
+        1392,
+        1393,
+        1394,
+        1395,
+        1396,
+        1397,
+        1398,
+        1399,
+        1400,
+        1402,
+        1403,
+        1405,
+        1406,
+        1407,
+        1408,
+        1409,
+        1410,
+        1411,
+        1412,
+        1413,
+        1414,
+        1415,
+        1416,
+        1428,
+        1429,
+        1430,
+        1431,
+        1432,
+        1433,
+        1438,
+        1439,
+        1440,
+        1441,
+        1442,
+        1443,
+        1444,
+        1445,
+        1502,
+        1505,
+        1506,
+        1507,
+        1508,
+        1509,
+        1510,
+        1538,
+        1541,
+        1542,
+        1543,
+        1545,
+        1619,
+        1620,
+        1621,
+        1622,
+        1631,
+        1632,
+        1633,
+        1634,
+        1635,
+        1636,
+        1637,
+        1638,
+        1639,
+        1640,
+        1641,
+        1642,
+        1645,
+        1646,
+        1647,
+        1648,
+        1649,
+        1650,
+        1651,
+        1652,
+        1653,
+        1654,
+        1655,
+        1656,
+        1658,
+        1659,
+        1661,
+        1662,
+        1664,
+        1666,
+        1667,
+        1668,
+        1669,
+        1670,
+        1671,
+        1672,
+        1673,
+        1674,
+        1675,
+        1676,
+        1677,
+        1678,
+        1679,
+        1680,
+        1681,
+        1682,
+        1683,
+        1684,
+        1696,
+        1697,
+        1698,
+        1703,
+        1704,
+        1705,
+        1706,
+        1707,
+        1708,
+        1709,
+        1710,
+        1711,
+        1712,
+        1713,
+        1714,
+        1715,
+        1716,
+        1717,
+        1718,
+        1719,
+        1720,
+        1725,
+        1731,
+        1732,
+        1733,
+        1734,
+        1735,
+        1737,
+        1739,
+        1740,
+        1745,
+        1746,
+        1747,
+        1748,
+        1749,
+        1751,
+        1761,
+        1830,
+        1831,
+        1844,
+        1845,
+        1846,
+        1850,
+        1851,
+        1854,
+        1855,
+        1858,
+        1860,
+        1865,
+        1866,
+        1867,
+        1869,
+        1870,
+        1871,
+        1874,
+        1875,
+        1876,
+        1877,
+        1878,
+        1882,
+        1883,
+        1888,
+        1889,
+        1892,
+        1900,
+        1901,
+        1902,
+        1903,
+        1904,
+        1909,
+        2819,
+        2820,
+        2821,
+        2822,
+        2895,
+        2896,
+        2897,
+        2898,
+        2899,
+        2900,
+        2901,
+        2902,
+        2903,
+        2945,
+        2946,
+        2974,
+        2975,
+        2976,
+        2977,
+        2978,
+        2979,
+        2980,
+        2981,
+        2982,
+        2983,
+        2984,
+        2985,
+        2986,
+        2987,
+        2988,
+        2989,
+        2990,
+        2991,
+        2992,
+        2993,
+        2994,
+        2995,
+        2996,
+        3002,
+        3013
+    ],
+    "leftLeg": [
+        995,
+        998,
+        999,
+        1002,
+        1004,
+        1005,
+        1008,
+        1010,
+        1012,
+        1015,
+        1016,
+        1018,
+        1019,
+        1043,
+        1044,
+        1047,
+        1048,
+        1049,
+        1050,
+        1051,
+        1052,
+        1053,
+        1054,
+        1055,
+        1056,
+        1057,
+        1058,
+        1059,
+        1060,
+        1061,
+        1062,
+        1063,
+        1064,
+        1065,
+        1066,
+        1067,
+        1068,
+        1069,
+        1070,
+        1071,
+        1072,
+        1073,
+        1074,
+        1075,
+        1076,
+        1077,
+        1078,
+        1079,
+        1080,
+        1081,
+        1082,
+        1083,
+        1084,
+        1085,
+        1086,
+        1087,
+        1088,
+        1089,
+        1090,
+        1091,
+        1092,
+        1093,
+        1094,
+        1095,
+        1096,
+        1097,
+        1098,
+        1099,
+        1100,
+        1101,
+        1102,
+        1103,
+        1104,
+        1105,
+        1106,
+        1107,
+        1108,
+        1109,
+        1110,
+        1111,
+        1112,
+        1113,
+        1114,
+        1115,
+        1116,
+        1117,
+        1118,
+        1119,
+        1120,
+        1121,
+        1122,
+        1123,
+        1124,
+        1125,
+        1126,
+        1127,
+        1128,
+        1129,
+        1130,
+        1131,
+        1132,
+        1133,
+        1134,
+        1135,
+        1136,
+        1148,
+        1149,
+        1150,
+        1151,
+        1152,
+        1153,
+        1154,
+        1155,
+        1156,
+        1157,
+        1158,
+        1175,
+        1176,
+        1177,
+        1178,
+        1179,
+        1180,
+        1181,
+        1182,
+        1183,
+        1369,
+        1370,
+        1371,
+        1372,
+        1373,
+        1374,
+        1375,
+        1464,
+        1465,
+        1466,
+        1467,
+        1468,
+        1469,
+        1470,
+        1471,
+        1472,
+        1473,
+        1474,
+        1522,
+        1523,
+        1524,
+        1525,
+        1526,
+        1527,
+        1528,
+        1529,
+        1530,
+        1531,
+        1532,
+        3174,
+        3175,
+        3176,
+        3177,
+        3178,
+        3179,
+        3180,
+        3181,
+        3182,
+        3183,
+        3184,
+        3185,
+        3186,
+        3187,
+        3188,
+        3189,
+        3190,
+        3191,
+        3192,
+        3193,
+        3194,
+        3195,
+        3196,
+        3197,
+        3198,
+        3199,
+        3200,
+        3201,
+        3202,
+        3203,
+        3204,
+        3205,
+        3206,
+        3207,
+        3208,
+        3209,
+        3210,
+        3319,
+        3320,
+        3321,
+        3322,
+        3323,
+        3324,
+        3325,
+        3326,
+        3327,
+        3328,
+        3329,
+        3330,
+        3331,
+        3332,
+        3333,
+        3334,
+        3335,
+        3432,
+        3433,
+        3434,
+        3435,
+        3436,
+        3469,
+        3472,
+        3473,
+        3474
+    ],
+    "leftToeBase": [
+        3211,
+        3212,
+        3213,
+        3214,
+        3215,
+        3216,
+        3217,
+        3218,
+        3219,
+        3220,
+        3221,
+        3222,
+        3223,
+        3224,
+        3225,
+        3226,
+        3227,
+        3228,
+        3229,
+        3230,
+        3231,
+        3232,
+        3233,
+        3234,
+        3235,
+        3236,
+        3237,
+        3238,
+        3239,
+        3240,
+        3241,
+        3242,
+        3243,
+        3244,
+        3245,
+        3246,
+        3247,
+        3248,
+        3249,
+        3250,
+        3251,
+        3252,
+        3253,
+        3254,
+        3255,
+        3256,
+        3257,
+        3258,
+        3259,
+        3260,
+        3261,
+        3262,
+        3263,
+        3264,
+        3265,
+        3266,
+        3267,
+        3268,
+        3269,
+        3270,
+        3271,
+        3272,
+        3273,
+        3274,
+        3275,
+        3276,
+        3277,
+        3278,
+        3279,
+        3280,
+        3281,
+        3282,
+        3283,
+        3284,
+        3285,
+        3286,
+        3287,
+        3288,
+        3289,
+        3290,
+        3291,
+        3292,
+        3293,
+        3294,
+        3295,
+        3296,
+        3297,
+        3298,
+        3299,
+        3300,
+        3301,
+        3302,
+        3303,
+        3304,
+        3305,
+        3306,
+        3307,
+        3308,
+        3309,
+        3310,
+        3311,
+        3312,
+        3313,
+        3314,
+        3315,
+        3316,
+        3317,
+        3318,
+        3336,
+        3337,
+        3340,
+        3342,
+        3344,
+        3346,
+        3348,
+        3350,
+        3352,
+        3354,
+        3357,
+        3358,
+        3360,
+        3362
+    ],
+    "leftFoot": [
+        3327,
+        3328,
+        3329,
+        3330,
+        3331,
+        3332,
+        3333,
+        3334,
+        3335,
+        3336,
+        3337,
+        3338,
+        3339,
+        3340,
+        3341,
+        3342,
+        3343,
+        3344,
+        3345,
+        3346,
+        3347,
+        3348,
+        3349,
+        3350,
+        3351,
+        3352,
+        3353,
+        3354,
+        3355,
+        3356,
+        3357,
+        3358,
+        3359,
+        3360,
+        3361,
+        3362,
+        3363,
+        3364,
+        3365,
+        3366,
+        3367,
+        3368,
+        3369,
+        3370,
+        3371,
+        3372,
+        3373,
+        3374,
+        3375,
+        3376,
+        3377,
+        3378,
+        3379,
+        3380,
+        3381,
+        3382,
+        3383,
+        3384,
+        3385,
+        3386,
+        3387,
+        3388,
+        3389,
+        3390,
+        3391,
+        3392,
+        3393,
+        3394,
+        3395,
+        3396,
+        3397,
+        3398,
+        3399,
+        3400,
+        3401,
+        3402,
+        3403,
+        3404,
+        3405,
+        3406,
+        3407,
+        3408,
+        3409,
+        3410,
+        3411,
+        3412,
+        3413,
+        3414,
+        3415,
+        3416,
+        3417,
+        3418,
+        3419,
+        3420,
+        3421,
+        3422,
+        3423,
+        3424,
+        3425,
+        3426,
+        3427,
+        3428,
+        3429,
+        3430,
+        3431,
+        3432,
+        3433,
+        3434,
+        3435,
+        3436,
+        3437,
+        3438,
+        3439,
+        3440,
+        3441,
+        3442,
+        3443,
+        3444,
+        3445,
+        3446,
+        3447,
+        3448,
+        3449,
+        3450,
+        3451,
+        3452,
+        3453,
+        3454,
+        3455,
+        3456,
+        3457,
+        3458,
+        3459,
+        3460,
+        3461,
+        3462,
+        3463,
+        3464,
+        3465,
+        3466,
+        3467,
+        3468,
+        3469
+    ],
+    "spine1": [
+        598,
+        599,
+        600,
+        601,
+        610,
+        611,
+        612,
+        613,
+        614,
+        615,
+        616,
+        617,
+        618,
+        619,
+        620,
+        621,
+        642,
+        645,
+        646,
+        647,
+        652,
+        653,
+        658,
+        659,
+        660,
+        661,
+        668,
+        669,
+        670,
+        671,
+        684,
+        685,
+        686,
+        687,
+        688,
+        689,
+        690,
+        691,
+        692,
+        722,
+        723,
+        724,
+        725,
+        736,
+        750,
+        751,
+        761,
+        764,
+        766,
+        767,
+        794,
+        795,
+        891,
+        892,
+        893,
+        894,
+        925,
+        926,
+        927,
+        928,
+        929,
+        940,
+        941,
+        942,
+        943,
+        1190,
+        1191,
+        1192,
+        1193,
+        1194,
+        1195,
+        1196,
+        1197,
+        1200,
+        1201,
+        1202,
+        1212,
+        1236,
+        1252,
+        1253,
+        1254,
+        1255,
+        1268,
+        1269,
+        1270,
+        1329,
+        1330,
+        1348,
+        1349,
+        1351,
+        1420,
+        1421,
+        1423,
+        1424,
+        1425,
+        1426,
+        1436,
+        1437,
+        1756,
+        1757,
+        1758,
+        2839,
+        2840,
+        2841,
+        2842,
+        2843,
+        2844,
+        2845,
+        2846,
+        2847,
+        2848,
+        2849,
+        2850,
+        2851,
+        2870,
+        2871,
+        2883,
+        2906,
+        2908,
+        3014,
+        3017,
+        3025,
+        3030,
+        3033,
+        3034,
+        3037,
+        3039,
+        3040,
+        3041,
+        3042,
+        3043,
+        3044,
+        3076,
+        3077,
+        3079,
+        3480,
+        3505,
+        3511,
+        4086,
+        4087,
+        4088,
+        4089,
+        4098,
+        4099,
+        4100,
+        4101,
+        4102,
+        4103,
+        4104,
+        4105,
+        4106,
+        4107,
+        4108,
+        4109,
+        4130,
+        4131,
+        4134,
+        4135,
+        4140,
+        4141,
+        4146,
+        4147,
+        4148,
+        4149,
+        4156,
+        4157,
+        4158,
+        4159,
+        4172,
+        4173,
+        4174,
+        4175,
+        4176,
+        4177,
+        4178,
+        4179,
+        4180,
+        4210,
+        4211,
+        4212,
+        4213,
+        4225,
+        4239,
+        4240,
+        4249,
+        4250,
+        4255,
+        4256,
+        4282,
+        4283,
+        4377,
+        4378,
+        4379,
+        4380,
+        4411,
+        4412,
+        4413,
+        4414,
+        4415,
+        4426,
+        4427,
+        4428,
+        4429,
+        4676,
+        4677,
+        4678,
+        4679,
+        4680,
+        4681,
+        4682,
+        4683,
+        4686,
+        4687,
+        4688,
+        4695,
+        4719,
+        4735,
+        4736,
+        4737,
+        4740,
+        4751,
+        4752,
+        4753,
+        4824,
+        4825,
+        4828,
+        4893,
+        4894,
+        4895,
+        4897,
+        4898,
+        4899,
+        4908,
+        4909,
+        5223,
+        5224,
+        5225,
+        6300,
+        6301,
+        6302,
+        6303,
+        6304,
+        6305,
+        6306,
+        6307,
+        6308,
+        6309,
+        6310,
+        6311,
+        6312,
+        6331,
+        6332,
+        6342,
+        6366,
+        6367,
+        6475,
+        6477,
+        6478,
+        6481,
+        6482,
+        6485,
+        6487,
+        6488,
+        6489,
+        6490,
+        6491,
+        6878
+    ],
+    "spine2": [
+        570,
+        571,
+        572,
+        573,
+        584,
+        585,
+        586,
+        587,
+        588,
+        589,
+        590,
+        591,
+        592,
+        593,
+        594,
+        595,
+        596,
+        597,
+        602,
+        603,
+        604,
+        605,
+        606,
+        607,
+        608,
+        609,
+        622,
+        623,
+        624,
+        625,
+        638,
+        639,
+        640,
+        641,
+        643,
+        644,
+        648,
+        649,
+        650,
+        651,
+        666,
+        667,
+        672,
+        673,
+        674,
+        675,
+        680,
+        681,
+        682,
+        683,
+        693,
+        694,
+        695,
+        696,
+        697,
+        698,
+        699,
+        700,
+        701,
+        702,
+        703,
+        704,
+        713,
+        714,
+        715,
+        716,
+        717,
+        726,
+        727,
+        728,
+        729,
+        730,
+        731,
+        732,
+        733,
+        735,
+        737,
+        738,
+        739,
+        740,
+        741,
+        742,
+        743,
+        744,
+        745,
+        746,
+        747,
+        748,
+        749,
+        752,
+        753,
+        754,
+        755,
+        756,
+        757,
+        758,
+        759,
+        760,
+        762,
+        763,
+        803,
+        804,
+        805,
+        806,
+        811,
+        812,
+        813,
+        814,
+        817,
+        818,
+        819,
+        820,
+        821,
+        824,
+        825,
+        826,
+        827,
+        828,
+        895,
+        896,
+        930,
+        931,
+        1198,
+        1199,
+        1213,
+        1214,
+        1215,
+        1216,
+        1217,
+        1218,
+        1219,
+        1220,
+        1235,
+        1237,
+        1256,
+        1257,
+        1271,
+        1272,
+        1273,
+        1279,
+        1280,
+        1283,
+        1284,
+        1285,
+        1286,
+        1287,
+        1288,
+        1289,
+        1290,
+        1291,
+        1292,
+        1293,
+        1294,
+        1295,
+        1296,
+        1297,
+        1298,
+        1299,
+        1300,
+        1301,
+        1302,
+        1303,
+        1304,
+        1305,
+        1306,
+        1307,
+        1308,
+        1309,
+        1312,
+        1313,
+        1319,
+        1320,
+        1346,
+        1347,
+        1350,
+        1352,
+        1401,
+        1417,
+        1418,
+        1419,
+        1422,
+        1427,
+        1434,
+        1435,
+        1503,
+        1504,
+        1536,
+        1537,
+        1544,
+        1545,
+        1753,
+        1754,
+        1755,
+        1759,
+        1760,
+        1761,
+        1762,
+        1763,
+        1808,
+        1809,
+        1810,
+        1811,
+        1816,
+        1817,
+        1818,
+        1819,
+        1820,
+        1834,
+        1835,
+        1836,
+        1837,
+        1838,
+        1839,
+        1868,
+        1879,
+        1880,
+        2812,
+        2813,
+        2852,
+        2853,
+        2854,
+        2855,
+        2856,
+        2857,
+        2858,
+        2859,
+        2860,
+        2861,
+        2862,
+        2863,
+        2864,
+        2865,
+        2866,
+        2867,
+        2868,
+        2869,
+        2872,
+        2875,
+        2876,
+        2877,
+        2878,
+        2881,
+        2882,
+        2884,
+        2885,
+        2886,
+        2904,
+        2905,
+        2907,
+        2931,
+        2932,
+        2933,
+        2934,
+        2935,
+        2936,
+        2937,
+        2941,
+        2950,
+        2951,
+        2952,
+        2953,
+        2954,
+        2955,
+        2956,
+        2957,
+        2958,
+        2959,
+        2960,
+        2961,
+        2962,
+        2963,
+        2964,
+        2965,
+        2966,
+        2967,
+        2968,
+        2969,
+        2970,
+        2971,
+        2972,
+        2973,
+        2997,
+        2998,
+        3006,
+        3007,
+        3012,
+        3015,
+        3026,
+        3027,
+        3028,
+        3029,
+        3031,
+        3032,
+        3035,
+        3036,
+        3038,
+        3059,
+        3060,
+        3061,
+        3062,
+        3063,
+        3064,
+        3065,
+        3066,
+        3067,
+        3073,
+        3074,
+        3075,
+        3078,
+        3168,
+        3169,
+        3171,
+        3470,
+        3471,
+        3482,
+        3483,
+        3495,
+        3496,
+        3497,
+        3498,
+        3506,
+        3508,
+        4058,
+        4059,
+        4060,
+        4061,
+        4072,
+        4073,
+        4074,
+        4075,
+        4076,
+        4077,
+        4078,
+        4079,
+        4080,
+        4081,
+        4082,
+        4083,
+        4084,
+        4085,
+        4090,
+        4091,
+        4092,
+        4093,
+        4094,
+        4095,
+        4096,
+        4097,
+        4110,
+        4111,
+        4112,
+        4113,
+        4126,
+        4127,
+        4128,
+        4129,
+        4132,
+        4133,
+        4136,
+        4137,
+        4138,
+        4139,
+        4154,
+        4155,
+        4160,
+        4161,
+        4162,
+        4163,
+        4168,
+        4169,
+        4170,
+        4171,
+        4181,
+        4182,
+        4183,
+        4184,
+        4185,
+        4186,
+        4187,
+        4188,
+        4189,
+        4190,
+        4191,
+        4192,
+        4201,
+        4202,
+        4203,
+        4204,
+        4207,
+        4214,
+        4215,
+        4216,
+        4217,
+        4218,
+        4219,
+        4220,
+        4221,
+        4223,
+        4224,
+        4226,
+        4227,
+        4228,
+        4229,
+        4230,
+        4231,
+        4232,
+        4233,
+        4234,
+        4235,
+        4236,
+        4237,
+        4238,
+        4241,
+        4242,
+        4243,
+        4244,
+        4245,
+        4246,
+        4247,
+        4248,
+        4251,
+        4252,
+        4291,
+        4292,
+        4293,
+        4294,
+        4299,
+        4300,
+        4301,
+        4302,
+        4305,
+        4306,
+        4307,
+        4308,
+        4309,
+        4312,
+        4313,
+        4314,
+        4315,
+        4381,
+        4382,
+        4416,
+        4417,
+        4684,
+        4685,
+        4696,
+        4697,
+        4698,
+        4699,
+        4700,
+        4701,
+        4702,
+        4703,
+        4718,
+        4720,
+        4738,
+        4739,
+        4754,
+        4755,
+        4756,
+        4761,
+        4762,
+        4765,
+        4766,
+        4767,
+        4768,
+        4769,
+        4770,
+        4771,
+        4772,
+        4773,
+        4774,
+        4775,
+        4776,
+        4777,
+        4778,
+        4779,
+        4780,
+        4781,
+        4782,
+        4783,
+        4784,
+        4785,
+        4786,
+        4787,
+        4788,
+        4789,
+        4792,
+        4793,
+        4799,
+        4800,
+        4822,
+        4823,
+        4826,
+        4827,
+        4874,
+        4890,
+        4891,
+        4892,
+        4896,
+        4900,
+        4907,
+        4910,
+        4975,
+        4976,
+        5007,
+        5008,
+        5013,
+        5014,
+        5222,
+        5226,
+        5227,
+        5228,
+        5229,
+        5230,
+        5269,
+        5270,
+        5271,
+        5272,
+        5277,
+        5278,
+        5279,
+        5280,
+        5281,
+        5295,
+        5296,
+        5297,
+        5298,
+        5299,
+        5300,
+        5329,
+        5340,
+        5341,
+        6273,
+        6274,
+        6313,
+        6314,
+        6315,
+        6316,
+        6317,
+        6318,
+        6319,
+        6320,
+        6321,
+        6322,
+        6323,
+        6324,
+        6325,
+        6326,
+        6327,
+        6328,
+        6329,
+        6330,
+        6333,
+        6336,
+        6337,
+        6340,
+        6341,
+        6343,
+        6344,
+        6345,
+        6363,
+        6364,
+        6365,
+        6390,
+        6391,
+        6392,
+        6393,
+        6394,
+        6395,
+        6396,
+        6398,
+        6409,
+        6410,
+        6411,
+        6412,
+        6413,
+        6414,
+        6415,
+        6416,
+        6417,
+        6418,
+        6419,
+        6420,
+        6421,
+        6422,
+        6423,
+        6424,
+        6425,
+        6426,
+        6427,
+        6428,
+        6429,
+        6430,
+        6431,
+        6432,
+        6456,
+        6457,
+        6465,
+        6466,
+        6476,
+        6479,
+        6480,
+        6483,
+        6484,
+        6486,
+        6496,
+        6497,
+        6498,
+        6499,
+        6500,
+        6501,
+        6502,
+        6503,
+        6879
+    ],
+    "leftShoulder": [
+        591,
+        604,
+        605,
+        606,
+        609,
+        634,
+        635,
+        636,
+        637,
+        674,
+        706,
+        707,
+        708,
+        709,
+        710,
+        711,
+        712,
+        713,
+        715,
+        717,
+        730,
+        733,
+        734,
+        735,
+        781,
+        782,
+        783,
+        1238,
+        1239,
+        1240,
+        1241,
+        1242,
+        1243,
+        1244,
+        1245,
+        1290,
+        1291,
+        1294,
+        1316,
+        1317,
+        1318,
+        1401,
+        1402,
+        1403,
+        1404,
+        1509,
+        1535,
+        1545,
+        1808,
+        1810,
+        1811,
+        1812,
+        1813,
+        1814,
+        1815,
+        1818,
+        1819,
+        1821,
+        1822,
+        1823,
+        1824,
+        1825,
+        1826,
+        1827,
+        1828,
+        1829,
+        1830,
+        1831,
+        1832,
+        1833,
+        1837,
+        1840,
+        1841,
+        1842,
+        1843,
+        1844,
+        1845,
+        1846,
+        1847,
+        1848,
+        1849,
+        1850,
+        1851,
+        1852,
+        1853,
+        1854,
+        1855,
+        1856,
+        1857,
+        1858,
+        1859,
+        1861,
+        1862,
+        1863,
+        1864,
+        1872,
+        1873,
+        1880,
+        1881,
+        1884,
+        1885,
+        1886,
+        1887,
+        1890,
+        1891,
+        1893,
+        1894,
+        1895,
+        1896,
+        1897,
+        1898,
+        1899,
+        2879,
+        2880,
+        2881,
+        2886,
+        2887,
+        2888,
+        2889,
+        2890,
+        2891,
+        2892,
+        2893,
+        2894,
+        2903,
+        2938,
+        2939,
+        2940,
+        2941,
+        2942,
+        2943,
+        2944,
+        2945,
+        2946,
+        2947,
+        2948,
+        2949,
+        2965,
+        2967,
+        2969,
+        2999,
+        3000,
+        3001,
+        3002,
+        3003,
+        3004,
+        3005,
+        3008,
+        3009,
+        3010,
+        3011
+    ],
+    "rightShoulder": [
+        4077,
+        4091,
+        4092,
+        4094,
+        4095,
+        4122,
+        4123,
+        4124,
+        4125,
+        4162,
+        4194,
+        4195,
+        4196,
+        4197,
+        4198,
+        4199,
+        4200,
+        4201,
+        4203,
+        4207,
+        4218,
+        4219,
+        4222,
+        4223,
+        4269,
+        4270,
+        4271,
+        4721,
+        4722,
+        4723,
+        4724,
+        4725,
+        4726,
+        4727,
+        4728,
+        4773,
+        4774,
+        4778,
+        4796,
+        4797,
+        4798,
+        4874,
+        4875,
+        4876,
+        4877,
+        4982,
+        5006,
+        5014,
+        5269,
+        5271,
+        5272,
+        5273,
+        5274,
+        5275,
+        5276,
+        5279,
+        5281,
+        5282,
+        5283,
+        5284,
+        5285,
+        5286,
+        5287,
+        5288,
+        5289,
+        5290,
+        5291,
+        5292,
+        5293,
+        5294,
+        5298,
+        5301,
+        5302,
+        5303,
+        5304,
+        5305,
+        5306,
+        5307,
+        5308,
+        5309,
+        5310,
+        5311,
+        5312,
+        5313,
+        5314,
+        5315,
+        5316,
+        5317,
+        5318,
+        5319,
+        5320,
+        5322,
+        5323,
+        5324,
+        5325,
+        5333,
+        5334,
+        5341,
+        5342,
+        5345,
+        5346,
+        5347,
+        5348,
+        5351,
+        5352,
+        5354,
+        5355,
+        5356,
+        5357,
+        5358,
+        5359,
+        5360,
+        6338,
+        6339,
+        6340,
+        6345,
+        6346,
+        6347,
+        6348,
+        6349,
+        6350,
+        6351,
+        6352,
+        6353,
+        6362,
+        6397,
+        6398,
+        6399,
+        6400,
+        6401,
+        6402,
+        6403,
+        6404,
+        6405,
+        6406,
+        6407,
+        6408,
+        6424,
+        6425,
+        6428,
+        6458,
+        6459,
+        6460,
+        6461,
+        6462,
+        6463,
+        6464,
+        6467,
+        6468,
+        6469,
+        6470
+    ],
+    "rightFoot": [
+        6727,
+        6728,
+        6729,
+        6730,
+        6731,
+        6732,
+        6733,
+        6734,
+        6735,
+        6736,
+        6737,
+        6738,
+        6739,
+        6740,
+        6741,
+        6742,
+        6743,
+        6744,
+        6745,
+        6746,
+        6747,
+        6748,
+        6749,
+        6750,
+        6751,
+        6752,
+        6753,
+        6754,
+        6755,
+        6756,
+        6757,
+        6758,
+        6759,
+        6760,
+        6761,
+        6762,
+        6763,
+        6764,
+        6765,
+        6766,
+        6767,
+        6768,
+        6769,
+        6770,
+        6771,
+        6772,
+        6773,
+        6774,
+        6775,
+        6776,
+        6777,
+        6778,
+        6779,
+        6780,
+        6781,
+        6782,
+        6783,
+        6784,
+        6785,
+        6786,
+        6787,
+        6788,
+        6789,
+        6790,
+        6791,
+        6792,
+        6793,
+        6794,
+        6795,
+        6796,
+        6797,
+        6798,
+        6799,
+        6800,
+        6801,
+        6802,
+        6803,
+        6804,
+        6805,
+        6806,
+        6807,
+        6808,
+        6809,
+        6810,
+        6811,
+        6812,
+        6813,
+        6814,
+        6815,
+        6816,
+        6817,
+        6818,
+        6819,
+        6820,
+        6821,
+        6822,
+        6823,
+        6824,
+        6825,
+        6826,
+        6827,
+        6828,
+        6829,
+        6830,
+        6831,
+        6832,
+        6833,
+        6834,
+        6835,
+        6836,
+        6837,
+        6838,
+        6839,
+        6840,
+        6841,
+        6842,
+        6843,
+        6844,
+        6845,
+        6846,
+        6847,
+        6848,
+        6849,
+        6850,
+        6851,
+        6852,
+        6853,
+        6854,
+        6855,
+        6856,
+        6857,
+        6858,
+        6859,
+        6860,
+        6861,
+        6862,
+        6863,
+        6864,
+        6865,
+        6866,
+        6867,
+        6868,
+        6869
+    ],
+    "head": [
+        0,
+        1,
+        2,
+        3,
+        4,
+        5,
+        6,
+        7,
+        8,
+        9,
+        10,
+        11,
+        12,
+        13,
+        14,
+        15,
+        16,
+        17,
+        18,
+        19,
+        20,
+        21,
+        22,
+        23,
+        24,
+        25,
+        26,
+        27,
+        28,
+        29,
+        30,
+        31,
+        32,
+        33,
+        34,
+        35,
+        36,
+        37,
+        38,
+        39,
+        40,
+        41,
+        42,
+        43,
+        44,
+        45,
+        46,
+        47,
+        48,
+        49,
+        50,
+        51,
+        52,
+        53,
+        54,
+        55,
+        56,
+        57,
+        58,
+        59,
+        60,
+        61,
+        62,
+        63,
+        64,
+        65,
+        66,
+        67,
+        68,
+        69,
+        70,
+        71,
+        72,
+        73,
+        74,
+        75,
+        76,
+        77,
+        78,
+        79,
+        80,
+        81,
+        82,
+        83,
+        84,
+        85,
+        86,
+        87,
+        88,
+        89,
+        90,
+        91,
+        92,
+        93,
+        94,
+        95,
+        96,
+        97,
+        98,
+        99,
+        100,
+        101,
+        102,
+        103,
+        104,
+        105,
+        106,
+        107,
+        108,
+        109,
+        110,
+        111,
+        112,
+        113,
+        114,
+        115,
+        116,
+        117,
+        118,
+        119,
+        120,
+        121,
+        122,
+        123,
+        124,
+        125,
+        126,
+        127,
+        128,
+        129,
+        130,
+        131,
+        132,
+        133,
+        134,
+        135,
+        136,
+        137,
+        138,
+        139,
+        140,
+        141,
+        142,
+        143,
+        144,
+        145,
+        146,
+        147,
+        148,
+        149,
+        154,
+        155,
+        156,
+        157,
+        158,
+        159,
+        160,
+        161,
+        162,
+        163,
+        164,
+        165,
+        166,
+        167,
+        168,
+        169,
+        170,
+        171,
+        172,
+        173,
+        176,
+        177,
+        178,
+        179,
+        180,
+        181,
+        182,
+        183,
+        184,
+        185,
+        186,
+        187,
+        188,
+        189,
+        190,
+        191,
+        192,
+        193,
+        194,
+        195,
+        196,
+        197,
+        198,
+        199,
+        200,
+        201,
+        202,
+        203,
+        204,
+        205,
+        220,
+        221,
+        225,
+        226,
+        227,
+        228,
+        229,
+        230,
+        231,
+        232,
+        233,
+        234,
+        235,
+        236,
+        237,
+        238,
+        239,
+        240,
+        241,
+        242,
+        243,
+        244,
+        245,
+        246,
+        247,
+        248,
+        249,
+        250,
+        251,
+        252,
+        253,
+        254,
+        255,
+        258,
+        259,
+        260,
+        261,
+        262,
+        263,
+        264,
+        265,
+        266,
+        267,
+        268,
+        269,
+        270,
+        271,
+        272,
+        273,
+        274,
+        275,
+        276,
+        277,
+        278,
+        279,
+        280,
+        281,
+        282,
+        283,
+        286,
+        287,
+        288,
+        289,
+        290,
+        291,
+        292,
+        293,
+        294,
+        295,
+        303,
+        304,
+        306,
+        307,
+        310,
+        311,
+        312,
+        313,
+        314,
+        315,
+        316,
+        317,
+        318,
+        319,
+        320,
+        321,
+        322,
+        323,
+        324,
+        325,
+        326,
+        327,
+        328,
+        329,
+        330,
+        331,
+        332,
+        335,
+        336,
+        337,
+        338,
+        339,
+        340,
+        341,
+        342,
+        343,
+        344,
+        345,
+        346,
+        347,
+        348,
+        349,
+        350,
+        351,
+        352,
+        353,
+        354,
+        355,
+        356,
+        357,
+        358,
+        359,
+        360,
+        361,
+        362,
+        363,
+        364,
+        365,
+        366,
+        367,
+        368,
+        369,
+        370,
+        371,
+        372,
+        373,
+        374,
+        375,
+        376,
+        377,
+        378,
+        379,
+        380,
+        381,
+        382,
+        383,
+        384,
+        385,
+        386,
+        387,
+        388,
+        389,
+        390,
+        391,
+        392,
+        393,
+        394,
+        395,
+        396,
+        397,
+        398,
+        399,
+        400,
+        401,
+        402,
+        403,
+        404,
+        405,
+        406,
+        407,
+        408,
+        409,
+        410,
+        411,
+        412,
+        413,
+        414,
+        415,
+        416,
+        417,
+        418,
+        419,
+        420,
+        421,
+        422,
+        427,
+        428,
+        429,
+        430,
+        431,
+        432,
+        433,
+        434,
+        435,
+        436,
+        437,
+        438,
+        439,
+        442,
+        443,
+        444,
+        445,
+        446,
+        447,
+        448,
+        449,
+        450,
+        454,
+        455,
+        456,
+        457,
+        458,
+        459,
+        461,
+        462,
+        463,
+        464,
+        465,
+        466,
+        467,
+        468,
+        469,
+        470,
+        471,
+        472,
+        473,
+        474,
+        475,
+        476,
+        477,
+        478,
+        479,
+        480,
+        481,
+        482,
+        483,
+        484,
+        485,
+        486,
+        487,
+        488,
+        489,
+        490,
+        491,
+        492,
+        493,
+        494,
+        495,
+        496,
+        497,
+        498,
+        499,
+        500,
+        501,
+        502,
+        503,
+        504,
+        505,
+        506,
+        507,
+        508,
+        509,
+        510,
+        511,
+        512,
+        513,
+        514,
+        515,
+        516,
+        517,
+        518,
+        519,
+        520,
+        521,
+        522,
+        523,
+        524,
+        525,
+        526,
+        527,
+        528,
+        529,
+        530,
+        531,
+        532,
+        533,
+        534,
+        535,
+        536,
+        537,
+        538,
+        539,
+        540,
+        541,
+        542,
+        543,
+        544,
+        545,
+        546,
+        547,
+        548,
+        549,
+        550,
+        551,
+        552,
+        553,
+        554,
+        555,
+        556,
+        557,
+        558,
+        559,
+        560,
+        561,
+        562,
+        563,
+        564,
+        565,
+        566,
+        567,
+        568,
+        569,
+        574,
+        575,
+        576,
+        577,
+        578,
+        579,
+        580,
+        581,
+        582,
+        583,
+        1764,
+        1765,
+        1766,
+        1770,
+        1771,
+        1772,
+        1773,
+        1774,
+        1775,
+        1776,
+        1777,
+        1778,
+        1905,
+        1906,
+        1907,
+        1908,
+        2779,
+        2780,
+        2781,
+        2782,
+        2783,
+        2784,
+        2785,
+        2786,
+        2787,
+        2788,
+        2789,
+        2790,
+        2791,
+        2792,
+        2793,
+        2794,
+        2795,
+        2796,
+        2797,
+        2798,
+        2799,
+        2800,
+        2801,
+        2802,
+        2803,
+        2804,
+        2805,
+        2806,
+        2807,
+        2808,
+        2809,
+        2810,
+        2811,
+        2814,
+        2815,
+        2816,
+        2817,
+        2818,
+        3045,
+        3046,
+        3047,
+        3048,
+        3051,
+        3052,
+        3053,
+        3054,
+        3055,
+        3056,
+        3058,
+        3069,
+        3070,
+        3071,
+        3072,
+        3161,
+        3162,
+        3163,
+        3165,
+        3166,
+        3167,
+        3485,
+        3486,
+        3487,
+        3488,
+        3489,
+        3490,
+        3491,
+        3492,
+        3493,
+        3494,
+        3499,
+        3512,
+        3513,
+        3514,
+        3515,
+        3516,
+        3517,
+        3518,
+        3519,
+        3520,
+        3521,
+        3522,
+        3523,
+        3524,
+        3525,
+        3526,
+        3527,
+        3528,
+        3529,
+        3530,
+        3531,
+        3532,
+        3533,
+        3534,
+        3535,
+        3536,
+        3537,
+        3538,
+        3539,
+        3540,
+        3541,
+        3542,
+        3543,
+        3544,
+        3545,
+        3546,
+        3547,
+        3548,
+        3549,
+        3550,
+        3551,
+        3552,
+        3553,
+        3554,
+        3555,
+        3556,
+        3557,
+        3558,
+        3559,
+        3560,
+        3561,
+        3562,
+        3563,
+        3564,
+        3565,
+        3566,
+        3567,
+        3568,
+        3569,
+        3570,
+        3571,
+        3572,
+        3573,
+        3574,
+        3575,
+        3576,
+        3577,
+        3578,
+        3579,
+        3580,
+        3581,
+        3582,
+        3583,
+        3584,
+        3585,
+        3586,
+        3587,
+        3588,
+        3589,
+        3590,
+        3591,
+        3592,
+        3593,
+        3594,
+        3595,
+        3596,
+        3597,
+        3598,
+        3599,
+        3600,
+        3601,
+        3602,
+        3603,
+        3604,
+        3605,
+        3606,
+        3607,
+        3608,
+        3609,
+        3610,
+        3611,
+        3612,
+        3613,
+        3614,
+        3615,
+        3616,
+        3617,
+        3618,
+        3619,
+        3620,
+        3621,
+        3622,
+        3623,
+        3624,
+        3625,
+        3626,
+        3627,
+        3628,
+        3629,
+        3630,
+        3631,
+        3632,
+        3633,
+        3634,
+        3635,
+        3636,
+        3637,
+        3638,
+        3639,
+        3640,
+        3641,
+        3642,
+        3643,
+        3644,
+        3645,
+        3646,
+        3647,
+        3648,
+        3649,
+        3650,
+        3651,
+        3652,
+        3653,
+        3654,
+        3655,
+        3656,
+        3657,
+        3658,
+        3659,
+        3660,
+        3661,
+        3666,
+        3667,
+        3668,
+        3669,
+        3670,
+        3671,
+        3672,
+        3673,
+        3674,
+        3675,
+        3676,
+        3677,
+        3678,
+        3679,
+        3680,
+        3681,
+        3682,
+        3683,
+        3684,
+        3685,
+        3688,
+        3689,
+        3690,
+        3691,
+        3692,
+        3693,
+        3694,
+        3695,
+        3696,
+        3697,
+        3698,
+        3699,
+        3700,
+        3701,
+        3702,
+        3703,
+        3704,
+        3705,
+        3706,
+        3707,
+        3708,
+        3709,
+        3710,
+        3711,
+        3712,
+        3713,
+        3714,
+        3715,
+        3716,
+        3717,
+        3732,
+        3733,
+        3737,
+        3738,
+        3739,
+        3740,
+        3741,
+        3742,
+        3743,
+        3744,
+        3745,
+        3746,
+        3747,
+        3748,
+        3749,
+        3750,
+        3751,
+        3752,
+        3753,
+        3754,
+        3755,
+        3756,
+        3757,
+        3758,
+        3759,
+        3760,
+        3761,
+        3762,
+        3763,
+        3764,
+        3765,
+        3766,
+        3767,
+        3770,
+        3771,
+        3772,
+        3773,
+        3774,
+        3775,
+        3776,
+        3777,
+        3778,
+        3779,
+        3780,
+        3781,
+        3782,
+        3783,
+        3784,
+        3785,
+        3786,
+        3787,
+        3788,
+        3789,
+        3790,
+        3791,
+        3792,
+        3793,
+        3794,
+        3795,
+        3798,
+        3799,
+        3800,
+        3801,
+        3802,
+        3803,
+        3804,
+        3805,
+        3806,
+        3807,
+        3815,
+        3816,
+        3819,
+        3820,
+        3821,
+        3822,
+        3823,
+        3824,
+        3825,
+        3826,
+        3827,
+        3828,
+        3829,
+        3830,
+        3831,
+        3832,
+        3833,
+        3834,
+        3835,
+        3836,
+        3837,
+        3838,
+        3841,
+        3842,
+        3843,
+        3844,
+        3845,
+        3846,
+        3847,
+        3848,
+        3849,
+        3850,
+        3851,
+        3852,
+        3853,
+        3854,
+        3855,
+        3856,
+        3857,
+        3858,
+        3859,
+        3860,
+        3861,
+        3862,
+        3863,
+        3864,
+        3865,
+        3866,
+        3867,
+        3868,
+        3869,
+        3870,
+        3871,
+        3872,
+        3873,
+        3874,
+        3875,
+        3876,
+        3877,
+        3878,
+        3879,
+        3880,
+        3881,
+        3882,
+        3883,
+        3884,
+        3885,
+        3886,
+        3887,
+        3888,
+        3889,
+        3890,
+        3891,
+        3892,
+        3893,
+        3894,
+        3895,
+        3896,
+        3897,
+        3898,
+        3899,
+        3900,
+        3901,
+        3902,
+        3903,
+        3904,
+        3905,
+        3906,
+        3907,
+        3908,
+        3909,
+        3910,
+        3911,
+        3912,
+        3913,
+        3914,
+        3915,
+        3916,
+        3917,
+        3922,
+        3923,
+        3924,
+        3925,
+        3926,
+        3927,
+        3928,
+        3929,
+        3930,
+        3931,
+        3932,
+        3933,
+        3936,
+        3937,
+        3938,
+        3939,
+        3940,
+        3941,
+        3945,
+        3946,
+        3947,
+        3948,
+        3949,
+        3950,
+        3951,
+        3952,
+        3953,
+        3954,
+        3955,
+        3956,
+        3957,
+        3958,
+        3959,
+        3960,
+        3961,
+        3962,
+        3963,
+        3964,
+        3965,
+        3966,
+        3967,
+        3968,
+        3969,
+        3970,
+        3971,
+        3972,
+        3973,
+        3974,
+        3975,
+        3976,
+        3977,
+        3978,
+        3979,
+        3980,
+        3981,
+        3982,
+        3983,
+        3984,
+        3985,
+        3986,
+        3987,
+        3988,
+        3989,
+        3990,
+        3991,
+        3992,
+        3993,
+        3994,
+        3995,
+        3996,
+        3997,
+        3998,
+        3999,
+        4000,
+        4001,
+        4002,
+        4003,
+        4004,
+        4005,
+        4006,
+        4007,
+        4008,
+        4009,
+        4010,
+        4011,
+        4012,
+        4013,
+        4014,
+        4015,
+        4016,
+        4017,
+        4018,
+        4019,
+        4020,
+        4021,
+        4022,
+        4023,
+        4024,
+        4025,
+        4026,
+        4027,
+        4028,
+        4029,
+        4030,
+        4031,
+        4032,
+        4033,
+        4034,
+        4035,
+        4036,
+        4037,
+        4038,
+        4039,
+        4040,
+        4041,
+        4042,
+        4043,
+        4044,
+        4045,
+        4046,
+        4047,
+        4048,
+        4049,
+        4050,
+        4051,
+        4052,
+        4053,
+        4054,
+        4055,
+        4056,
+        4057,
+        4062,
+        4063,
+        4064,
+        4065,
+        4066,
+        4067,
+        4068,
+        4069,
+        4070,
+        4071,
+        5231,
+        5232,
+        5233,
+        5235,
+        5236,
+        5237,
+        5238,
+        5239,
+        5240,
+        5241,
+        5242,
+        5243,
+        5366,
+        5367,
+        5368,
+        5369,
+        6240,
+        6241,
+        6242,
+        6243,
+        6244,
+        6245,
+        6246,
+        6247,
+        6248,
+        6249,
+        6250,
+        6251,
+        6252,
+        6253,
+        6254,
+        6255,
+        6256,
+        6257,
+        6258,
+        6259,
+        6260,
+        6261,
+        6262,
+        6263,
+        6264,
+        6265,
+        6266,
+        6267,
+        6268,
+        6269,
+        6270,
+        6271,
+        6272,
+        6275,
+        6276,
+        6277,
+        6278,
+        6279,
+        6492,
+        6493,
+        6494,
+        6495,
+        6880,
+        6881,
+        6882,
+        6883,
+        6884,
+        6885,
+        6886,
+        6887,
+        6888,
+        6889
+    ],
+    "rightArm": [
+        4114,
+        4115,
+        4116,
+        4117,
+        4122,
+        4125,
+        4168,
+        4171,
+        4204,
+        4205,
+        4206,
+        4207,
+        4257,
+        4258,
+        4259,
+        4260,
+        4261,
+        4262,
+        4263,
+        4264,
+        4265,
+        4266,
+        4267,
+        4268,
+        4272,
+        4273,
+        4274,
+        4275,
+        4276,
+        4277,
+        4278,
+        4279,
+        4280,
+        4281,
+        4714,
+        4715,
+        4716,
+        4717,
+        4741,
+        4742,
+        4743,
+        4744,
+        4756,
+        4763,
+        4764,
+        4790,
+        4791,
+        4794,
+        4795,
+        4816,
+        4817,
+        4818,
+        4819,
+        4830,
+        4831,
+        4832,
+        4833,
+        4849,
+        4850,
+        4851,
+        4852,
+        4853,
+        4854,
+        4855,
+        4856,
+        4857,
+        4858,
+        4859,
+        4860,
+        4861,
+        4862,
+        4863,
+        4864,
+        4865,
+        4866,
+        4867,
+        4868,
+        4869,
+        4870,
+        4871,
+        4872,
+        4873,
+        4876,
+        4877,
+        4878,
+        4879,
+        4880,
+        4881,
+        4882,
+        4883,
+        4884,
+        4885,
+        4886,
+        4887,
+        4888,
+        4889,
+        4901,
+        4902,
+        4903,
+        4904,
+        4905,
+        4906,
+        4911,
+        4912,
+        4913,
+        4914,
+        4915,
+        4916,
+        4917,
+        4918,
+        4974,
+        4977,
+        4978,
+        4979,
+        4980,
+        4981,
+        4982,
+        5009,
+        5010,
+        5011,
+        5012,
+        5014,
+        5088,
+        5089,
+        5090,
+        5091,
+        5100,
+        5101,
+        5102,
+        5103,
+        5104,
+        5105,
+        5106,
+        5107,
+        5108,
+        5109,
+        5110,
+        5111,
+        5114,
+        5115,
+        5116,
+        5117,
+        5118,
+        5119,
+        5120,
+        5121,
+        5122,
+        5123,
+        5124,
+        5125,
+        5128,
+        5129,
+        5130,
+        5131,
+        5134,
+        5135,
+        5136,
+        5137,
+        5138,
+        5139,
+        5140,
+        5141,
+        5142,
+        5143,
+        5144,
+        5145,
+        5146,
+        5147,
+        5148,
+        5149,
+        5150,
+        5151,
+        5152,
+        5153,
+        5165,
+        5166,
+        5167,
+        5172,
+        5173,
+        5174,
+        5175,
+        5176,
+        5177,
+        5178,
+        5179,
+        5180,
+        5181,
+        5182,
+        5183,
+        5184,
+        5185,
+        5186,
+        5187,
+        5188,
+        5189,
+        5194,
+        5200,
+        5201,
+        5202,
+        5203,
+        5204,
+        5206,
+        5208,
+        5209,
+        5214,
+        5215,
+        5216,
+        5217,
+        5218,
+        5220,
+        5229,
+        5292,
+        5293,
+        5303,
+        5306,
+        5309,
+        5311,
+        5314,
+        5315,
+        5318,
+        5319,
+        5321,
+        5326,
+        5327,
+        5328,
+        5330,
+        5331,
+        5332,
+        5335,
+        5336,
+        5337,
+        5338,
+        5339,
+        5343,
+        5344,
+        5349,
+        5350,
+        5353,
+        5361,
+        5362,
+        5363,
+        5364,
+        5365,
+        5370,
+        6280,
+        6281,
+        6282,
+        6283,
+        6354,
+        6355,
+        6356,
+        6357,
+        6358,
+        6359,
+        6360,
+        6361,
+        6362,
+        6404,
+        6405,
+        6433,
+        6434,
+        6435,
+        6436,
+        6437,
+        6438,
+        6439,
+        6440,
+        6441,
+        6442,
+        6443,
+        6444,
+        6445,
+        6446,
+        6447,
+        6448,
+        6449,
+        6450,
+        6451,
+        6452,
+        6453,
+        6454,
+        6455,
+        6461,
+        6471
+    ],
+    "leftHandIndex1": [
+        2027,
+        2028,
+        2029,
+        2030,
+        2037,
+        2038,
+        2039,
+        2040,
+        2057,
+        2067,
+        2068,
+        2123,
+        2124,
+        2125,
+        2126,
+        2127,
+        2128,
+        2129,
+        2130,
+        2132,
+        2145,
+        2146,
+        2152,
+        2153,
+        2154,
+        2156,
+        2157,
+        2158,
+        2159,
+        2160,
+        2161,
+        2162,
+        2163,
+        2164,
+        2165,
+        2166,
+        2167,
+        2168,
+        2169,
+        2177,
+        2178,
+        2179,
+        2181,
+        2186,
+        2187,
+        2190,
+        2191,
+        2204,
+        2205,
+        2215,
+        2216,
+        2217,
+        2218,
+        2219,
+        2220,
+        2232,
+        2233,
+        2245,
+        2246,
+        2247,
+        2258,
+        2259,
+        2261,
+        2262,
+        2263,
+        2269,
+        2270,
+        2272,
+        2273,
+        2274,
+        2276,
+        2277,
+        2280,
+        2281,
+        2282,
+        2283,
+        2291,
+        2292,
+        2293,
+        2294,
+        2295,
+        2296,
+        2297,
+        2298,
+        2299,
+        2300,
+        2301,
+        2302,
+        2303,
+        2304,
+        2305,
+        2306,
+        2307,
+        2308,
+        2309,
+        2310,
+        2311,
+        2312,
+        2313,
+        2314,
+        2315,
+        2316,
+        2317,
+        2318,
+        2319,
+        2320,
+        2321,
+        2322,
+        2323,
+        2324,
+        2325,
+        2326,
+        2327,
+        2328,
+        2329,
+        2330,
+        2331,
+        2332,
+        2333,
+        2334,
+        2335,
+        2336,
+        2337,
+        2338,
+        2339,
+        2340,
+        2341,
+        2342,
+        2343,
+        2344,
+        2345,
+        2346,
+        2347,
+        2348,
+        2349,
+        2350,
+        2351,
+        2352,
+        2353,
+        2354,
+        2355,
+        2356,
+        2357,
+        2358,
+        2359,
+        2360,
+        2361,
+        2362,
+        2363,
+        2364,
+        2365,
+        2366,
+        2367,
+        2368,
+        2369,
+        2370,
+        2371,
+        2372,
+        2373,
+        2374,
+        2375,
+        2376,
+        2377,
+        2378,
+        2379,
+        2380,
+        2381,
+        2382,
+        2383,
+        2384,
+        2385,
+        2386,
+        2387,
+        2388,
+        2389,
+        2390,
+        2391,
+        2392,
+        2393,
+        2394,
+        2395,
+        2396,
+        2397,
+        2398,
+        2399,
+        2400,
+        2401,
+        2402,
+        2403,
+        2404,
+        2405,
+        2406,
+        2407,
+        2408,
+        2409,
+        2410,
+        2411,
+        2412,
+        2413,
+        2414,
+        2415,
+        2416,
+        2417,
+        2418,
+        2419,
+        2420,
+        2421,
+        2422,
+        2423,
+        2424,
+        2425,
+        2426,
+        2427,
+        2428,
+        2429,
+        2430,
+        2431,
+        2432,
+        2433,
+        2434,
+        2435,
+        2436,
+        2437,
+        2438,
+        2439,
+        2440,
+        2441,
+        2442,
+        2443,
+        2444,
+        2445,
+        2446,
+        2447,
+        2448,
+        2449,
+        2450,
+        2451,
+        2452,
+        2453,
+        2454,
+        2455,
+        2456,
+        2457,
+        2458,
+        2459,
+        2460,
+        2461,
+        2462,
+        2463,
+        2464,
+        2465,
+        2466,
+        2467,
+        2468,
+        2469,
+        2470,
+        2471,
+        2472,
+        2473,
+        2474,
+        2475,
+        2476,
+        2477,
+        2478,
+        2479,
+        2480,
+        2481,
+        2482,
+        2483,
+        2484,
+        2485,
+        2486,
+        2487,
+        2488,
+        2489,
+        2490,
+        2491,
+        2492,
+        2493,
+        2494,
+        2495,
+        2496,
+        2497,
+        2498,
+        2499,
+        2500,
+        2501,
+        2502,
+        2503,
+        2504,
+        2505,
+        2506,
+        2507,
+        2508,
+        2509,
+        2510,
+        2511,
+        2512,
+        2513,
+        2514,
+        2515,
+        2516,
+        2517,
+        2518,
+        2519,
+        2520,
+        2521,
+        2522,
+        2523,
+        2524,
+        2525,
+        2526,
+        2527,
+        2528,
+        2529,
+        2530,
+        2531,
+        2532,
+        2533,
+        2534,
+        2535,
+        2536,
+        2537,
+        2538,
+        2539,
+        2540,
+        2541,
+        2542,
+        2543,
+        2544,
+        2545,
+        2546,
+        2547,
+        2548,
+        2549,
+        2550,
+        2551,
+        2552,
+        2553,
+        2554,
+        2555,
+        2556,
+        2557,
+        2558,
+        2559,
+        2560,
+        2561,
+        2562,
+        2563,
+        2564,
+        2565,
+        2566,
+        2567,
+        2568,
+        2569,
+        2570,
+        2571,
+        2572,
+        2573,
+        2574,
+        2575,
+        2576,
+        2577,
+        2578,
+        2579,
+        2580,
+        2581,
+        2582,
+        2583,
+        2584,
+        2585,
+        2586,
+        2587,
+        2588,
+        2589,
+        2590,
+        2591,
+        2592,
+        2593,
+        2594,
+        2596,
+        2597,
+        2599,
+        2600,
+        2601,
+        2602,
+        2603,
+        2604,
+        2606,
+        2607,
+        2609,
+        2610,
+        2611,
+        2612,
+        2613,
+        2614,
+        2615,
+        2616,
+        2617,
+        2618,
+        2619,
+        2620,
+        2621,
+        2622,
+        2623,
+        2624,
+        2625,
+        2626,
+        2627,
+        2628,
+        2629,
+        2630,
+        2631,
+        2632,
+        2633,
+        2634,
+        2635,
+        2636,
+        2637,
+        2638,
+        2639,
+        2640,
+        2641,
+        2642,
+        2643,
+        2644,
+        2645,
+        2646,
+        2647,
+        2648,
+        2649,
+        2650,
+        2651,
+        2652,
+        2653,
+        2654,
+        2655,
+        2656,
+        2657,
+        2658,
+        2659,
+        2660,
+        2661,
+        2662,
+        2663,
+        2664,
+        2665,
+        2666,
+        2667,
+        2668,
+        2669,
+        2670,
+        2671,
+        2672,
+        2673,
+        2674,
+        2675,
+        2676,
+        2677,
+        2678,
+        2679,
+        2680,
+        2681,
+        2682,
+        2683,
+        2684,
+        2685,
+        2686,
+        2687,
+        2688,
+        2689,
+        2690,
+        2691,
+        2692,
+        2693,
+        2694,
+        2695,
+        2696
+    ],
+    "rightLeg": [
+        4481,
+        4482,
+        4485,
+        4486,
+        4491,
+        4492,
+        4493,
+        4495,
+        4498,
+        4500,
+        4501,
+        4505,
+        4506,
+        4529,
+        4532,
+        4533,
+        4534,
+        4535,
+        4536,
+        4537,
+        4538,
+        4539,
+        4540,
+        4541,
+        4542,
+        4543,
+        4544,
+        4545,
+        4546,
+        4547,
+        4548,
+        4549,
+        4550,
+        4551,
+        4552,
+        4553,
+        4554,
+        4555,
+        4556,
+        4557,
+        4558,
+        4559,
+        4560,
+        4561,
+        4562,
+        4563,
+        4564,
+        4565,
+        4566,
+        4567,
+        4568,
+        4569,
+        4570,
+        4571,
+        4572,
+        4573,
+        4574,
+        4575,
+        4576,
+        4577,
+        4578,
+        4579,
+        4580,
+        4581,
+        4582,
+        4583,
+        4584,
+        4585,
+        4586,
+        4587,
+        4588,
+        4589,
+        4590,
+        4591,
+        4592,
+        4593,
+        4594,
+        4595,
+        4596,
+        4597,
+        4598,
+        4599,
+        4600,
+        4601,
+        4602,
+        4603,
+        4604,
+        4605,
+        4606,
+        4607,
+        4608,
+        4609,
+        4610,
+        4611,
+        4612,
+        4613,
+        4614,
+        4615,
+        4616,
+        4617,
+        4618,
+        4619,
+        4620,
+        4621,
+        4622,
+        4634,
+        4635,
+        4636,
+        4637,
+        4638,
+        4639,
+        4640,
+        4641,
+        4642,
+        4643,
+        4644,
+        4661,
+        4662,
+        4663,
+        4664,
+        4665,
+        4666,
+        4667,
+        4668,
+        4669,
+        4842,
+        4843,
+        4844,
+        4845,
+        4846,
+        4847,
+        4848,
+        4937,
+        4938,
+        4939,
+        4940,
+        4941,
+        4942,
+        4943,
+        4944,
+        4945,
+        4946,
+        4947,
+        4993,
+        4994,
+        4995,
+        4996,
+        4997,
+        4998,
+        4999,
+        5000,
+        5001,
+        5002,
+        5003,
+        6574,
+        6575,
+        6576,
+        6577,
+        6578,
+        6579,
+        6580,
+        6581,
+        6582,
+        6583,
+        6584,
+        6585,
+        6586,
+        6587,
+        6588,
+        6589,
+        6590,
+        6591,
+        6592,
+        6593,
+        6594,
+        6595,
+        6596,
+        6597,
+        6598,
+        6599,
+        6600,
+        6601,
+        6602,
+        6603,
+        6604,
+        6605,
+        6606,
+        6607,
+        6608,
+        6609,
+        6610,
+        6719,
+        6720,
+        6721,
+        6722,
+        6723,
+        6724,
+        6725,
+        6726,
+        6727,
+        6728,
+        6729,
+        6730,
+        6731,
+        6732,
+        6733,
+        6734,
+        6735,
+        6832,
+        6833,
+        6834,
+        6835,
+        6836,
+        6869,
+        6870,
+        6871,
+        6872
+    ],
+    "rightHandIndex1": [
+        5488,
+        5489,
+        5490,
+        5491,
+        5498,
+        5499,
+        5500,
+        5501,
+        5518,
+        5528,
+        5529,
+        5584,
+        5585,
+        5586,
+        5587,
+        5588,
+        5589,
+        5590,
+        5591,
+        5592,
+        5606,
+        5607,
+        5613,
+        5615,
+        5616,
+        5617,
+        5618,
+        5619,
+        5620,
+        5621,
+        5622,
+        5623,
+        5624,
+        5625,
+        5626,
+        5627,
+        5628,
+        5629,
+        5630,
+        5638,
+        5639,
+        5640,
+        5642,
+        5647,
+        5648,
+        5650,
+        5651,
+        5665,
+        5666,
+        5676,
+        5677,
+        5678,
+        5679,
+        5680,
+        5681,
+        5693,
+        5694,
+        5706,
+        5707,
+        5708,
+        5719,
+        5721,
+        5722,
+        5723,
+        5724,
+        5730,
+        5731,
+        5733,
+        5734,
+        5735,
+        5737,
+        5738,
+        5741,
+        5742,
+        5743,
+        5744,
+        5752,
+        5753,
+        5754,
+        5755,
+        5756,
+        5757,
+        5758,
+        5759,
+        5760,
+        5761,
+        5762,
+        5763,
+        5764,
+        5765,
+        5766,
+        5767,
+        5768,
+        5769,
+        5770,
+        5771,
+        5772,
+        5773,
+        5774,
+        5775,
+        5776,
+        5777,
+        5778,
+        5779,
+        5780,
+        5781,
+        5782,
+        5783,
+        5784,
+        5785,
+        5786,
+        5787,
+        5788,
+        5789,
+        5790,
+        5791,
+        5792,
+        5793,
+        5794,
+        5795,
+        5796,
+        5797,
+        5798,
+        5799,
+        5800,
+        5801,
+        5802,
+        5803,
+        5804,
+        5805,
+        5806,
+        5807,
+        5808,
+        5809,
+        5810,
+        5811,
+        5812,
+        5813,
+        5814,
+        5815,
+        5816,
+        5817,
+        5818,
+        5819,
+        5820,
+        5821,
+        5822,
+        5823,
+        5824,
+        5825,
+        5826,
+        5827,
+        5828,
+        5829,
+        5830,
+        5831,
+        5832,
+        5833,
+        5834,
+        5835,
+        5836,
+        5837,
+        5838,
+        5839,
+        5840,
+        5841,
+        5842,
+        5843,
+        5844,
+        5845,
+        5846,
+        5847,
+        5848,
+        5849,
+        5850,
+        5851,
+        5852,
+        5853,
+        5854,
+        5855,
+        5856,
+        5857,
+        5858,
+        5859,
+        5860,
+        5861,
+        5862,
+        5863,
+        5864,
+        5865,
+        5866,
+        5867,
+        5868,
+        5869,
+        5870,
+        5871,
+        5872,
+        5873,
+        5874,
+        5875,
+        5876,
+        5877,
+        5878,
+        5879,
+        5880,
+        5881,
+        5882,
+        5883,
+        5884,
+        5885,
+        5886,
+        5887,
+        5888,
+        5889,
+        5890,
+        5891,
+        5892,
+        5893,
+        5894,
+        5895,
+        5896,
+        5897,
+        5898,
+        5899,
+        5900,
+        5901,
+        5902,
+        5903,
+        5904,
+        5905,
+        5906,
+        5907,
+        5908,
+        5909,
+        5910,
+        5911,
+        5912,
+        5913,
+        5914,
+        5915,
+        5916,
+        5917,
+        5918,
+        5919,
+        5920,
+        5921,
+        5922,
+        5923,
+        5924,
+        5925,
+        5926,
+        5927,
+        5928,
+        5929,
+        5930,
+        5931,
+        5932,
+        5933,
+        5934,
+        5935,
+        5936,
+        5937,
+        5938,
+        5939,
+        5940,
+        5941,
+        5942,
+        5943,
+        5944,
+        5945,
+        5946,
+        5947,
+        5948,
+        5949,
+        5950,
+        5951,
+        5952,
+        5953,
+        5954,
+        5955,
+        5956,
+        5957,
+        5958,
+        5959,
+        5960,
+        5961,
+        5962,
+        5963,
+        5964,
+        5965,
+        5966,
+        5967,
+        5968,
+        5969,
+        5970,
+        5971,
+        5972,
+        5973,
+        5974,
+        5975,
+        5976,
+        5977,
+        5978,
+        5979,
+        5980,
+        5981,
+        5982,
+        5983,
+        5984,
+        5985,
+        5986,
+        5987,
+        5988,
+        5989,
+        5990,
+        5991,
+        5992,
+        5993,
+        5994,
+        5995,
+        5996,
+        5997,
+        5998,
+        5999,
+        6000,
+        6001,
+        6002,
+        6003,
+        6004,
+        6005,
+        6006,
+        6007,
+        6008,
+        6009,
+        6010,
+        6011,
+        6012,
+        6013,
+        6014,
+        6015,
+        6016,
+        6017,
+        6018,
+        6019,
+        6020,
+        6021,
+        6022,
+        6023,
+        6024,
+        6025,
+        6026,
+        6027,
+        6028,
+        6029,
+        6030,
+        6031,
+        6032,
+        6033,
+        6034,
+        6035,
+        6036,
+        6037,
+        6038,
+        6039,
+        6040,
+        6041,
+        6042,
+        6043,
+        6044,
+        6045,
+        6046,
+        6047,
+        6048,
+        6049,
+        6050,
+        6051,
+        6052,
+        6053,
+        6054,
+        6055,
+        6058,
+        6059,
+        6060,
+        6061,
+        6062,
+        6063,
+        6064,
+        6065,
+        6068,
+        6069,
+        6070,
+        6071,
+        6072,
+        6073,
+        6074,
+        6075,
+        6076,
+        6077,
+        6078,
+        6079,
+        6080,
+        6081,
+        6082,
+        6083,
+        6084,
+        6085,
+        6086,
+        6087,
+        6088,
+        6089,
+        6090,
+        6091,
+        6092,
+        6093,
+        6094,
+        6095,
+        6096,
+        6097,
+        6098,
+        6099,
+        6100,
+        6101,
+        6102,
+        6103,
+        6104,
+        6105,
+        6106,
+        6107,
+        6108,
+        6109,
+        6110,
+        6111,
+        6112,
+        6113,
+        6114,
+        6115,
+        6116,
+        6117,
+        6118,
+        6119,
+        6120,
+        6121,
+        6122,
+        6123,
+        6124,
+        6125,
+        6126,
+        6127,
+        6128,
+        6129,
+        6130,
+        6131,
+        6132,
+        6133,
+        6134,
+        6135,
+        6136,
+        6137,
+        6138,
+        6139,
+        6140,
+        6141,
+        6142,
+        6143,
+        6144,
+        6145,
+        6146,
+        6147,
+        6148,
+        6149,
+        6150,
+        6151,
+        6152,
+        6153,
+        6154,
+        6155,
+        6156,
+        6157
+    ],
+    "leftForeArm": [
+        1546,
+        1547,
+        1548,
+        1549,
+        1550,
+        1551,
+        1552,
+        1553,
+        1554,
+        1555,
+        1556,
+        1557,
+        1558,
+        1559,
+        1560,
+        1561,
+        1562,
+        1563,
+        1564,
+        1565,
+        1566,
+        1567,
+        1568,
+        1569,
+        1570,
+        1571,
+        1572,
+        1573,
+        1574,
+        1575,
+        1576,
+        1577,
+        1578,
+        1579,
+        1580,
+        1581,
+        1582,
+        1583,
+        1584,
+        1585,
+        1586,
+        1587,
+        1588,
+        1589,
+        1590,
+        1591,
+        1592,
+        1593,
+        1594,
+        1595,
+        1596,
+        1597,
+        1598,
+        1599,
+        1600,
+        1601,
+        1602,
+        1603,
+        1604,
+        1605,
+        1606,
+        1607,
+        1608,
+        1609,
+        1610,
+        1611,
+        1612,
+        1613,
+        1614,
+        1615,
+        1616,
+        1617,
+        1618,
+        1620,
+        1621,
+        1623,
+        1624,
+        1625,
+        1626,
+        1627,
+        1628,
+        1629,
+        1630,
+        1643,
+        1644,
+        1646,
+        1647,
+        1650,
+        1651,
+        1654,
+        1655,
+        1657,
+        1658,
+        1659,
+        1660,
+        1661,
+        1662,
+        1663,
+        1664,
+        1665,
+        1666,
+        1685,
+        1686,
+        1687,
+        1688,
+        1689,
+        1690,
+        1691,
+        1692,
+        1693,
+        1694,
+        1695,
+        1699,
+        1700,
+        1701,
+        1702,
+        1721,
+        1722,
+        1723,
+        1724,
+        1725,
+        1726,
+        1727,
+        1728,
+        1729,
+        1730,
+        1732,
+        1736,
+        1738,
+        1741,
+        1742,
+        1743,
+        1744,
+        1750,
+        1752,
+        1900,
+        1909,
+        1910,
+        1911,
+        1912,
+        1913,
+        1914,
+        1915,
+        1916,
+        1917,
+        1918,
+        1919,
+        1920,
+        1921,
+        1922,
+        1923,
+        1924,
+        1925,
+        1926,
+        1927,
+        1928,
+        1929,
+        1930,
+        1931,
+        1932,
+        1933,
+        1934,
+        1935,
+        1936,
+        1937,
+        1938,
+        1939,
+        1940,
+        1941,
+        1942,
+        1943,
+        1944,
+        1945,
+        1946,
+        1947,
+        1948,
+        1949,
+        1950,
+        1951,
+        1952,
+        1953,
+        1954,
+        1955,
+        1956,
+        1957,
+        1958,
+        1959,
+        1960,
+        1961,
+        1962,
+        1963,
+        1964,
+        1965,
+        1966,
+        1967,
+        1968,
+        1969,
+        1970,
+        1971,
+        1972,
+        1973,
+        1974,
+        1975,
+        1976,
+        1977,
+        1978,
+        1979,
+        1980,
+        2019,
+        2059,
+        2060,
+        2073,
+        2089,
+        2098,
+        2099,
+        2100,
+        2101,
+        2102,
+        2103,
+        2104,
+        2105,
+        2106,
+        2107,
+        2108,
+        2109,
+        2110,
+        2111,
+        2112,
+        2147,
+        2148,
+        2206,
+        2207,
+        2208,
+        2209,
+        2228,
+        2230,
+        2234,
+        2235,
+        2241,
+        2242,
+        2243,
+        2244,
+        2279,
+        2286,
+        2873,
+        2874
+    ],
+    "rightForeArm": [
+        5015,
+        5016,
+        5017,
+        5018,
+        5019,
+        5020,
+        5021,
+        5022,
+        5023,
+        5024,
+        5025,
+        5026,
+        5027,
+        5028,
+        5029,
+        5030,
+        5031,
+        5032,
+        5033,
+        5034,
+        5035,
+        5036,
+        5037,
+        5038,
+        5039,
+        5040,
+        5041,
+        5042,
+        5043,
+        5044,
+        5045,
+        5046,
+        5047,
+        5048,
+        5049,
+        5050,
+        5051,
+        5052,
+        5053,
+        5054,
+        5055,
+        5056,
+        5057,
+        5058,
+        5059,
+        5060,
+        5061,
+        5062,
+        5063,
+        5064,
+        5065,
+        5066,
+        5067,
+        5068,
+        5069,
+        5070,
+        5071,
+        5072,
+        5073,
+        5074,
+        5075,
+        5076,
+        5077,
+        5078,
+        5079,
+        5080,
+        5081,
+        5082,
+        5083,
+        5084,
+        5085,
+        5086,
+        5087,
+        5090,
+        5091,
+        5092,
+        5093,
+        5094,
+        5095,
+        5096,
+        5097,
+        5098,
+        5099,
+        5112,
+        5113,
+        5116,
+        5117,
+        5120,
+        5121,
+        5124,
+        5125,
+        5126,
+        5127,
+        5128,
+        5129,
+        5130,
+        5131,
+        5132,
+        5133,
+        5134,
+        5135,
+        5154,
+        5155,
+        5156,
+        5157,
+        5158,
+        5159,
+        5160,
+        5161,
+        5162,
+        5163,
+        5164,
+        5168,
+        5169,
+        5170,
+        5171,
+        5190,
+        5191,
+        5192,
+        5193,
+        5194,
+        5195,
+        5196,
+        5197,
+        5198,
+        5199,
+        5202,
+        5205,
+        5207,
+        5210,
+        5211,
+        5212,
+        5213,
+        5219,
+        5221,
+        5361,
+        5370,
+        5371,
+        5372,
+        5373,
+        5374,
+        5375,
+        5376,
+        5377,
+        5378,
+        5379,
+        5380,
+        5381,
+        5382,
+        5383,
+        5384,
+        5385,
+        5386,
+        5387,
+        5388,
+        5389,
+        5390,
+        5391,
+        5392,
+        5393,
+        5394,
+        5395,
+        5396,
+        5397,
+        5398,
+        5399,
+        5400,
+        5401,
+        5402,
+        5403,
+        5404,
+        5405,
+        5406,
+        5407,
+        5408,
+        5409,
+        5410,
+        5411,
+        5412,
+        5413,
+        5414,
+        5415,
+        5416,
+        5417,
+        5418,
+        5419,
+        5420,
+        5421,
+        5422,
+        5423,
+        5424,
+        5425,
+        5426,
+        5427,
+        5428,
+        5429,
+        5430,
+        5431,
+        5432,
+        5433,
+        5434,
+        5435,
+        5436,
+        5437,
+        5438,
+        5439,
+        5440,
+        5441,
+        5480,
+        5520,
+        5521,
+        5534,
+        5550,
+        5559,
+        5560,
+        5561,
+        5562,
+        5563,
+        5564,
+        5565,
+        5566,
+        5567,
+        5568,
+        5569,
+        5570,
+        5571,
+        5572,
+        5573,
+        5608,
+        5609,
+        5667,
+        5668,
+        5669,
+        5670,
+        5689,
+        5691,
+        5695,
+        5696,
+        5702,
+        5703,
+        5704,
+        5705,
+        5740,
+        5747,
+        6334,
+        6335
+    ],
+    "neck": [
+        148,
+        150,
+        151,
+        152,
+        153,
+        172,
+        174,
+        175,
+        201,
+        202,
+        204,
+        205,
+        206,
+        207,
+        208,
+        209,
+        210,
+        211,
+        212,
+        213,
+        214,
+        215,
+        216,
+        217,
+        218,
+        219,
+        222,
+        223,
+        224,
+        225,
+        256,
+        257,
+        284,
+        285,
+        295,
+        296,
+        297,
+        298,
+        299,
+        300,
+        301,
+        302,
+        303,
+        304,
+        305,
+        306,
+        307,
+        308,
+        309,
+        333,
+        334,
+        423,
+        424,
+        425,
+        426,
+        440,
+        441,
+        451,
+        452,
+        453,
+        460,
+        461,
+        571,
+        572,
+        824,
+        825,
+        826,
+        827,
+        828,
+        829,
+        1279,
+        1280,
+        1312,
+        1313,
+        1319,
+        1320,
+        1331,
+        3049,
+        3050,
+        3057,
+        3058,
+        3059,
+        3068,
+        3164,
+        3661,
+        3662,
+        3663,
+        3664,
+        3665,
+        3685,
+        3686,
+        3687,
+        3714,
+        3715,
+        3716,
+        3717,
+        3718,
+        3719,
+        3720,
+        3721,
+        3722,
+        3723,
+        3724,
+        3725,
+        3726,
+        3727,
+        3728,
+        3729,
+        3730,
+        3731,
+        3734,
+        3735,
+        3736,
+        3737,
+        3768,
+        3769,
+        3796,
+        3797,
+        3807,
+        3808,
+        3809,
+        3810,
+        3811,
+        3812,
+        3813,
+        3814,
+        3815,
+        3816,
+        3817,
+        3818,
+        3819,
+        3839,
+        3840,
+        3918,
+        3919,
+        3920,
+        3921,
+        3934,
+        3935,
+        3942,
+        3943,
+        3944,
+        3950,
+        4060,
+        4061,
+        4312,
+        4313,
+        4314,
+        4315,
+        4761,
+        4762,
+        4792,
+        4793,
+        4799,
+        4800,
+        4807
+    ],
+    "rightToeBase": [
+        6611,
+        6612,
+        6613,
+        6614,
+        6615,
+        6616,
+        6617,
+        6618,
+        6619,
+        6620,
+        6621,
+        6622,
+        6623,
+        6624,
+        6625,
+        6626,
+        6627,
+        6628,
+        6629,
+        6630,
+        6631,
+        6632,
+        6633,
+        6634,
+        6635,
+        6636,
+        6637,
+        6638,
+        6639,
+        6640,
+        6641,
+        6642,
+        6643,
+        6644,
+        6645,
+        6646,
+        6647,
+        6648,
+        6649,
+        6650,
+        6651,
+        6652,
+        6653,
+        6654,
+        6655,
+        6656,
+        6657,
+        6658,
+        6659,
+        6660,
+        6661,
+        6662,
+        6663,
+        6664,
+        6665,
+        6666,
+        6667,
+        6668,
+        6669,
+        6670,
+        6671,
+        6672,
+        6673,
+        6674,
+        6675,
+        6676,
+        6677,
+        6678,
+        6679,
+        6680,
+        6681,
+        6682,
+        6683,
+        6684,
+        6685,
+        6686,
+        6687,
+        6688,
+        6689,
+        6690,
+        6691,
+        6692,
+        6693,
+        6694,
+        6695,
+        6696,
+        6697,
+        6698,
+        6699,
+        6700,
+        6701,
+        6702,
+        6703,
+        6704,
+        6705,
+        6706,
+        6707,
+        6708,
+        6709,
+        6710,
+        6711,
+        6712,
+        6713,
+        6714,
+        6715,
+        6716,
+        6717,
+        6718,
+        6736,
+        6739,
+        6741,
+        6743,
+        6745,
+        6747,
+        6749,
+        6750,
+        6752,
+        6754,
+        6757,
+        6758,
+        6760,
+        6762
+    ],
+    "spine": [
+        616,
+        617,
+        630,
+        631,
+        632,
+        633,
+        654,
+        655,
+        656,
+        657,
+        662,
+        663,
+        664,
+        665,
+        720,
+        721,
+        765,
+        766,
+        767,
+        768,
+        796,
+        797,
+        798,
+        799,
+        889,
+        890,
+        916,
+        917,
+        918,
+        919,
+        921,
+        922,
+        923,
+        924,
+        925,
+        926,
+        1188,
+        1189,
+        1211,
+        1212,
+        1248,
+        1249,
+        1250,
+        1251,
+        1264,
+        1265,
+        1266,
+        1267,
+        1323,
+        1324,
+        1325,
+        1326,
+        1327,
+        1328,
+        1332,
+        1333,
+        1334,
+        1335,
+        1336,
+        1344,
+        1345,
+        1481,
+        1482,
+        1483,
+        1484,
+        1485,
+        1486,
+        1487,
+        1488,
+        1489,
+        1490,
+        1491,
+        1492,
+        1493,
+        1494,
+        1495,
+        1496,
+        1767,
+        2823,
+        2824,
+        2825,
+        2826,
+        2827,
+        2828,
+        2829,
+        2830,
+        2831,
+        2832,
+        2833,
+        2834,
+        2835,
+        2836,
+        2837,
+        2838,
+        2839,
+        2840,
+        2841,
+        2842,
+        2843,
+        2844,
+        2845,
+        2847,
+        2848,
+        2851,
+        3016,
+        3017,
+        3018,
+        3019,
+        3020,
+        3023,
+        3024,
+        3124,
+        3173,
+        3476,
+        3477,
+        3478,
+        3480,
+        3500,
+        3501,
+        3502,
+        3504,
+        3509,
+        3511,
+        4103,
+        4104,
+        4118,
+        4119,
+        4120,
+        4121,
+        4142,
+        4143,
+        4144,
+        4145,
+        4150,
+        4151,
+        4152,
+        4153,
+        4208,
+        4209,
+        4253,
+        4254,
+        4255,
+        4256,
+        4284,
+        4285,
+        4286,
+        4287,
+        4375,
+        4376,
+        4402,
+        4403,
+        4405,
+        4406,
+        4407,
+        4408,
+        4409,
+        4410,
+        4411,
+        4412,
+        4674,
+        4675,
+        4694,
+        4695,
+        4731,
+        4732,
+        4733,
+        4734,
+        4747,
+        4748,
+        4749,
+        4750,
+        4803,
+        4804,
+        4805,
+        4806,
+        4808,
+        4809,
+        4810,
+        4811,
+        4812,
+        4820,
+        4821,
+        4953,
+        4954,
+        4955,
+        4956,
+        4957,
+        4958,
+        4959,
+        4960,
+        4961,
+        4962,
+        4963,
+        4964,
+        4965,
+        4966,
+        4967,
+        4968,
+        5234,
+        6284,
+        6285,
+        6286,
+        6287,
+        6288,
+        6289,
+        6290,
+        6291,
+        6292,
+        6293,
+        6294,
+        6295,
+        6296,
+        6297,
+        6298,
+        6299,
+        6300,
+        6301,
+        6302,
+        6303,
+        6304,
+        6305,
+        6306,
+        6308,
+        6309,
+        6312,
+        6472,
+        6473,
+        6474,
+        6545,
+        6874,
+        6875,
+        6876,
+        6878
+    ],
+    "leftUpLeg": [
+        833,
+        834,
+        838,
+        839,
+        847,
+        848,
+        849,
+        850,
+        851,
+        852,
+        853,
+        854,
+        870,
+        871,
+        872,
+        873,
+        874,
+        875,
+        876,
+        877,
+        878,
+        879,
+        880,
+        881,
+        897,
+        898,
+        899,
+        900,
+        901,
+        902,
+        903,
+        904,
+        905,
+        906,
+        907,
+        908,
+        909,
+        910,
+        911,
+        912,
+        913,
+        914,
+        915,
+        933,
+        934,
+        935,
+        936,
+        944,
+        945,
+        946,
+        947,
+        948,
+        949,
+        950,
+        951,
+        952,
+        953,
+        954,
+        955,
+        956,
+        957,
+        958,
+        959,
+        960,
+        961,
+        962,
+        963,
+        964,
+        965,
+        966,
+        967,
+        968,
+        969,
+        970,
+        971,
+        972,
+        973,
+        974,
+        975,
+        976,
+        977,
+        978,
+        979,
+        980,
+        981,
+        982,
+        983,
+        984,
+        985,
+        986,
+        987,
+        988,
+        989,
+        990,
+        991,
+        992,
+        993,
+        994,
+        995,
+        996,
+        997,
+        998,
+        999,
+        1000,
+        1001,
+        1002,
+        1003,
+        1004,
+        1005,
+        1006,
+        1007,
+        1008,
+        1009,
+        1010,
+        1011,
+        1012,
+        1013,
+        1014,
+        1015,
+        1016,
+        1017,
+        1018,
+        1019,
+        1020,
+        1021,
+        1022,
+        1023,
+        1024,
+        1025,
+        1026,
+        1027,
+        1028,
+        1029,
+        1030,
+        1031,
+        1032,
+        1033,
+        1034,
+        1035,
+        1036,
+        1037,
+        1038,
+        1039,
+        1040,
+        1041,
+        1042,
+        1043,
+        1044,
+        1045,
+        1046,
+        1137,
+        1138,
+        1139,
+        1140,
+        1141,
+        1142,
+        1143,
+        1144,
+        1145,
+        1146,
+        1147,
+        1148,
+        1159,
+        1160,
+        1161,
+        1162,
+        1163,
+        1164,
+        1165,
+        1166,
+        1167,
+        1168,
+        1169,
+        1170,
+        1171,
+        1172,
+        1173,
+        1174,
+        1184,
+        1185,
+        1186,
+        1187,
+        1221,
+        1222,
+        1223,
+        1224,
+        1225,
+        1226,
+        1227,
+        1228,
+        1229,
+        1230,
+        1262,
+        1263,
+        1274,
+        1275,
+        1276,
+        1277,
+        1321,
+        1322,
+        1354,
+        1359,
+        1360,
+        1361,
+        1362,
+        1365,
+        1366,
+        1367,
+        1368,
+        1451,
+        1452,
+        1453,
+        1455,
+        1456,
+        1457,
+        1458,
+        1459,
+        1460,
+        1461,
+        1462,
+        1463,
+        1475,
+        1477,
+        1478,
+        1479,
+        1480,
+        1498,
+        1499,
+        1500,
+        1501,
+        1511,
+        1512,
+        1513,
+        1514,
+        1516,
+        1517,
+        1518,
+        1519,
+        1520,
+        1521,
+        1522,
+        1533,
+        1534,
+        3125,
+        3126,
+        3127,
+        3128,
+        3131,
+        3132,
+        3133,
+        3134,
+        3135,
+        3475,
+        3479
+    ],
+    "leftHand": [
+        1981,
+        1982,
+        1983,
+        1984,
+        1985,
+        1986,
+        1987,
+        1988,
+        1989,
+        1990,
+        1991,
+        1992,
+        1993,
+        1994,
+        1995,
+        1996,
+        1997,
+        1998,
+        1999,
+        2000,
+        2001,
+        2002,
+        2003,
+        2004,
+        2005,
+        2006,
+        2007,
+        2008,
+        2009,
+        2010,
+        2011,
+        2012,
+        2013,
+        2014,
+        2015,
+        2016,
+        2017,
+        2018,
+        2019,
+        2020,
+        2021,
+        2022,
+        2023,
+        2024,
+        2025,
+        2026,
+        2031,
+        2032,
+        2033,
+        2034,
+        2035,
+        2036,
+        2041,
+        2042,
+        2043,
+        2044,
+        2045,
+        2046,
+        2047,
+        2048,
+        2049,
+        2050,
+        2051,
+        2052,
+        2053,
+        2054,
+        2055,
+        2056,
+        2057,
+        2058,
+        2059,
+        2060,
+        2061,
+        2062,
+        2063,
+        2064,
+        2065,
+        2066,
+        2069,
+        2070,
+        2071,
+        2072,
+        2073,
+        2074,
+        2075,
+        2076,
+        2077,
+        2078,
+        2079,
+        2080,
+        2081,
+        2082,
+        2083,
+        2084,
+        2085,
+        2086,
+        2087,
+        2088,
+        2089,
+        2090,
+        2091,
+        2092,
+        2093,
+        2094,
+        2095,
+        2096,
+        2097,
+        2098,
+        2099,
+        2100,
+        2101,
+        2107,
+        2111,
+        2113,
+        2114,
+        2115,
+        2116,
+        2117,
+        2118,
+        2119,
+        2120,
+        2121,
+        2122,
+        2127,
+        2130,
+        2131,
+        2132,
+        2133,
+        2134,
+        2135,
+        2136,
+        2137,
+        2138,
+        2139,
+        2140,
+        2141,
+        2142,
+        2143,
+        2144,
+        2149,
+        2150,
+        2151,
+        2152,
+        2155,
+        2160,
+        2163,
+        2164,
+        2170,
+        2171,
+        2172,
+        2173,
+        2174,
+        2175,
+        2176,
+        2177,
+        2178,
+        2179,
+        2180,
+        2182,
+        2183,
+        2184,
+        2185,
+        2188,
+        2189,
+        2191,
+        2192,
+        2193,
+        2194,
+        2195,
+        2196,
+        2197,
+        2198,
+        2199,
+        2200,
+        2201,
+        2202,
+        2203,
+        2207,
+        2209,
+        2210,
+        2211,
+        2212,
+        2213,
+        2214,
+        2221,
+        2222,
+        2223,
+        2224,
+        2225,
+        2226,
+        2227,
+        2228,
+        2229,
+        2231,
+        2234,
+        2236,
+        2237,
+        2238,
+        2239,
+        2240,
+        2246,
+        2247,
+        2248,
+        2249,
+        2250,
+        2251,
+        2252,
+        2253,
+        2254,
+        2255,
+        2256,
+        2257,
+        2258,
+        2259,
+        2260,
+        2262,
+        2263,
+        2264,
+        2265,
+        2266,
+        2267,
+        2268,
+        2269,
+        2270,
+        2271,
+        2274,
+        2275,
+        2276,
+        2277,
+        2278,
+        2279,
+        2284,
+        2285,
+        2287,
+        2288,
+        2289,
+        2290,
+        2293,
+        2595,
+        2598,
+        2605,
+        2608,
+        2697,
+        2698,
+        2699,
+        2700,
+        2701,
+        2702,
+        2703,
+        2704,
+        2705,
+        2706,
+        2707,
+        2708,
+        2709,
+        2710,
+        2711,
+        2712,
+        2713,
+        2714,
+        2715,
+        2716,
+        2717,
+        2718,
+        2719,
+        2720,
+        2721,
+        2722,
+        2723,
+        2724,
+        2725,
+        2726,
+        2727,
+        2728,
+        2729,
+        2730,
+        2731,
+        2732,
+        2733,
+        2734,
+        2735,
+        2736,
+        2737,
+        2738,
+        2739,
+        2740,
+        2741,
+        2742,
+        2743,
+        2744,
+        2745,
+        2746,
+        2747,
+        2748,
+        2749,
+        2750,
+        2751,
+        2752,
+        2753,
+        2754,
+        2755,
+        2756,
+        2757,
+        2758,
+        2759,
+        2760,
+        2761,
+        2762,
+        2763,
+        2764,
+        2765,
+        2766,
+        2767,
+        2768,
+        2769,
+        2770,
+        2771,
+        2772,
+        2773,
+        2774,
+        2775,
+        2776,
+        2777,
+        2778
+    ],
+    "hips": [
+        631,
+        632,
+        654,
+        657,
+        662,
+        665,
+        676,
+        677,
+        678,
+        679,
+        705,
+        720,
+        796,
+        799,
+        800,
+        801,
+        802,
+        807,
+        808,
+        809,
+        810,
+        815,
+        816,
+        822,
+        823,
+        830,
+        831,
+        832,
+        833,
+        834,
+        835,
+        836,
+        837,
+        838,
+        839,
+        840,
+        841,
+        842,
+        843,
+        844,
+        845,
+        846,
+        855,
+        856,
+        857,
+        858,
+        859,
+        860,
+        861,
+        862,
+        863,
+        864,
+        865,
+        866,
+        867,
+        868,
+        869,
+        871,
+        878,
+        881,
+        882,
+        883,
+        884,
+        885,
+        886,
+        887,
+        888,
+        889,
+        890,
+        912,
+        915,
+        916,
+        917,
+        918,
+        919,
+        920,
+        932,
+        937,
+        938,
+        939,
+        1163,
+        1166,
+        1203,
+        1204,
+        1205,
+        1206,
+        1207,
+        1208,
+        1209,
+        1210,
+        1246,
+        1247,
+        1262,
+        1263,
+        1276,
+        1277,
+        1278,
+        1321,
+        1336,
+        1337,
+        1338,
+        1339,
+        1353,
+        1354,
+        1361,
+        1362,
+        1363,
+        1364,
+        1446,
+        1447,
+        1448,
+        1449,
+        1450,
+        1454,
+        1476,
+        1497,
+        1511,
+        1513,
+        1514,
+        1515,
+        1533,
+        1534,
+        1539,
+        1540,
+        1768,
+        1769,
+        1779,
+        1780,
+        1781,
+        1782,
+        1783,
+        1784,
+        1785,
+        1786,
+        1787,
+        1788,
+        1789,
+        1790,
+        1791,
+        1792,
+        1793,
+        1794,
+        1795,
+        1796,
+        1797,
+        1798,
+        1799,
+        1800,
+        1801,
+        1802,
+        1803,
+        1804,
+        1805,
+        1806,
+        1807,
+        2909,
+        2910,
+        2911,
+        2912,
+        2913,
+        2914,
+        2915,
+        2916,
+        2917,
+        2918,
+        2919,
+        2920,
+        2921,
+        2922,
+        2923,
+        2924,
+        2925,
+        2926,
+        2927,
+        2928,
+        2929,
+        2930,
+        3018,
+        3019,
+        3021,
+        3022,
+        3080,
+        3081,
+        3082,
+        3083,
+        3084,
+        3085,
+        3086,
+        3087,
+        3088,
+        3089,
+        3090,
+        3091,
+        3092,
+        3093,
+        3094,
+        3095,
+        3096,
+        3097,
+        3098,
+        3099,
+        3100,
+        3101,
+        3102,
+        3103,
+        3104,
+        3105,
+        3106,
+        3107,
+        3108,
+        3109,
+        3110,
+        3111,
+        3112,
+        3113,
+        3114,
+        3115,
+        3116,
+        3117,
+        3118,
+        3119,
+        3120,
+        3121,
+        3122,
+        3123,
+        3124,
+        3128,
+        3129,
+        3130,
+        3136,
+        3137,
+        3138,
+        3139,
+        3140,
+        3141,
+        3142,
+        3143,
+        3144,
+        3145,
+        3146,
+        3147,
+        3148,
+        3149,
+        3150,
+        3151,
+        3152,
+        3153,
+        3154,
+        3155,
+        3156,
+        3157,
+        3158,
+        3159,
+        3160,
+        3170,
+        3172,
+        3481,
+        3484,
+        3500,
+        3502,
+        3503,
+        3507,
+        3510,
+        4120,
+        4121,
+        4142,
+        4143,
+        4150,
+        4151,
+        4164,
+        4165,
+        4166,
+        4167,
+        4193,
+        4208,
+        4284,
+        4285,
+        4288,
+        4289,
+        4290,
+        4295,
+        4296,
+        4297,
+        4298,
+        4303,
+        4304,
+        4310,
+        4311,
+        4316,
+        4317,
+        4318,
+        4319,
+        4320,
+        4321,
+        4322,
+        4323,
+        4324,
+        4325,
+        4326,
+        4327,
+        4328,
+        4329,
+        4330,
+        4331,
+        4332,
+        4341,
+        4342,
+        4343,
+        4344,
+        4345,
+        4346,
+        4347,
+        4348,
+        4349,
+        4350,
+        4351,
+        4352,
+        4353,
+        4354,
+        4355,
+        4356,
+        4364,
+        4365,
+        4368,
+        4369,
+        4370,
+        4371,
+        4372,
+        4373,
+        4374,
+        4375,
+        4376,
+        4398,
+        4399,
+        4402,
+        4403,
+        4404,
+        4405,
+        4406,
+        4418,
+        4423,
+        4424,
+        4425,
+        4649,
+        4650,
+        4689,
+        4690,
+        4691,
+        4692,
+        4693,
+        4729,
+        4730,
+        4745,
+        4746,
+        4759,
+        4760,
+        4801,
+        4812,
+        4813,
+        4814,
+        4815,
+        4829,
+        4836,
+        4837,
+        4919,
+        4920,
+        4921,
+        4922,
+        4923,
+        4927,
+        4969,
+        4983,
+        4984,
+        4986,
+        5004,
+        5005,
+        5244,
+        5245,
+        5246,
+        5247,
+        5248,
+        5249,
+        5250,
+        5251,
+        5252,
+        5253,
+        5254,
+        5255,
+        5256,
+        5257,
+        5258,
+        5259,
+        5260,
+        5261,
+        5262,
+        5263,
+        5264,
+        5265,
+        5266,
+        5267,
+        5268,
+        6368,
+        6369,
+        6370,
+        6371,
+        6372,
+        6373,
+        6374,
+        6375,
+        6376,
+        6377,
+        6378,
+        6379,
+        6380,
+        6381,
+        6382,
+        6383,
+        6384,
+        6385,
+        6386,
+        6387,
+        6388,
+        6389,
+        6473,
+        6474,
+        6504,
+        6505,
+        6506,
+        6507,
+        6508,
+        6509,
+        6510,
+        6511,
+        6512,
+        6513,
+        6514,
+        6515,
+        6516,
+        6517,
+        6518,
+        6519,
+        6520,
+        6521,
+        6522,
+        6523,
+        6524,
+        6525,
+        6526,
+        6527,
+        6528,
+        6529,
+        6530,
+        6531,
+        6532,
+        6533,
+        6534,
+        6535,
+        6536,
+        6537,
+        6538,
+        6539,
+        6540,
+        6541,
+        6542,
+        6543,
+        6544,
+        6545,
+        6549,
+        6550,
+        6551,
+        6557,
+        6558,
+        6559,
+        6560,
+        6561,
+        6562,
+        6563,
+        6564,
+        6565,
+        6566,
+        6567,
+        6568,
+        6569,
+        6570,
+        6571,
+        6572,
+        6573
+    ]
+}
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/smplx_lite.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/smplx_lite.py
new file mode 100644
index 0000000000000000000000000000000000000000..93c7172fb68fba8882a1e64bd404cfb21e3081ad
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/smplx_lite.py
@@ -0,0 +1,329 @@
+"""Lightweight SMPLX body models for fast FK and skinning on joint/vertex subsets."""
+
+from pathlib import Path
+from time import time
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import einsum, rearrange
+from smplx.utils import Struct, to_np, to_tensor
+
+from gear_sonic.trl.utils.rotation_conversion import (
+    axis_angle_to_matrix,
+    rotation_6d_to_matrix,
+)
+
+
+class SmplxLite(nn.Module):
+    def __init__(
+        self,
+        model_path="data/body_models/smplx",
+        gender="neutral",
+        num_betas=10,
+    ):
+        super().__init__()
+
+        # Load the model
+        model_path = Path(model_path)
+        if model_path.is_dir():
+            smplx_path = Path(model_path) / f"SMPLX_{gender.upper()}.npz"
+        else:
+            smplx_path = model_path
+        assert smplx_path.exists()
+        model_data = np.load(smplx_path, allow_pickle=True)
+
+        data_struct = Struct(**model_data)
+        self.faces = data_struct.f  # (F, 3)
+
+        self.register_smpl_buffers(data_struct, num_betas)
+        # self.register_smplh_buffers(data_struct, num_pca_comps, flat_hand_mean)
+        # self.register_smplx_buffers(data_struct)
+        self.register_fast_skeleton_computing_buffers()
+
+        # default_pose (99,) for torch.cat([global_orient, body_pose, default_pose])
+        other_default_pose = torch.cat(
+            [
+                torch.zeros(9),
+                to_tensor(data_struct.hands_meanl).float(),
+                to_tensor(data_struct.hands_meanr).float(),
+            ]
+        )
+        self.register_buffer("other_default_pose", other_default_pose, False)
+
+    def register_smpl_buffers(self, data_struct, num_betas):
+        # shapedirs, (V, 3, N_betas), V=10475 for SMPLX
+        shapedirs = to_tensor(to_np(data_struct.shapedirs[:, :, :num_betas])).float()
+        self.register_buffer("shapedirs", shapedirs, False)
+
+        # v_template, (V, 3)
+        v_template = to_tensor(to_np(data_struct.v_template)).float()
+        self.register_buffer("v_template", v_template, False)
+
+        # J_regressor, (J, V), J=55 for SMPLX
+        J_regressor = to_tensor(to_np(data_struct.J_regressor)).float()
+        self.register_buffer("J_regressor", J_regressor, False)
+
+        # posedirs, (54*9, V, 3), note that the first global_orient is not included
+        posedirs = to_tensor(to_np(data_struct.posedirs)).float()  # (V, 3, 54*9)
+        posedirs = rearrange(posedirs, "v c n -> n v c")
+        self.register_buffer("posedirs", posedirs, False)
+
+        # lbs_weights, (V, J), J=55
+        lbs_weights = to_tensor(to_np(data_struct.weights)).float()
+        self.register_buffer("lbs_weights", lbs_weights, False)
+
+        # parents, (J), long
+        parents = to_tensor(to_np(data_struct.kintree_table[0])).long()
+        parents[0] = -1
+        self.register_buffer("parents", parents, False)
+
+    def register_smplh_buffers(self, data_struct, num_pca_comps, flat_hand_mean):
+        # hand_pca, (N_pca, 45)
+        left_hand_components = to_tensor(data_struct.hands_componentsl[:num_pca_comps]).float()
+        right_hand_components = to_tensor(data_struct.hands_componentsr[:num_pca_comps]).float()
+        self.register_buffer("left_hand_components", left_hand_components, False)
+        self.register_buffer("right_hand_components", right_hand_components, False)
+
+        # hand_mean, (45,)
+        left_hand_mean = to_tensor(data_struct.hands_meanl).float()
+        right_hand_mean = to_tensor(data_struct.hands_meanr).float()
+        if not flat_hand_mean:
+            left_hand_mean = torch.zeros_like(left_hand_mean)
+            right_hand_mean = torch.zeros_like(right_hand_mean)
+        self.register_buffer("left_hand_mean", left_hand_mean, False)
+        self.register_buffer("right_hand_mean", right_hand_mean, False)
+
+    def register_smplx_buffers(self, data_struct):
+        # expr_dirs, (V, 3, N_expr)
+        expr_dirs = to_tensor(to_np(data_struct.shapedirs[:, :, 300:310])).float()
+        self.register_buffer("expr_dirs", expr_dirs, False)
+
+    def register_fast_skeleton_computing_buffers(self):
+        # For fast computing of skeleton under beta
+        J_template = self.J_regressor @ self.v_template  # (J, 3)
+        J_shapedirs = torch.einsum("jv, vcd -> jcd", self.J_regressor, self.shapedirs)  # (J, 3, 10)
+        self.register_buffer("J_template", J_template, False)
+        self.register_buffer("J_shapedirs", J_shapedirs, False)
+
+    def get_skeleton(self, betas):
+        return self.J_template + einsum(betas, self.J_shapedirs, "... k, j c k -> ... j c")
+
+    # Parents needs to be a list in order to torch.compile the forward() method
+    # properly, so we have this method to create the list.
+    #
+    # Called in hmr4d.utils.smplx_utils.make_smplx().
+    def create_parents_list(self):
+        self.parents_list = self.parents.tolist()
+
+    def forward(
+        self,
+        body_pose,
+        betas,
+        global_orient,
+        transl=None,
+        rotation_type="aa",
+    ):
+        """
+        Args:
+            body_pose: (B, L, 63)
+            betas: (B, L, 10)
+            global_orient: (B, L, 3)
+            transl: (B, L, 3)
+        Returns:
+            vertices: (B, L, V, 3)
+        """
+        # 1. Convert [global_orient, body_pose, other_default_pose] to rot_mats
+        other_default_pose = self.other_default_pose  # (99,)
+        other_default_pose[:] = 0.0
+        if rotation_type == "aa":
+            other_default_pose = other_default_pose.expand(*body_pose.shape[:-1], -1)
+            full_pose = torch.cat([global_orient, body_pose, other_default_pose], dim=-1)
+            rot_mats = axis_angle_to_matrix(full_pose.reshape(*full_pose.shape[:-1], 55, 3))
+            del full_pose, other_default_pose
+        else:
+            assert rotation_type == "r6d"  # useful when doing smplify
+            other_default_pose = axis_angle_to_matrix(other_default_pose.view(33, 3))
+            part_full_pose = torch.cat([global_orient, body_pose], dim=-1)
+            rot_mats = rotation_6d_to_matrix(part_full_pose.view(*part_full_pose.shape[:-1], 22, 6))
+            other_default_pose = other_default_pose.expand(*rot_mats.shape[:-3], -1, -1, -1)
+            rot_mats = torch.cat([rot_mats, other_default_pose], dim=-3)
+            del part_full_pose, other_default_pose
+
+        # 2. Forward Kinematics
+        J = self.get_skeleton(betas)  # (*, 55, 3)
+        if not hasattr(self, "parents_list"):
+            self.create_parents_list()
+        human_joints_info = {
+            "J": J[0].cpu(),
+            "parents_list": self.parents_list,
+            "rot_mats": rot_mats[0].cpu(),
+        }
+        torch.save(human_joints_info, "data/human_joints_info.pkl")
+        posed_joints, A = batch_rigid_transform_v2(rot_mats, J, self.parents_list)
+
+        # 3. Canonical v_posed = v_template + shaped_offsets + pose_offsets
+        pose_feature = rot_mats[..., 1:, :, :] - rot_mats.new([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
+        pose_feature = pose_feature.view(*pose_feature.shape[:-3], -1)  # (*, 55*3*3)
+        v_posed = (
+            self.v_template
+            + einsum(betas, self.shapedirs, "... k, v c k -> ... v c")
+            + einsum(pose_feature, self.posedirs, "... k, k v c -> ... v c")
+        )
+        del pose_feature, rot_mats
+
+        # 4. Skinning
+        T = einsum(self.lbs_weights, A, "v j, ... j c d -> ... v c d")
+        verts = einsum(T[..., :3, :3], v_posed, "... v c d, ... v d -> ... v c") + T[..., :3, 3]
+
+        # 5. Translation
+        if transl is not None:
+            verts = verts + transl[..., None, :]
+        return verts, posed_joints
+
+
+class SmplxLiteCoco17(SmplxLite):
+    """Output COCO17 joints (Faster, but cannot output vertices)"""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        # Compute mapping
+        smplx2smpl = torch.load("groot/rl/trl/utils/smplx/body_model/smplx2smpl_sparse.pt")
+        COCO17_regressor = torch.load(
+            "groot/rl/trl/utils/smplx/body_model/smpl_coco17_J_regressor.pt"
+        )
+        smplx2coco17 = torch.matmul(COCO17_regressor, smplx2smpl.to_dense())
+
+        jids, smplx_vids = torch.where(smplx2coco17 != 0)
+        smplx2coco17_interestd = torch.zeros([len(smplx_vids), 17])
+        for idx, (jid, smplx_vid) in enumerate(zip(jids, smplx_vids)):
+            smplx2coco17_interestd[idx, jid] = smplx2coco17[jid, smplx_vid]
+        self.register_buffer("smplx2coco17_interestd", smplx2coco17_interestd, False)  # (132, 17)
+
+        # Update to vertices of interest
+        self.v_template = self.v_template[smplx_vids].clone()  # (V', 3)
+        self.shapedirs = self.shapedirs[smplx_vids].clone()  # (V', 3, K)
+        self.posedirs = self.posedirs[:, smplx_vids].clone()  # (K, V', 3)
+        self.lbs_weights = self.lbs_weights[smplx_vids].clone()  # (V', J)
+
+    def forward(self, body_pose, betas, global_orient, transl):
+        """Returns: joints (*, 17, 3). (B, L) or  (B,) are both supported."""
+        # Use super class's forward to get verts
+        verts = super().forward(body_pose, betas, global_orient, transl)  # (*, 132, 3)
+        joints = einsum(self.smplx2coco17_interestd, verts, "v j, ... v c -> ... j c")
+        return joints
+
+
+class SmplxLiteV437Coco17(SmplxLite):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        # Compute mapping (COCO17)
+        smplx2smpl = torch.load("groot/rl/trl/utils/smplx/body_model/smplx2smpl_sparse.pt")
+        COCO17_regressor = torch.load(
+            "groot/rl/trl/utils/smplx/body_model/smpl_coco17_J_regressor.pt"
+        )
+        smplx2coco17 = torch.matmul(COCO17_regressor, smplx2smpl.to_dense())
+
+        jids, smplx_vids = torch.where(smplx2coco17 != 0)
+        smplx2coco17_interestd = torch.zeros([len(smplx_vids), 17])
+        for idx, (jid, smplx_vid) in enumerate(zip(jids, smplx_vids)):
+            smplx2coco17_interestd[idx, jid] = smplx2coco17[jid, smplx_vid]
+        self.register_buffer("smplx2coco17_interestd", smplx2coco17_interestd, False)  # (132, 17)
+        assert len(smplx_vids) == 132
+
+        # Verts437
+        smplx_vids2 = torch.load("groot/rl/trl/utils/smplx/body_model/smplx_verts437.pt")
+        smplx_vids = torch.cat([smplx_vids, smplx_vids2])
+
+        # Update to vertices of interest
+        self.v_template = self.v_template[smplx_vids].clone()  # (V', 3)
+        self.shapedirs = self.shapedirs[smplx_vids].clone()  # (V', 3, K)
+        self.posedirs = self.posedirs[:, smplx_vids].clone()  # (K, V', 3)
+        self.lbs_weights = self.lbs_weights[smplx_vids].clone()  # (V', J)
+
+    def forward(self, body_pose, betas, global_orient, transl):
+        """
+        Returns:
+            verts_437: (*, 437, 3)
+            joints (*, 17, 3). (B, L) or  (B,) are both supported.
+        """
+        # Use super class's forward to get verts
+        verts = super().forward(body_pose, betas, global_orient, transl)  # (*, 132+437, 3)
+
+        verts_437 = verts[..., 132:, :].clone()
+        joints = einsum(self.smplx2coco17_interestd, verts[..., :132, :], "v j, ... v c -> ... j c")
+        return verts_437, joints
+
+
+class SmplxLiteSmplN24(SmplxLite):
+    """Output SMPL(not smplx)-Neutral 24 joints (Faster, but cannot output vertices)"""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        # Compute mapping
+        smplx2smpl = torch.load("groot/rl/trl/utils/smplx/body_model/smplx2smpl_sparse.pt")
+        smpl2joints = torch.load("groot/rl/trl/utils/smplx/body_model/smpl_neutral_J_regressor.pt")
+        smplx2joints = torch.matmul(smpl2joints, smplx2smpl.to_dense())
+
+        jids, smplx_vids = torch.where(smplx2joints != 0)
+        smplx2joints_interested = torch.zeros([len(smplx_vids), smplx2joints.size(0)])
+        for idx, (jid, smplx_vid) in enumerate(zip(jids, smplx_vids)):
+            smplx2joints_interested[idx, jid] = smplx2joints[jid, smplx_vid]
+        self.register_buffer("smplx2joints_interested", smplx2joints_interested, False)  # (V', J)
+
+        # Update to vertices of interest
+        self.v_template = self.v_template[smplx_vids].clone()  # (V', 3)
+        self.shapedirs = self.shapedirs[smplx_vids].clone()  # (V', 3, K)
+        self.posedirs = self.posedirs[:, smplx_vids].clone()  # (K, V', 3)
+        self.lbs_weights = self.lbs_weights[smplx_vids].clone()  # (V', J)
+
+    def forward(self, body_pose, betas, global_orient, transl):
+        """Returns: joints (*, J, 3). (B, L) or  (B,) are both supported."""
+        # Use super class's forward to get verts
+        verts, posed_joints = super().forward(body_pose, betas, global_orient, transl)  # (*, V', 3)
+        joints = einsum(self.smplx2joints_interested, verts, "v j, ... v c -> ... j c")
+        return joints, posed_joints
+
+
+def batch_rigid_transform_v2(rot_mats, joints, parents):
+    """
+    Args:
+        rot_mats: (*, J, 3, 3)
+        joints: (*, J, 3)
+    """
+    # check shape, since sometimes beta has shape=1
+    rot_mats_shape_prefix = rot_mats.shape[:-3]
+    if rot_mats_shape_prefix != joints.shape[:-2]:
+        joints = joints.expand(*rot_mats_shape_prefix, -1, -1)
+
+    rel_joints = joints.clone()
+    rel_joints[..., 1:, :] -= joints[..., parents[1:], :]
+    transforms_mat = torch.cat([rot_mats, rel_joints[..., :, None]], dim=-1)  # (*, J, 3, 4)
+    transforms_mat = F.pad(transforms_mat, [0, 0, 0, 1], value=0.0)
+    transforms_mat[..., 3, 3] = 1.0  # (*, J, 4, 4)
+
+    transform_chain = [transforms_mat[..., 0, :, :]]
+    for i in range(1, len(parents)):
+        # Subtract the joint location at the rest pose
+        # No need for rotation, since it's identity when at rest
+        curr_res = torch.matmul(transform_chain[parents[i]], transforms_mat[..., i, :, :])
+        transform_chain.append(curr_res)
+
+    transforms = torch.stack(transform_chain, dim=-3)  # (*, J, 4, 4)
+
+    # The last column of the transformations contains the posed joints
+    posed_joints = transforms[..., :3, 3].clone()
+    rel_transforms = transforms.clone()
+    rel_transforms[..., :3, 3] -= einsum(
+        transforms[..., :3, :3], joints, "... j c d, ... j d -> ... j c"
+    )
+    return posed_joints, rel_transforms
+
+
+def sync_time():
+    torch.cuda.synchronize()
+    return time()
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/utils.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..cdbd71033af3d6e538d6577b2bac5514e8f4b1e0
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/body_model/utils.py
@@ -0,0 +1,564 @@
+"""SMPLH/SMPLX joint name constants, body-part indices, and OpenPose mapping."""
+
+import os
+
+import numpy as np
+import torch
+
+SMPLH_JOINT_NAMES = [
+    "pelvis",
+    "left_hip",
+    "right_hip",
+    "spine1",
+    "left_knee",
+    "right_knee",
+    "spine2",
+    "left_ankle",
+    "right_ankle",
+    "spine3",
+    "left_foot",
+    "right_foot",
+    "neck",
+    "left_collar",
+    "right_collar",
+    "head",
+    "left_shoulder",
+    "right_shoulder",
+    "left_elbow",
+    "right_elbow",
+    "left_wrist",
+    "right_wrist",
+    "left_index1",
+    "left_index2",
+    "left_index3",
+    "left_middle1",
+    "left_middle2",
+    "left_middle3",
+    "left_pinky1",
+    "left_pinky2",
+    "left_pinky3",
+    "left_ring1",
+    "left_ring2",
+    "left_ring3",
+    "left_thumb1",
+    "left_thumb2",
+    "left_thumb3",
+    "right_index1",
+    "right_index2",
+    "right_index3",
+    "right_middle1",
+    "right_middle2",
+    "right_middle3",
+    "right_pinky1",
+    "right_pinky2",
+    "right_pinky3",
+    "right_ring1",
+    "right_ring2",
+    "right_ring3",
+    "right_thumb1",
+    "right_thumb2",
+    "right_thumb3",
+    "nose",
+    "right_eye",
+    "left_eye",
+    "right_ear",
+    "left_ear",
+    "left_big_toe",
+    "left_small_toe",
+    "left_heel",
+    "right_big_toe",
+    "right_small_toe",
+    "right_heel",
+    "left_thumb",
+    "left_index",
+    "left_middle",
+    "left_ring",
+    "left_pinky",
+    "right_thumb",
+    "right_index",
+    "right_middle",
+    "right_ring",
+    "right_pinky",
+]
+
+SMPLH_LEFT_LEG = ["left_hip", "left_knee", "left_ankle", "left_foot"]
+SMPLH_RIGHT_LEG = ["right_hip", "right_knee", "right_ankle", "right_foot"]
+SMPLH_LEFT_ARM = ["left_collar", "left_shoulder", "left_elbow", "left_wrist"]
+SMPLH_RIGHT_ARM = ["right_collar", "right_shoulder", "right_elbow", "right_wrist"]
+SMPLH_HEAD = ["neck", "head"]
+SMPLH_SPINE = ["spine1", "spine2", "spine3"]
+
+# name to 21 index (without pelvis, hand, and extra)
+_name_2_idx = {j: i for i, j in enumerate(SMPLH_JOINT_NAMES[1:22])}
+SMPLH_PART_IDX = {
+    "left_leg": [_name_2_idx[x] for x in SMPLH_LEFT_LEG],
+    "right_leg": [_name_2_idx[x] for x in SMPLH_RIGHT_LEG],
+    "left_arm": [_name_2_idx[x] for x in SMPLH_LEFT_ARM],
+    "right_arm": [_name_2_idx[x] for x in SMPLH_RIGHT_ARM],
+    "two_legs": [_name_2_idx[x] for x in SMPLH_LEFT_LEG + SMPLH_RIGHT_LEG],
+    "left_arm_and_leg": [_name_2_idx[x] for x in SMPLH_LEFT_ARM + SMPLH_LEFT_LEG],
+    "right_arm_and_leg": [_name_2_idx[x] for x in SMPLH_RIGHT_ARM + SMPLH_RIGHT_LEG],
+}
+
+# name to full index
+_name_2_idx_full = {j: i for i, j in enumerate(SMPLH_JOINT_NAMES)}
+SMPLH_PART_IDX_FULL = {
+    "lower_body": [_name_2_idx_full[x] for x in ["pelvis"] + SMPLH_LEFT_LEG + SMPLH_RIGHT_LEG]
+}
+
+# ===== ⬇️ Fitting optimizer ⬇️ ===== #
+SMPL_JOINTS = {
+    "hips": 0,
+    "leftUpLeg": 1,
+    "rightUpLeg": 2,
+    "spine": 3,
+    "leftLeg": 4,
+    "rightLeg": 5,
+    "spine1": 6,
+    "leftFoot": 7,
+    "rightFoot": 8,
+    "spine2": 9,
+    "leftToeBase": 10,
+    "rightToeBase": 11,
+    "neck": 12,
+    "leftShoulder": 13,
+    "rightShoulder": 14,
+    "head": 15,
+    "leftArm": 16,
+    "rightArm": 17,
+    "leftForeArm": 18,
+    "rightForeArm": 19,
+    "leftHand": 20,
+    "rightHand": 21,
+}
+
+# chosen virtual mocap markers that are "keypoints" to work with
+KEYPT_VERTS = [
+    4404,
+    920,
+    3076,
+    3169,
+    823,
+    4310,
+    1010,
+    1085,
+    4495,
+    4569,
+    6615,
+    3217,
+    3313,
+    6713,
+    6785,
+    3383,
+    6607,
+    3207,
+    1241,
+    1508,
+    4797,
+    4122,
+    1618,
+    1569,
+    5135,
+    5040,
+    5691,
+    5636,
+    5404,
+    2230,
+    2173,
+    2108,
+    134,
+    3645,
+    6543,
+    3123,
+    3024,
+    4194,
+    1306,
+    182,
+    3694,
+    4294,
+    744,
+]
+
+
+# From https://github.com/vchoutas/smplify-x/blob/master/smplifyx/utils.py
+# Please see license for usage restrictions.
+def smpl_to_openpose(
+    model_type="smplx",
+    use_hands=True,
+    use_face=True,
+    use_face_contour=False,
+    openpose_format="coco25",
+):
+    """Returns the indices of the permutation that maps SMPL to OpenPose
+
+    Parameters
+    ----------
+    model_type: str, optional
+        The type of SMPL-like model that is used. The default mapping
+        returned is for the SMPLX model
+    use_hands: bool, optional
+        Flag for adding to the returned permutation the mapping for the
+        hand keypoints. Defaults to True
+    use_face: bool, optional
+        Flag for adding to the returned permutation the mapping for the
+        face keypoints. Defaults to True
+    use_face_contour: bool, optional
+        Flag for appending the facial contour keypoints. Defaults to False
+    openpose_format: bool, optional
+        The output format of OpenPose. For now only COCO-25 and COCO-19 is
+        supported. Defaults to 'coco25'
+
+    """
+    if openpose_format.lower() == "coco25":
+        if model_type == "smpl":
+            return np.array(
+                [
+                    24,
+                    12,
+                    17,
+                    19,
+                    21,
+                    16,
+                    18,
+                    20,
+                    0,
+                    2,
+                    5,
+                    8,
+                    1,
+                    4,
+                    7,
+                    25,
+                    26,
+                    27,
+                    28,
+                    29,
+                    30,
+                    31,
+                    32,
+                    33,
+                    34,
+                ],
+                dtype=np.int32,
+            )
+        elif model_type == "smplh":
+            body_mapping = np.array(
+                [
+                    52,
+                    12,
+                    17,
+                    19,
+                    21,
+                    16,
+                    18,
+                    20,
+                    0,
+                    2,
+                    5,
+                    8,
+                    1,
+                    4,
+                    7,
+                    53,
+                    54,
+                    55,
+                    56,
+                    57,
+                    58,
+                    59,
+                    60,
+                    61,
+                    62,
+                ],
+                dtype=np.int32,
+            )
+            mapping = [body_mapping]
+            if use_hands:
+                lhand_mapping = np.array(
+                    [
+                        20,
+                        34,
+                        35,
+                        36,
+                        63,
+                        22,
+                        23,
+                        24,
+                        64,
+                        25,
+                        26,
+                        27,
+                        65,
+                        31,
+                        32,
+                        33,
+                        66,
+                        28,
+                        29,
+                        30,
+                        67,
+                    ],
+                    dtype=np.int32,
+                )
+                rhand_mapping = np.array(
+                    [
+                        21,
+                        49,
+                        50,
+                        51,
+                        68,
+                        37,
+                        38,
+                        39,
+                        69,
+                        40,
+                        41,
+                        42,
+                        70,
+                        46,
+                        47,
+                        48,
+                        71,
+                        43,
+                        44,
+                        45,
+                        72,
+                    ],
+                    dtype=np.int32,
+                )
+                mapping += [lhand_mapping, rhand_mapping]
+            return np.concatenate(mapping)
+        # SMPLX
+        elif model_type == "smplx":
+            body_mapping = np.array(
+                [
+                    55,
+                    12,
+                    17,
+                    19,
+                    21,
+                    16,
+                    18,
+                    20,
+                    0,
+                    2,
+                    5,
+                    8,
+                    1,
+                    4,
+                    7,
+                    56,
+                    57,
+                    58,
+                    59,
+                    60,
+                    61,
+                    62,
+                    63,
+                    64,
+                    65,
+                ],
+                dtype=np.int32,
+            )
+            mapping = [body_mapping]
+            if use_hands:
+                lhand_mapping = np.array(
+                    [
+                        20,
+                        37,
+                        38,
+                        39,
+                        66,
+                        25,
+                        26,
+                        27,
+                        67,
+                        28,
+                        29,
+                        30,
+                        68,
+                        34,
+                        35,
+                        36,
+                        69,
+                        31,
+                        32,
+                        33,
+                        70,
+                    ],
+                    dtype=np.int32,
+                )
+                rhand_mapping = np.array(
+                    [
+                        21,
+                        52,
+                        53,
+                        54,
+                        71,
+                        40,
+                        41,
+                        42,
+                        72,
+                        43,
+                        44,
+                        45,
+                        73,
+                        49,
+                        50,
+                        51,
+                        74,
+                        46,
+                        47,
+                        48,
+                        75,
+                    ],
+                    dtype=np.int32,
+                )
+
+                mapping += [lhand_mapping, rhand_mapping]
+            if use_face:
+                #  end_idx = 127 + 17 * use_face_contour
+                face_mapping = np.arange(76, 127 + 17 * use_face_contour, dtype=np.int32)
+                mapping += [face_mapping]
+
+            return np.concatenate(mapping)
+        else:
+            raise ValueError("Unknown model type: {}".format(model_type))
+    elif openpose_format == "coco19":
+        if model_type == "smpl":
+            return np.array(
+                [24, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5, 8, 1, 4, 7, 25, 26, 27, 28],
+                dtype=np.int32,
+            )
+        elif model_type == "smplh":
+            body_mapping = np.array(
+                [52, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5, 8, 1, 4, 7, 53, 54, 55, 56],
+                dtype=np.int32,
+            )
+            mapping = [body_mapping]
+            if use_hands:
+                lhand_mapping = np.array(
+                    [
+                        20,
+                        34,
+                        35,
+                        36,
+                        57,
+                        22,
+                        23,
+                        24,
+                        58,
+                        25,
+                        26,
+                        27,
+                        59,
+                        31,
+                        32,
+                        33,
+                        60,
+                        28,
+                        29,
+                        30,
+                        61,
+                    ],
+                    dtype=np.int32,
+                )
+                rhand_mapping = np.array(
+                    [
+                        21,
+                        49,
+                        50,
+                        51,
+                        62,
+                        37,
+                        38,
+                        39,
+                        63,
+                        40,
+                        41,
+                        42,
+                        64,
+                        46,
+                        47,
+                        48,
+                        65,
+                        43,
+                        44,
+                        45,
+                        66,
+                    ],
+                    dtype=np.int32,
+                )
+                mapping += [lhand_mapping, rhand_mapping]
+            return np.concatenate(mapping)
+        # SMPLX
+        elif model_type == "smplx":
+            body_mapping = np.array(
+                [55, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5, 8, 1, 4, 7, 56, 57, 58, 59],
+                dtype=np.int32,
+            )
+            mapping = [body_mapping]
+            if use_hands:
+                lhand_mapping = np.array(
+                    [
+                        20,
+                        37,
+                        38,
+                        39,
+                        60,
+                        25,
+                        26,
+                        27,
+                        61,
+                        28,
+                        29,
+                        30,
+                        62,
+                        34,
+                        35,
+                        36,
+                        63,
+                        31,
+                        32,
+                        33,
+                        64,
+                    ],
+                    dtype=np.int32,
+                )
+                rhand_mapping = np.array(
+                    [
+                        21,
+                        52,
+                        53,
+                        54,
+                        65,
+                        40,
+                        41,
+                        42,
+                        66,
+                        43,
+                        44,
+                        45,
+                        67,
+                        49,
+                        50,
+                        51,
+                        68,
+                        46,
+                        47,
+                        48,
+                        69,
+                    ],
+                    dtype=np.int32,
+                )
+
+                mapping += [lhand_mapping, rhand_mapping]
+            if use_face:
+                face_mapping = np.arange(70, 70 + 51 + 17 * use_face_contour, dtype=np.int32)
+                mapping += [face_mapping]
+
+            return np.concatenate(mapping)
+        else:
+            raise ValueError("Unknown model type: {}".format(model_type))
+    else:
+        raise ValueError("Unknown joint format: {}".format(openpose_format))
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/smplx_utils.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/smplx_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..ce6189ca5622a35958ba64aa731cfa3e8c73ee86
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/smplx/smplx_utils.py
@@ -0,0 +1,515 @@
+"""SMPL/SMPLX body model utilities: creation and pose decomposition."""
+
+import os
+import pickle
+from pathlib import Path
+
+import numpy as np
+import smplx
+import torch
+import torch.nn.functional as F
+from smplx import SMPL, SMPLX, SMPLXLayer
+
+from gear_sonic.trl.utils.smplx.body_model import BodyModelSMPLH, BodyModelSMPLX
+from gear_sonic.trl.utils.smplx.body_model.smplx_lite import (
+    SmplxLiteCoco17,
+    SmplxLiteSmplN24,
+    SmplxLiteV437Coco17,
+)
+
+# fmt: off
+SMPLH_PARENTS = torch.tensor([-1,  0,  0,  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  9,  9, 12, 13, 14,
+                              16, 17, 18, 19, 20, 22, 23, 20, 25, 26, 20, 28, 29, 20, 31, 32, 20, 34,
+                              35, 21, 37, 38, 21, 40, 41, 21, 43, 44, 21, 46, 47, 21, 49, 50])
+# fmt: on
+
+
+def _ensure_body_models_downloaded(model_path: str) -> Path:
+    """Ensure body_models directory exists.
+
+    Args:
+        model_path: Path to body_models directory (e.g., "data/body_models")
+
+    Returns:
+        Path object to the body_models directory
+
+    Raises:
+        FileNotFoundError: If body models are not found at the given path.
+            Download SMPL/SMPLX body models from https://smpl-x.is.tue.mpg.de/
+            and place them under ``data/body_models/``.
+    """
+    model_path_obj = Path(model_path)
+
+    if model_path_obj.exists():
+        smplx_subdir = model_path_obj / "smplx"
+        if smplx_subdir.exists() and any(smplx_subdir.iterdir()):
+            return model_path_obj
+
+    project_root = Path(__file__).parent.parent.parent.parent.parent.parent
+    full_path = project_root / model_path
+    if full_path.exists():
+        smplx_subdir = full_path / "smplx"
+        if smplx_subdir.exists() and any(smplx_subdir.iterdir()):
+            return full_path
+
+    raise FileNotFoundError(
+        f"Body models not found at '{model_path}' or '{full_path}'.\n"
+        f"Download SMPL/SMPLX body models from https://smpl-x.is.tue.mpg.de/ "
+        f"and place them under data/body_models/ in the repository root.\n"
+        f"Expected structure:\n"
+        f"  data/body_models/smplx/SMPLX_NEUTRAL.npz\n"
+        f"  data/body_models/smplx/SMPLX_MALE.npz\n"
+        f"  data/body_models/smplx/SMPLX_FEMALE.npz"
+    )
+
+
+def make_smplx(type="neu_fullpose", **kwargs):
+    if type == "neu_fullpose":
+        model = smplx.create(
+            model_path="inputs/models/smplx/SMPLX_NEUTRAL.npz",
+            use_pca=False,
+            flat_hand_mean=True,
+            **kwargs,
+        )
+    elif type == "supermotion":
+        # SuperMotion is trained on BEDLAM dataset, the smplx config is the same except only 10 betas are used
+        bm_kwargs = {
+            "model_type": "smplx",
+            "gender": "neutral",
+            "num_pca_comps": 12,
+            "flat_hand_mean": False,
+        }
+        bm_kwargs.update(kwargs)
+        body_models_path = _ensure_body_models_downloaded("data/body_models")
+        model = BodyModelSMPLX(model_path=str(body_models_path), **bm_kwargs)
+    elif type == "supermotion_EVAL3DPW":
+        # SuperMotion is trained on BEDLAM dataset, the smplx config is the same except only 10 betas are used
+        bm_kwargs = {
+            "model_type": "smplx",
+            "gender": "neutral",
+            "num_pca_comps": 12,
+            "flat_hand_mean": True,
+        }
+        bm_kwargs.update(kwargs)
+        body_models_path = _ensure_body_models_downloaded("data/body_models")
+        model = BodyModelSMPLX(model_path=str(body_models_path), **bm_kwargs)
+    elif type == "supermotion_coco17":
+        # Fast but only predicts 17 joints
+        body_models_path = _ensure_body_models_downloaded("data/body_models")
+        model = SmplxLiteCoco17(model_path=str(body_models_path / "smplx"))
+    elif type == "supermotion_v437coco17":
+        # Predicts 437 verts and 17 joints
+        body_models_path = _ensure_body_models_downloaded("data/body_models")
+        model = SmplxLiteV437Coco17(model_path=str(body_models_path / "smplx"))
+    elif type == "supermotion_smpl24":
+        body_models_path = _ensure_body_models_downloaded("data/body_models")
+        model = SmplxLiteSmplN24(model_path=str(body_models_path / "smplx"))
+    elif type == "supermotion_smpl24_male":
+        body_models_path = _ensure_body_models_downloaded("data/body_models")
+        model = SmplxLiteSmplN24(gender="male", model_path=str(body_models_path / "smplx"))
+    elif type == "supermotion_smpl24_female":
+        body_models_path = _ensure_body_models_downloaded("data/body_models")
+        model = SmplxLiteSmplN24(gender="female", model_path=str(body_models_path / "smplx"))
+    elif type == "rich-smplx":
+        # https://github.com/paulchhuang/rich_toolkit/blob/main/smplx2images.py
+        bm_kwargs = {
+            "model_type": "smplx",
+            "gender": kwargs.get("gender", "male"),
+            "num_pca_comps": 12,
+            "flat_hand_mean": False,
+            # create_expression=True, create_jaw_pose=Ture
+        }
+        # A /smplx folder should exist under the model_path
+        body_models_path = _ensure_body_models_downloaded("data/body_models")
+        model = BodyModelSMPLX(model_path=str(body_models_path), **bm_kwargs)
+    elif type == "rich-smplh":
+        bm_kwargs = {
+            "model_type": "smplh",
+            "gender": kwargs.get("gender", "male"),
+            "use_pca": False,
+            "flat_hand_mean": True,
+        }
+        body_models_path = _ensure_body_models_downloaded("data/body_models")
+        model = BodyModelSMPLH(model_path=str(body_models_path), **bm_kwargs)
+
+    elif type in ["smplx-circle", "smplx-groundlink"]:
+        # don't use hand
+        body_models_path = _ensure_body_models_downloaded("data/body_models")
+        bm_kwargs = {
+            "model_path": str(body_models_path),
+            "model_type": "smplx",
+            "gender": kwargs.get("gender"),
+            "num_betas": 16,
+            "num_expression": 0,
+        }
+        model = BodyModelSMPLX(**bm_kwargs)
+
+    elif type == "smplx-motionx":
+        layer_args = {
+            "create_global_orient": False,
+            "create_body_pose": False,
+            "create_left_hand_pose": False,
+            "create_right_hand_pose": False,
+            "create_jaw_pose": False,
+            "create_leye_pose": False,
+            "create_reye_pose": False,
+            "create_betas": False,
+            "create_expression": False,
+            "create_transl": False,
+        }
+
+        body_models_path = _ensure_body_models_downloaded("data/body_models")
+        bm_kwargs = {
+            "model_type": "smplx",
+            "model_path": str(body_models_path),
+            "gender": "neutral",
+            "use_pca": False,
+            "use_face_contour": True,
+            **layer_args,
+        }
+        model = smplx.create(**bm_kwargs)
+
+    elif type == "smplx-samp":
+        # don't use hand
+        body_models_path = _ensure_body_models_downloaded("data/body_models")
+        bm_kwargs = {
+            "model_path": str(body_models_path),
+            "model_type": "smplx",
+            "gender": kwargs.get("gender"),
+            "num_betas": 10,
+            "num_expression": 0,
+        }
+        model = BodyModelSMPLX(**bm_kwargs)
+
+    elif type == "smplx-bedlam":
+        # don't use hand
+        bm_kwargs = {
+            "model_path": "data/body_models",
+            "model_type": "smplx",
+            "gender": kwargs.get("gender"),
+            "num_betas": 11,
+            "num_expression": 0,
+        }
+        model = BodyModelSMPLX(**bm_kwargs)
+
+    elif type in ["smplx-layer", "smplx-fit3d"]:
+        # Use layer
+        if type == "smplx-fit3d":
+            assert (
+                kwargs.get("gender") == "neutral"
+            ), "smplx-fit3d use neutral model: https://github.com/sminchisescu-research/imar_vision_datasets_tools/blob/e8c8f83ffac23cc36adf8ec8d0fd1c55679484ef/util/smplx_util.py#L15C34-L15C34"
+
+        body_models_path = _ensure_body_models_downloaded("data/body_models")
+        bm_kwargs = {
+            "model_path": str(body_models_path / "smplx"),
+            "gender": kwargs.get("gender"),
+            "num_betas": 10,
+            "num_expression": 10,
+        }
+        model = SMPLXLayer(**bm_kwargs)
+
+    elif type == "smpl":
+        body_models_path = _ensure_body_models_downloaded("data/body_models")
+        bm_kwargs = {
+            "model_path": str(body_models_path),
+            "model_type": "smpl",
+            "gender": "neutral",
+            "num_betas": 10,
+            "create_body_pose": False,
+            "create_betas": False,
+            "create_global_orient": False,
+            "create_transl": False,
+        }
+        bm_kwargs.update(kwargs)
+        # model = SMPL(**bm_kwargs)
+        model = BodyModelSMPLH(**bm_kwargs)
+    elif type == "smplh":
+        bm_kwargs = {
+            "model_type": "smplh",
+            "gender": kwargs.get("gender", "male"),
+            "use_pca": False,
+            "flat_hand_mean": False,
+        }
+        body_models_path = _ensure_body_models_downloaded("data/body_models")
+        model = BodyModelSMPLH(model_path=str(body_models_path), **bm_kwargs)
+
+    else:
+        raise NotImplementedError
+
+    if hasattr(model, "create_parents_list"):
+        model.create_parents_list()
+
+    return model
+
+
+def load_parents(npz_path="models/smplx/SMPLX_NEUTRAL.npz"):
+    smplx_struct = np.load("models/smplx/SMPLX_NEUTRAL.npz", allow_pickle=True)
+    parents = smplx_struct["kintree_table"][0].astype(np.long)
+    parents[0] = -1
+    return parents
+
+
+def load_smpl_faces(npz_path="models/smplh/SMPLH_FEMALE.pkl"):
+    with open(npz_path, "rb") as f:
+        smpl_model = pickle.load(f, encoding="latin1")
+    faces = np.array(smpl_model["f"].astype(np.int64))
+    return faces
+
+
+def decompose_fullpose(fullpose, model_type="smplx"):
+    assert model_type == "smplx"
+
+    fullpose_dict = {
+        "global_orient": fullpose[..., :3],
+        "body_pose": fullpose[..., 3:66],
+        "jaw_pose": fullpose[..., 66:69],
+        "leye_pose": fullpose[..., 69:72],
+        "reye_pose": fullpose[..., 72:75],
+        "left_hand_pose": fullpose[..., 75:120],
+        "right_hand_pose": fullpose[..., 120:165],
+    }
+
+    return fullpose_dict
+
+
+def compose_fullpose(fullpose_dict, model_type="smplx"):
+    assert model_type == "smplx"
+    fullpose = torch.cat(
+        [
+            fullpose_dict[k]
+            for k in [
+                "global_orient",
+                "body_pose",
+                "jaw_pose",
+                "leye_pose",
+                "reye_pose",
+                "left_hand_pose",
+                "right_hand_pose",
+            ]
+        ],
+        dim=-1,
+    )
+    return fullpose
+
+
+def compute_R_from_kinetree(rot_mats, parents):
+    """operation of lbs/batch_rigid_transform, focus on 3x3 R only
+    Parameters
+    ----------
+    rot_mats: torch.tensor BxNx3x3
+        Tensor of rotation matrices
+    parents : torch.tensor BxN
+        The kinematic tree of each object
+
+    Returns
+    -------
+    R : torch.tensor BxNx3x3
+        Tensor of rotation matrices
+    """
+    rot_mat_chain = [rot_mats[:, 0]]
+    for i in range(1, parents.shape[0]):
+        curr_res = torch.matmul(rot_mat_chain[parents[i]], rot_mats[:, i])
+        rot_mat_chain.append(curr_res)
+
+    R = torch.stack(rot_mat_chain, dim=1)
+    return R
+
+
+def compute_relR_from_kinetree(R, parents):
+    """Inverse operation of lbs/batch_rigid_transform, focus on 3x3 R only
+    Parameters
+    ----------
+    R : torch.tensor BxNx4x4 or BxNx3x3
+        Tensor of rotation matrices
+    parents : torch.tensor BxN
+        The kinematic tree of each object
+
+    Returns
+    -------
+    rot_mats: torch.tensor BxNx3x3
+        Tensor of rotation matrices
+    """
+    R = R[:, :, :3, :3]
+
+    Rp = R[:, parents]  # Rp[:, 0] is invalid
+    rot_mats = Rp.transpose(2, 3) @ R
+    rot_mats[:, 0] = R[:, 0]
+
+    return rot_mats
+
+
+def quat_mul(x, y):
+    """
+    Performs quaternion multiplication on arrays of quaternions
+
+    :param x: tensor of quaternions of shape (..., Nb of joints, 4)
+    :param y: tensor of quaternions of shape (..., Nb of joints, 4)
+    :return: The resulting quaternions
+    """
+    x0, x1, x2, x3 = x[..., 0:1], x[..., 1:2], x[..., 2:3], x[..., 3:4]
+    y0, y1, y2, y3 = y[..., 0:1], y[..., 1:2], y[..., 2:3], y[..., 3:4]
+
+    # res = np.concatenate(
+    #     [
+    #         y0 * x0 - y1 * x1 - y2 * x2 - y3 * x3,
+    #         y0 * x1 + y1 * x0 - y2 * x3 + y3 * x2,
+    #         y0 * x2 + y1 * x3 + y2 * x0 - y3 * x1,
+    #         y0 * x3 - y1 * x2 + y2 * x1 + y3 * x0,
+    #     ],
+    #     axis=-1,
+    # )
+    res = torch.cat(
+        [
+            y0 * x0 - y1 * x1 - y2 * x2 - y3 * x3,
+            y0 * x1 + y1 * x0 - y2 * x3 + y3 * x2,
+            y0 * x2 + y1 * x3 + y2 * x0 - y3 * x1,
+            y0 * x3 - y1 * x2 + y2 * x1 + y3 * x0,
+        ],
+        axis=-1,
+    )
+
+    return res
+
+
+def quat_inv(q):
+    """
+    Inverts a tensor of quaternions
+
+    :param q: quaternion tensor
+    :return: tensor of inverted quaternions
+    """
+    # res = np.asarray([1, -1, -1, -1], dtype=np.float32) * q
+    res = torch.tensor([1, -1, -1, -1], device=q.device).float() * q
+    return res
+
+
+def quat_mul_vec(q, x):
+    """
+    Performs multiplication of an array of 3D vectors by an array of quaternions (rotation).
+
+    :param q: tensor of quaternions of shape (..., Nb of joints, 4)
+    :param x: tensor of vectors of shape (..., Nb of joints, 3)
+    :return: the resulting array of rotated vectors
+    """
+    # t = 2.0 * np.cross(q[..., 1:], x)
+    t = 2.0 * torch.cross(q[..., 1:], x)
+    # res = x + q[..., 0][..., np.newaxis] * t + np.cross(q[..., 1:], t)
+    res = x + q[..., 0][..., None] * t + torch.cross(q[..., 1:], t)
+
+    return res
+
+
+def inverse_kinematics_motion(
+    global_pos,
+    global_rot,
+    parents=SMPLH_PARENTS,
+):
+    """
+    Args:
+        global_pos : (B, T, J-1, 3)
+        global_rot (q) : (B, T, J-1, 4)
+        parents : SMPLH_PARENTS
+    Returns:
+        local_pos : (B, T, J-1, 3)
+        local_rot (q) : (B, T, J-1, 4)
+    """
+    J = 22
+    local_pos = quat_mul_vec(
+        quat_inv(global_rot[..., parents[1:J], :]),
+        global_pos - global_pos[..., parents[1:J], :],
+    )
+    local_rot = (quat_mul(quat_inv(global_rot[..., parents[1:J], :]), global_rot),)
+    return local_pos, local_rot
+
+
+def transform_mat(R, t):
+    """Creates a batch of transformation matrices
+    Args:
+        - R: Bx3x3 array of a batch of rotation matrices
+        - t: Bx3x1 array of a batch of translation vectors
+    Returns:
+        - T: Bx4x4 Transformation matrix
+    """
+    # No padding left or right, only add an extra row
+    return torch.cat([F.pad(R, [0, 0, 0, 1]), F.pad(t, [0, 0, 0, 1], value=1)], dim=2)
+
+
+def normalize_joints(joints):
+    """
+    Args:
+        joints: (B, *, J, 3)
+    """
+    LR_hips_xy = joints[..., 2, [0, 1]] - joints[..., 1, [0, 1]]
+    LR_shoulders_xy = joints[..., 17, [0, 1]] - joints[..., 16, [0, 1]]
+    LR_xy = (LR_hips_xy + LR_shoulders_xy) / 2  # (B, *, J, 2)
+
+    x_dir = F.pad(F.normalize(LR_xy, 2, -1), (0, 1), "constant", 0)  # (B, *, 3)
+    z_dir = torch.zeros_like(x_dir)  # (B, *, 3)
+    z_dir[..., 2] = 1
+    y_dir = torch.cross(z_dir, x_dir, dim=-1)
+
+    joints_normalized = (joints - joints[..., [0], :]) @ torch.stack([x_dir, y_dir, z_dir], dim=-1)
+    return joints_normalized
+
+
+@torch.no_grad()
+def compute_Rt_af2az(joints, inverse=False):
+    """Assume z coord is upward
+    Args:
+        joints: (B, J, 3), in the start-frame
+    Returns:
+        R_af2az: (B, 3, 3)
+        t_af2az: (B, 3)
+    """
+    t_af2az = joints[:, 0, :].detach().clone()
+    t_af2az[:, 2] = 0  # do not modify z
+
+    LR_xy = joints[:, 2, [0, 1]] - joints[:, 1, [0, 1]]  # (B, 2)
+    I_mask = LR_xy.pow(2).sum(-1) < 1e-4  # do not rotate, when can't decided the face direction
+    x_dir = F.pad(F.normalize(LR_xy, 2, -1), (0, 1), "constant", 0)  # (B, 3)
+    z_dir = torch.zeros_like(x_dir)
+    z_dir[..., 2] = 1
+    y_dir = torch.cross(z_dir, x_dir, dim=-1)
+    R_af2az = torch.stack([x_dir, y_dir, z_dir], dim=-1)  # (B, 3, 3)
+    R_af2az[I_mask] = torch.eye(3).to(R_af2az)
+
+    if inverse:
+        R_az2af = R_af2az.transpose(1, 2)
+        t_az2af = -(R_az2af @ t_af2az.unsqueeze(2)).squeeze(2)
+        return R_az2af, t_az2af
+    else:
+        return R_af2az, t_af2az
+
+
+def finite_difference_forward(x, dim_t=1, dup_last=True):
+    if dim_t == 1:
+        v = x[:, 1:] - x[:, :-1]
+        if dup_last:
+            v = torch.cat([v, v[:, [-1]]], dim=1)
+    else:
+        raise NotImplementedError
+
+    return v
+
+
+def compute_joints_zero(betas, gender):
+    """
+    Args:
+        betas: (16)
+        gender: 'male' or 'female'
+    Returns:
+        joints_zero: (22, 3)
+    """
+    body_model = {
+        "male": make_smplx(type="humor", gender="male"),
+        "female": make_smplx(type="humor", gender="female"),
+    }
+
+    smpl_params = {
+        "root_orient": torch.zeros((1, 3)),
+        "pose_body": torch.zeros((1, 63)),
+        "betas": betas[None],
+        "trans": torch.zeros(1, 3),
+    }
+    joints_zero = body_model[gender](**smpl_params).Jtr[0, :22]
+    return joints_zero
diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/utils/torch_transform.py b/GR00T-WholeBodyControl/gear_sonic/trl/utils/torch_transform.py
new file mode 100644
index 0000000000000000000000000000000000000000..4c03c758399ccff896b4684e1467b76444961492
--- /dev/null
+++ b/GR00T-WholeBodyControl/gear_sonic/trl/utils/torch_transform.py
@@ -0,0 +1,654 @@
+"""PyTorch quaternion / rotation-matrix arithmetic (scalar-first wxyz convention).
+
+Wraps kornia_transform conversions and adds JIT-compiled helpers for
+quaternion apply, inverse, slerp, and SMPL joint computation.
+"""
+
+# This file assumes w x y z quaternion format
+import os
+import numpy as np
+import torch
+
+# Check environment variable to enable/disable torch.jit.script
+USE_JIT_TORCH_TRANSFORM = os.getenv("USE_JIT_TORCH_TRANSFORM", "1").lower() in ("1", "true", "yes")
+
+
+def conditional_jit_script(func):
+    """Conditionally apply torch.jit.script based on USE_JIT_TORCH_TRANSFORM env var"""
+    if USE_JIT_TORCH_TRANSFORM:
+        return torch.jit.script(func)
+    return func
+
+
+if __name__ != "__main__":
+    from .kornia_transform import (
+        angle_axis_to_quaternion,
+        angle_axis_to_rotation_matrix,
+        quaternion_to_angle_axis,
+        quaternion_to_rotation_matrix,
+        rotation_matrix_to_angle_axis,
+        rotation_matrix_to_quaternion,
+    )
+else:
+    from kornia_transform import (
+        angle_axis_to_quaternion,
+        angle_axis_to_rotation_matrix,
+        quaternion_to_angle_axis,
+        quaternion_to_rotation_matrix,
+        rotation_matrix_to_angle_axis,
+        rotation_matrix_to_quaternion,
+    )
+
+import torch.nn.functional as F
+
+
+def normalize(x, eps: float = 1e-9):
+    return x / x.norm(p=2, dim=-1).clamp(min=eps, max=None).unsqueeze(-1)
+
+
+@conditional_jit_script
+def quat_mul(a, b):
+    assert a.shape == b.shape
+    shape = a.shape
+    a = a.reshape(-1, 4)
+    b = b.reshape(-1, 4)
+
+    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)
+    return torch.stack([w, x, y, z], dim=-1).view(shape)
+
+
+@conditional_jit_script
+def quat_conjugate(a):
+    shape = a.shape
+    a = a.reshape(-1, 4)
+    return torch.cat((a[:, 0:1], -a[:, 1:]), dim=-1).view(shape)
+
+
+@conditional_jit_script
+def quat_inv(a):
+    return normalize(quat_conjugate(a))
+
+
+@conditional_jit_script
+def quat_apply(a, b):
+    shape = b.shape
+    a = a.reshape(-1, 4)
+    b = b.reshape(-1, 3)
+    xyz = a[:, 1:].clone()
+    t = xyz.cross(b, dim=-1) * 2
+    return (b + a[:, 0:1].clone() * t + xyz.cross(t, dim=-1)).view(shape)
+
+
+@conditional_jit_script
+def quat_angle(a, eps: float = 1e-6):
+    shape = a.shape
+    a = a.reshape(-1, 4)
+    s = 2 * (a[:, 0] ** 2) - 1
+    s = s.clamp(-1 + eps, 1 - eps)
+    s = s.acos()
+    return s.view(shape[:-1])
+
+
+@conditional_jit_script
+def quat_angle_diff(quat1, quat2):
+    return quat_angle(quat_mul(quat1, quat_conjugate(quat2)))
+
+
+@conditional_jit_script
+def torch_safe_atan2(y, x, eps: float = 1e-8):
+    y = y.clone()
+    y[(y.abs() < eps) & (x.abs() < eps)] += eps
+    return torch.atan2(y, x)
+
+
+@conditional_jit_script
+def ypr_euler_from_quat(
+    q, handle_singularity: bool = False, eps: float = 1e-6, singular_eps: float = 1e-6
+):
+    """
+    convert quaternion to yaw-pitch-roll euler angles
+    """
+    yaw_atany = 2 * (q[..., 0] * q[..., 3] + q[..., 1] * q[..., 2])
+    yaw_atanx = 1 - 2 * (q[..., 2] * q[..., 2] + q[..., 3] * q[..., 3])
+    roll_atany = 2 * (q[..., 0] * q[..., 1] + q[..., 2] * q[..., 3])
+    roll_atanx = 1 - 2 * (q[..., 1] * q[..., 1] + q[..., 2] * q[..., 2])
+    yaw = torch_safe_atan2(yaw_atany, yaw_atanx, eps)
+    pitch = torch.asin(
+        torch.clamp(
+            2 * (q[..., 0] * q[..., 2] - q[..., 1] * q[..., 3]),
+            min=-1 + eps,
+            max=1 - eps,
+        )
+    )
+    roll = torch_safe_atan2(roll_atany, roll_atanx, eps)
+
+    if handle_singularity:
+        """handle two special cases"""
+        # Gimbal lock detection: test = w*z - x*y approaches ±0.5 when pitch → ±90°
+        test = q[..., 0] * q[..., 2] - q[..., 1] * q[..., 3]
+        # north pole, pitch ~= 90 degrees
+        np_ind = test > 0.5 - singular_eps
+        if torch.any(np_ind):
+            # print('ypr_euler_from_quat singularity -- north pole!')
+            roll[np_ind] = 0.0
+            pitch[np_ind].clamp_max_(0.5 * np.pi)
+            yaw_atany = q[..., 3][np_ind]
+            yaw_atanx = q[..., 0][np_ind]
+            yaw[np_ind] = 2 * torch_safe_atan2(yaw_atany, yaw_atanx, eps)
+        # south pole, pitch ~= -90 degrees
+        sp_ind = test < -0.5 + singular_eps
+        if torch.any(sp_ind):
+            # print('ypr_euler_from_quat singularity -- south pole!')
+            roll[sp_ind] = 0.0
+            pitch[sp_ind].clamp_min_(-0.5 * np.pi)
+            yaw_atany = q[..., 3][sp_ind]
+            yaw_atanx = q[..., 0][sp_ind]
+            yaw[sp_ind] = 2 * torch_safe_atan2(yaw_atany, yaw_atanx, eps)
+
+    return torch.stack([roll, pitch, yaw], dim=-1)
+
+
+@conditional_jit_script
+def quat_from_ypr_euler(angles):
+    """
+    convert yaw-pitch-roll euler angles to quaternion
+    """
+    half_ang = angles * 0.5
+    sin = torch.sin(half_ang)
+    cos = torch.cos(half_ang)
+    q = torch.stack(
+        [
+            cos[..., 0] * cos[..., 1] * cos[..., 2] + sin[..., 0] * sin[..., 1] * sin[..., 2],
+            sin[..., 0] * cos[..., 1] * cos[..., 2] - cos[..., 0] * sin[..., 1] * sin[..., 2],
+            cos[..., 0] * sin[..., 1] * cos[..., 2] + sin[..., 0] * cos[..., 1] * sin[..., 2],
+            cos[..., 0] * cos[..., 1] * sin[..., 2] - sin[..., 0] * sin[..., 1] * cos[..., 2],
+        ],
+        dim=-1,
+    )
+    return q
+
+
+def quat_between_two_vec(v1, v2, eps: float = 1e-6):
+    """
+    quaternion for rotating v1 to v2
+    """
+    orig_shape = v1.shape
+    v1 = v1.reshape(-1, 3)
+    v2 = v2.reshape(-1, 3)
+    dot = (v1 * v2).sum(-1)
+    cross = torch.cross(v1, v2, dim=-1)
+    out = torch.cat([(1 + dot).unsqueeze(-1), cross], dim=-1)
+    # handle v1 & v2 with same direction
+    sind = dot > 1 - eps
+    out[sind] = torch.tensor([1.0, 0.0, 0.0, 0.0], device=v1.device)
+    # handle v1 & v2 with opposite direction
+    nind = dot < -1 + eps
+    if torch.any(nind):
+        vx = torch.tensor([1.0, 0.0, 0.0], device=v1.device)
+        vxdot = (v1 * vx).sum(-1).abs()
+        nxind = nind & (vxdot < 1 - eps)
+        if torch.any(nxind):
+            out[nxind] = angle_axis_to_quaternion(
+                normalize(torch.cross(vx.expand_as(v1[nxind]), v1[nxind], dim=-1)) * np.pi
+            )
+        # handle v1 & v2 with opposite direction and they are parallel to x axis
+        pind = nind & (vxdot >= 1 - eps)
+        if torch.any(pind):
+            vy = torch.tensor([0.0, 1.0, 0.0], device=v1.device)
+            out[pind] = angle_axis_to_quaternion(
+                normalize(torch.cross(vy.expand_as(v1[pind]), v1[pind], dim=-1)) * np.pi
+            )
+    # normalize and reshape
+    out = normalize(out).view(orig_shape[:-1] + (4,))
+    return out
+
+
+@conditional_jit_script
+def get_yaw(q, eps: float = 1e-6):
+    yaw_atany = 2 * (q[..., 0] * q[..., 3] + q[..., 1] * q[..., 2])
+    yaw_atanx = 1 - 2 * (q[..., 2] * q[..., 2] + q[..., 3] * q[..., 3])
+    yaw = torch_safe_atan2(yaw_atany, yaw_atanx, eps)
+    return yaw
+
+
+import torch
+
+
+def swing_twist_decomposition_around_z_torch(
+    q: torch.Tensor,
+    eps: float = 1e-8,
+):
+    """
+    PyTorch version of your SciPy swing-twist decomposition around world Z.
+
+    Args:
+        q: (..., 4) quaternion in [w, x, y, z] order (scalar-first).
+        eps: numerical epsilon.
+
+    Returns:
+        q_swing: (..., 4) quaternion [w, x, y, z] (scalar-first), the swing component.
+        heading: (..., 2) 2D heading vector from the twist rotation matrix: [r00, r10].
+                 (equivalently [cos(yaw), sin(yaw)]).
+        q_twist: (..., 4) quaternion [w, x, y, z] (scalar-first), the twist about Z.
+    """
+    assert q.shape[-1] == 4, "q must have shape (..., 4) in [w,x,y,z] order"
+
+    # Extract components (scalar-first convention)
+    w = q[..., 0]
+    x = q[..., 1]
+    y = q[..., 2]
+    z = q[..., 3]
+
+    # --- Build twist quaternion by projecting vector part onto world Z ---
+    # SciPy code does: q_twist ∝ [0,0,z,w] in (x,y,z,w) order.
+    # In scalar-first [w,x,y,z], that's: [w, 0, 0, z].
+    # Normalize using only w,z (same as norm of [0,0,z,w]).
+    n2 = w * w + z * z
+    inv_n = torch.rsqrt(n2 + eps)
+
+    # Handle degenerate case like SciPy: if norm ~ 0 -> identity
+    # Here: if n2 is extremely small, set to identity twist.
+    deg = n2 < eps
+    w_t = w * inv_n
+    z_t = z * inv_n
+    w_t = torch.where(deg, torch.ones_like(w_t), w_t)
+    z_t = torch.where(deg, torch.zeros_like(z_t), z_t)
+
+    q_twist = torch.stack([w_t, torch.zeros_like(w_t), torch.zeros_like(w_t), z_t], dim=-1)
+
+    # --- Swing = twist^{-1} ⊗ q ---
+    # Quaternion inverse for unit quaternion: conj
+    q_twist_inv = torch.stack([w_t, -torch.zeros_like(w_t), -torch.zeros_like(w_t), -z_t], dim=-1)
+
+    def quat_mul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        """Hamilton product for [w,x,y,z] quaternions."""
+        aw, ax, ay, az = a.unbind(dim=-1)
+        bw, bx, by, bz = b.unbind(dim=-1)
+        return torch.stack(
+            [
+                aw * bw - ax * bx - ay * by - az * bz,
+                aw * bx + ax * bw + ay * bz - az * by,
+                aw * by - ax * bz + ay * bw + az * bx,
+                aw * bz + ax * by - ay * bx + az * bw,
+            ],
+            dim=-1,
+        )
+
+    q_swing = quat_mul(q_twist_inv, q)
+
+    # --- Heading vector from twist rotation matrix ---
+    # For twist about Z with [w,0,0,z], yaw = 2*atan2(z,w)
+    # and the rotated +x axis is [cos(yaw), sin(yaw)].
+    yaw = 2.0 * torch.atan2(z_t, w_t)
+    heading = torch.stack([torch.cos(yaw), torch.sin(yaw)], dim=-1)
+
+    return q_swing, heading, q_twist
+
+
+def swing_twist_decomposition_around_z_np(rot_in):
+    import numpy as np
+    from scipy.spatial.transform import Rotation as R
+
+    quat_in = rot_in.as_quat(scalar_first=False)
+
+    # Project vector part to gravity
+    quat_vec_projected_to_gravity = np.array([0, 0, 1]) * quat_in[2]
+
+    # Take scalar part and append the new projected vector part
+    q_twist = np.append(quat_vec_projected_to_gravity, quat_in[3])
+
+    # Normalize it, now is just a rotation around yaw
+    # There is a degenerate case around 180 degree rotation, where the new z axis points downwards
+    norm = np.linalg.norm(q_twist)
+    if np.isclose(norm, 0):
+        q_twist = np.array([0, 0, 0, 1])
+    else:
+        q_twist = q_twist / norm
+    q_twist = R.from_quat(q_twist, scalar_first=False)
+
+    # q_swing is the rest of the rotation
+    q_swing = q_twist.inv() * rot_in
+
+    # Get heading vector from q_twist represented as rot matrix
+    r_twist = q_twist.as_matrix()
+    heading = np.array([r_twist[0][0], r_twist[1][0]])
+
+    return q_swing, heading
+
+
+@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)
+
+
+@conditional_jit_script
+def get_yaw_q(q):
+    yaw = get_yaw(q)
+    angle_axis = torch.cat(
+        [torch.zeros(yaw.shape + (2,), device=q.device), yaw.unsqueeze(-1)], dim=-1
+    )
+    heading_q = angle_axis_to_quaternion(angle_axis)
+    return heading_q
+
+
+@conditional_jit_script
+def get_heading(q, eps: float = 1e-6):
+    heading_atany = q[..., 3]
+    heading_atanx = q[..., 0]
+    heading = 2 * torch_safe_atan2(heading_atany, heading_atanx, eps)
+    return heading
+
+
+@conditional_jit_script
+def get_heading_twist(q, eps: float = 1e-6):
+    w = q[..., 0]
+    z = q[..., 3]
+    s = torch.rsqrt(w * w + z * z + eps)  # 1/sqrt(...)
+    w = w * s
+    z = z * s
+    return 2 * torch_safe_atan2(z, w, eps)
+
+
+@conditional_jit_script
+def calc_heading_from_projecting_x(q):
+    ref_dir = torch.zeros((q.shape[0], 3), dtype=q.dtype, device=q.device)
+    ref_dir[..., 0] = 1
+    rot_dir = quat_apply(q, ref_dir)
+
+    heading = torch.atan2(rot_dir[..., 1], rot_dir[..., 0])
+    return heading
+
+
+def get_heading_q(q):
+    # Zero out x,y quaternion components to extract pure yaw (Z-axis rotation),
+    # then re-normalize to a valid unit quaternion.
+    # This will cause discontinuities or ill-defined heading when the robot is upside down, which does not often
+    # happen for humanoid robots.
+    q_new = q.clone()
+    q_new[..., 1] = 0
+    q_new[..., 2] = 0
+    q_new = normalize(q_new)
+    return q_new
+
+
+def get_y_heading_q(q):
+    q_new = q.clone()
+    q_new[..., 1] = 0
+    q_new[..., 3] = 0
+    q_new = normalize(q_new)
+    return q_new
+
+
+@conditional_jit_script
+def heading_to_vec(h_theta):
+    v = torch.stack([torch.cos(h_theta), torch.sin(h_theta)], dim=-1)
+    return v
+
+
+@conditional_jit_script
+def vec_to_heading(h_vec):
+    h_theta = torch_safe_atan2(h_vec[..., 1], h_vec[..., 0])
+    return h_theta
+
+
+@conditional_jit_script
+def heading_to_quat(h_theta):
+    angle_axis = torch.cat(
+        [
+            torch.zeros(h_theta.shape + (2,), device=h_theta.device),
+            h_theta.unsqueeze(-1),
+        ],
+        dim=-1,
+    )
+    heading_q = angle_axis_to_quaternion(angle_axis)
+    return heading_q
+
+
+def deheading_quat(q, heading_q=None):
+    if heading_q is None:
+        heading_q = get_heading_q(q)
+    dq = quat_mul(quat_conjugate(heading_q), q)
+    return dq
+
+
+@conditional_jit_script
+def rotmat_to_rot6d(mat):
+    rot6d = torch.cat([mat[..., 0], mat[..., 1]], dim=-1)
+    return rot6d
+
+
+# @conditional_jit_script
+def rot6d_to_rotmat(rot6d, eps: float = 1e-8):
+    a1 = rot6d[..., :3].clone()
+    a2 = rot6d[..., 3:].clone()
+    ind = torch.norm(a1, dim=-1) < eps
+    a1[ind] = torch.tensor([1.0, 0.0, 0.0], device=a1.device)
+    b1 = normalize(a1)
+
+    b2 = normalize(a2 - (b1 * a2).sum(dim=-1).unsqueeze(-1) * b1)
+    ind = torch.norm(b2, dim=-1) < eps
+    b2[ind] = torch.tensor([0.0, 1.0, 0.0], device=b2.device)
+
+    b3 = torch.cross(b1, b2, dim=-1)
+    mat = torch.stack([b1, b2, b3], dim=-1)
+    return mat
+
+
+@conditional_jit_script
+def angle_axis_to_rot6d(aa):
+    return rotmat_to_rot6d(angle_axis_to_rotation_matrix(aa))
+
+
+@conditional_jit_script
+def rot6d_to_angle_axis(rot6d):
+    return rotation_matrix_to_angle_axis(rot6d_to_rotmat(rot6d))
+
+
+@conditional_jit_script
+def quat_to_rot6d(q):
+    return rotmat_to_rot6d(quaternion_to_rotation_matrix(q))
+
+
+@conditional_jit_script
+def rot6d_to_quat(rot6d):
+    return rotation_matrix_to_quaternion(rot6d_to_rotmat(rot6d))
+
+
+@conditional_jit_script
+def make_transform(rot, trans, rot_type: str = "rotmat"):
+    if rot_type == "axis_angle":
+        rot = angle_axis_to_rotation_matrix(rot)
+    elif rot_type == "6d":
+        rot = rot6d_to_rotmat(rot)
+    transform = torch.eye(4).to(trans.device).repeat(rot.shape[:-2] + (1, 1))
+    transform[..., :3, :3] = rot
+    transform[..., :3, 3] = trans
+    return transform
+
+
+@conditional_jit_script
+def transform_trans(transform_mat, trans):
+    trans = torch.cat((trans, torch.ones_like(trans[..., :1])), dim=-1)[..., None, :]
+    while len(transform_mat.shape) < len(trans.shape):
+        transform_mat = transform_mat.unsqueeze(-3)
+    trans_new = torch.matmul(trans, transform_mat.transpose(-2, -1))[..., 0, :3]
+    return trans_new
+
+
+@conditional_jit_script
+def transform_rot(transform_mat, rot):
+    rot_qmat = angle_axis_to_rotation_matrix(rot)
+    while len(transform_mat.shape) < len(rot_qmat.shape):
+        transform_mat = transform_mat.unsqueeze(-3)
+    rot_qmat_new = torch.matmul(transform_mat[..., :3, :3], rot_qmat)
+    rot_new = rotation_matrix_to_angle_axis(rot_qmat_new)
+    return rot_new
+
+
+@conditional_jit_script
+def inverse_transform(transform_mat):
+    transform_inv = torch.zeros_like(transform_mat)
+    transform_inv[..., :3, :3] = transform_mat[..., :3, :3].transpose(-2, -1)
+    transform_inv[..., :3, 3] = -torch.matmul(
+        transform_mat[..., :3, 3].unsqueeze(-2), transform_mat[..., :3, :3]
+    ).squeeze(-2)
+    transform_inv[..., 3, 3] = 1.0
+    return transform_inv
+
+
+def batch_compute_similarity_transform_torch(S1, S2):
+    """
+    Computes a similarity transform (sR, t) that takes
+    a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
+    where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
+    i.e. solves the orthogonal Procrutes problem.
+    """
+    if len(S1.shape) > 3:
+        orig_shape = S1.shape
+        S1 = S1.reshape(-1, *S1.shape[-2:])
+        S2 = S2.reshape(-1, *S2.shape[-2:])
+    else:
+        orig_shape = None
+
+    transposed = False
+    if S1.shape[0] != 3 and S1.shape[0] != 2:
+        S1 = S1.permute(0, 2, 1)
+        S2 = S2.permute(0, 2, 1)
+        transposed = True
+    assert S2.shape[1] == S1.shape[1]
+
+    # 1. Remove mean.
+    mu1 = S1.mean(axis=-1, keepdims=True)
+    mu2 = S2.mean(axis=-1, keepdims=True)
+
+    X1 = S1 - mu1
+    X2 = S2 - mu2
+
+    # 2. Compute variance of X1 used for scale.
+    var1 = torch.sum(X1**2, dim=1).sum(dim=1)
+
+    # 3. The outer product of X1 and X2.
+    K = X1.bmm(X2.permute(0, 2, 1))
+
+    # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are
+    # singular vectors of K.
+    U, s, V = torch.svd(K)
+
+    # Construct Z that fixes the orientation of R to get det(R)=1.
+    Z = torch.eye(U.shape[1], device=S1.device).unsqueeze(0)
+    Z = Z.repeat(U.shape[0], 1, 1)
+    Z[:, -1, -1] *= torch.sign(torch.det(U.bmm(V.permute(0, 2, 1))))
+
+    # Construct R.
+    R = V.bmm(Z.bmm(U.permute(0, 2, 1)))
+
+    # 5. Recover scale.
+    scale = torch.cat([torch.trace(x).unsqueeze(0) for x in R.bmm(K)]) / var1
+
+    # 6. Recover translation.
+    t = mu2 - (scale.unsqueeze(-1).unsqueeze(-1) * (R.bmm(mu1)))
+
+    # 7. Error:
+    S1_hat = scale.unsqueeze(-1).unsqueeze(-1) * R.bmm(S1) + t
+
+    if transposed:
+        S1_hat = S1_hat.permute(0, 2, 1)
+
+    if orig_shape is not None:
+        S1_hat = S1_hat.reshape(orig_shape)
+
+    return S1_hat
+
+
+human_joints_info = None
+
+
+def compute_human_joints(
+    body_pose,
+    global_orient,
+    human_joints_info_path="gear_sonic/data/human/human_joints_info.pkl",
+    use_thumb_joints=True,
+):
+    """
+    Compute SMPL joint positions using forward kinematics.
+
+    Args:
+        body_pose: Body pose in axis-angle format (*, 63)
+        global_orient: Global orientation in axis-angle format (*, 3)
+        J: Rest pose joint positions (55, 3) - from human_joints_info.pkl
+        parents_list: List of parent joint indices - from human_joints_info.pkl
+
+    Returns:
+        posed_joints: Joint positions after applying pose (*, 55, 3)
+    """
+
+    global human_joints_info
+
+    if human_joints_info is None:
+        human_joints_info = torch.load(human_joints_info_path)
+    J = human_joints_info["J"]
+    parents_list = human_joints_info["parents_list"]
+    rot_mats = human_joints_info["rot_mats"]
+
+    device = body_pose.device
+    J = J.to(device)
+
+    # Build full pose: [global_orient(3), body_pose(63), zeros for rest(99)]
+    other_pose = torch.zeros(*body_pose.shape[:-1], 99, device=device)
+    full_pose = torch.cat([global_orient, body_pose, other_pose], dim=-1)
+    rot_mats = angle_axis_to_rotation_matrix(full_pose.reshape(*full_pose.shape[:-1], 55, 3))
+    # rot_mats = axis_angle_to_matrix(full_pose.reshape(*full_pose.shape[:-1], 55, 3))
+
+    # Forward kinematics
+    J = J.expand(*rot_mats.shape[:-3], -1, -1)
+    rel_joints = J.clone()
+    rel_joints[..., 1:, :] -= J[..., parents_list[1:], :]
+
+    transforms_mat = F.pad(
+        torch.cat([rot_mats, rel_joints[..., :, None]], dim=-1), [0, 0, 0, 1], value=0.0
+    )
+    transforms_mat[..., 3, 3] = 1.0
+
+    transform_chain = [transforms_mat[..., 0, :, :]]
+    for i in range(1, len(parents_list)):
+        transform_chain.append(
+            torch.matmul(transform_chain[parents_list[i]], transforms_mat[..., i, :, :])
+        )
+
+    joints = torch.stack(transform_chain, dim=-3)[..., :3, 3]
+
+    # First 22 SMPL joints are the main body; optionally append thumb tips at SMPL indices 39, 54
+    output_joint_index = np.arange(22)
+    if use_thumb_joints:
+        output_joint_index = np.concatenate([output_joint_index, np.array([39, 54])])
+    joints = joints[:, output_joint_index]
+    return joints