Add files using upload-large-folder tool

GR00T-WholeBodyControl/.github/workflows/docs.yml

@@ -0,0 +1,74 @@
+name: Build and Deploy Documentation
+
+on:
+  push:
+    branches:
+      - main
+      - gear-sonic
+    paths:
+      - "docs/**"
+      - ".github/workflows/docs.yml"
+      - ".gitattributes"
+  workflow_dispatch:
+
+# Allow only one concurrent deployment; cancel in-flight runs
+concurrency:
+  group: "pages"
+  cancel-in-progress: true
+
+jobs:
+  build:
+    name: Build Sphinx Docs
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          lfs: false
+
+      - name: Restore docs static assets (bypass git-lfs smudge)
+        run: |
+          # git-lfs on the runner rewrites files tracked by *.png/*.gif
+          # even when our .gitattributes override removes filter=lfs.
+          # Use git cat-file to write real binary content directly from
+          # the object store, bypassing all smudge filters.
+          git ls-tree -r HEAD -- docs/source/_static \
+            | awk '{print $3, $4}' \
+            | while IFS=" " read -r hash path; do
+                git cat-file blob "$hash" > "$path"
+              done
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+          cache: "pip"
+          cache-dependency-path: "docs/requirements.txt"
+
+      - name: Install documentation dependencies
+        run: pip install -r docs/requirements.txt
+
+      - name: Build HTML documentation
+        run: sphinx-build -b html docs/source docs/build/html
+
+      - name: Upload Pages artifact
+        if: github.ref == 'refs/heads/main' || github.ref == 'refs/heads/gear-sonic'
+        uses: actions/upload-pages-artifact@v3
+        with:
+          path: docs/build/html
+
+  deploy:
+    name: Deploy to GitHub Pages
+    needs: build
+    if: github.ref == 'refs/heads/main' || github.ref == 'refs/heads/gear-sonic'
+    runs-on: ubuntu-latest
+    permissions:
+      pages: write
+      id-token: write
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+    steps:
+      - name: Deploy to GitHub Pages
+        id: deployment
+        uses: actions/deploy-pages@v4

GR00T-WholeBodyControl/decoupled_wbc/control/base/env.py

@@ -0,0 +1,45 @@
+import gymnasium as gym
+
+
+class Env:
+    """Base interface for all environments in the Gr00t framework"""
+
+    def observe(self) -> dict[str, any]:
+        """Read the current state of this environment
+
+        Returns:
+            dict: A dictionary of observations
+        """
+        pass
+
+    def queue_action(self, action: dict[str, any]):
+        """Queue an action to be executed
+
+        Args:
+            action: A dictionary of action parameters
+        """
+        pass
+
+    def reset(self, **kwargs):
+        """Reset this environment to initial state"""
+        pass
+
+    def observation_space(self) -> gym.Space:
+        """Get the observation space of this environment
+
+        Returns:
+            gym.Space: The observation space
+        """
+        pass
+
+    def action_space(self) -> gym.Space:
+        """Get the action space of this environment
+
+        Returns:
+            gym.Space: The action space
+        """
+        pass
+
+    def close(self):
+        """Close and clean up this environment"""
+        pass

GR00T-WholeBodyControl/decoupled_wbc/control/base/humanoid_env.py

@@ -0,0 +1,60 @@
+from abc import abstractmethod
+
+from decoupled_wbc.control.base.env import Env
+from decoupled_wbc.control.base.sensor import Sensor
+from decoupled_wbc.control.robot_model.robot_model import RobotModel
+
+
+class Hands:
+    """Container class for left and right hand environments.
+
+    Attributes:
+        left: Environment for the left hand
+        right: Environment for the right hand
+    """
+
+    left: Env
+    right: Env
+
+
+class HumanoidEnv(Env):
+    """Base class for humanoid robot environments.
+
+    This class provides the interface for accessing the robot's body, hands, and sensors.
+    """
+
+    def body(self) -> Env:
+        """Get the robot's body environment.
+
+        Returns:
+            Env: The body environment
+        """
+        pass
+
+    def hands(self) -> Hands:
+        """Get the robot's hands.
+
+        Returns:
+            Hands: Container with left and right hand environments
+        """
+        pass
+
+    def sensors(self) -> dict[str, Sensor]:
+        """Get the sensors of this environment
+
+        Returns:
+            dict: A dictionary of sensors
+        """
+        pass
+
+    @abstractmethod
+    def robot_model(self) -> RobotModel:
+        """Get the robot model of this environment
+        This robot model is used to dispatch whole body actions to body
+        and hand actuators and to reconstruct proprioceptive
+        observations from body and hands.
+
+        Returns:
+            RobotModel: The robot model
+        """
+        pass

GR00T-WholeBodyControl/decoupled_wbc/control/base/policy.py

@@ -0,0 +1,47 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+
+class Policy(ABC):
+    """Base class for implementing control policies in the Gr00t framework.
+
+    A Policy defines how an agent should behave in an environment by mapping observations
+    to actions. This abstract base class provides the interface that all concrete policy
+    implementations must follow.
+    """
+
+    def set_goal(self, goal: dict[str, any]):
+        """Set the command from the planner that the policy should follow.
+
+        Args:
+            goal: Dictionary containing high-level commands or goals from the planner
+        """
+        pass
+
+    def set_observation(self, observation: dict[str, any]):
+        """Update the policy's current observation of the environment.
+
+        Args:
+            observation: Dictionary containing the current state/observation of the environment
+        """
+        self.observation = observation
+
+    @abstractmethod
+    def get_action(self, time: Optional[float] = None) -> dict[str, any]:
+        """Compute and return the next action at the specified time, based on current observation
+        and planner command.
+
+        Args:
+            time: Optional "monotonic time" for time-dependent policies
+
+        Returns:
+            Dictionary containing the action to be executed
+        """
+
+    def close(self):
+        """Clean up any resources used by the policy."""
+        pass
+
+    def reset(self):
+        """Reset the policy to its initial state."""
+        pass

GR00T-WholeBodyControl/decoupled_wbc/control/base/sensor.py

@@ -0,0 +1,35 @@
+import gymnasium as gym
+
+
+class Sensor:
+    """Base class for implementing sensors in the Gr00t framework.
+
+    A Sensor provides information about a specific sensor on the robot (e.g. camera, IMU,
+    force sensor). This abstract base class defines the interface that all concrete sensor
+    implementations must follow.
+    """
+
+    def read(self, **kwargs) -> any:
+        """Read the current sensor value.
+
+        Args:
+            **kwargs: Additional parameters specific to the sensor implementation
+                (e.g. camera resolution, sampling rate)
+
+        Returns:
+            The sensor reading value (e.g. image data, acceleration measurements)
+        """
+        pass
+
+    def observation_space(self) -> gym.Space:
+        """Get the observation space of this sensor.
+
+        Returns:
+            gym.Space: The observation space defining the shape and bounds of sensor readings
+                (e.g. image dimensions for camera, measurement ranges for IMU)
+        """
+        pass
+
+    def close(self):
+        """Clean up any resources used by the sensor."""
+        pass

GR00T-WholeBodyControl/decoupled_wbc/control/main/config_template.py

@@ -0,0 +1,45 @@
+from dataclasses import asdict, dataclass
+from typing import Any
+
+
+@dataclass
+class ArgsConfig:
+    """Args Config for running the data collection loop."""
+
+    def update(
+        self,
+        config_dict: dict,
+        strict: bool = False,
+        skip_keys: list[str] = [],
+        allowed_keys: list[str] | None = None,
+    ):
+        for k, v in config_dict.items():
+            if k in skip_keys:
+                continue
+            if allowed_keys is not None and k not in allowed_keys:
+                continue
+            if strict and not hasattr(self, k):
+                raise ValueError(f"Config {k} not found in {self.__class__.__name__}")
+            if not strict and not hasattr(self, k):
+                continue
+            setattr(self, k, v)
+
+    @classmethod
+    def from_dict(
+        cls,
+        config_dict: dict,
+        strict: bool = False,
+        skip_keys: list[str] = [],
+        allowed_keys: list[str] | None = None,
+    ):
+        instance = cls()
+        instance.update(
+            config_dict=config_dict, strict=strict, skip_keys=skip_keys, allowed_keys=allowed_keys
+        )
+        return instance
+
+    def to_dict(self):
+        return asdict(self)
+
+    def get(self, key: str, default: Any = None):
+        return getattr(self, key) if hasattr(self, key) else default

GR00T-WholeBodyControl/decoupled_wbc/control/main/constants.py

@@ -0,0 +1,16 @@
+IMAGE_TOPIC_NAME = "realsense/color/image_raw"
+STATE_TOPIC_NAME = "G1Env/env_state_act"
+CONTROL_GOAL_TOPIC = "ControlPolicy/upper_body_pose"
+ROBOT_CONFIG_TOPIC = "WBCPolicy/robot_config"
+KEYBOARD_INPUT_TOPIC = "/keyboard_input"
+LOCO_MANIP_TASK_STATUS_TOPIC = "LocoManipPolicy/task_status"
+LOCO_NAV_TASK_STATUS_TOPIC = "NavigationPolicy/task_status"
+LOWER_BODY_POLICY_STATUS_TOPIC = "ControlPolicy/lower_body_policy_status"
+JOINT_SAFETY_STATUS_TOPIC = "ControlPolicy/joint_safety_status"
+
+
+DEFAULT_NAV_CMD = [0.0, 0.0, 0.0]
+DEFAULT_BASE_HEIGHT = 0.74
+DEFAULT_WRIST_POSE = [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0] * 2  # x, y, z + w, x, y, z
+
+DEFAULT_MODEL_SERVER_PORT = 5555  # port used to host the model server

GR00T-WholeBodyControl/decoupled_wbc/control/policy/g1_decoupled_whole_body_policy.py

@@ -0,0 +1,157 @@
+import time as time_module
+from typing import Optional
+
+import numpy as np
+from pinocchio import rpy
+
+from decoupled_wbc.control.base.policy import Policy
+from decoupled_wbc.control.main.constants import DEFAULT_NAV_CMD
+
+
+class G1DecoupledWholeBodyPolicy(Policy):
+    """
+    This class implements a whole-body policy for the G1 robot by combining an upper-body
+    policy and a lower-body RL-based policy.
+    It is designed to work with the G1 robot's specific configuration and control requirements.
+    """
+
+    def __init__(
+        self,
+        robot_model,
+        lower_body_policy: Policy,
+        upper_body_policy: Policy,
+    ):
+        self.robot_model = robot_model
+        self.lower_body_policy = lower_body_policy
+        self.upper_body_policy = upper_body_policy
+        self.last_goal_time = time_module.monotonic()
+        self.is_in_teleop_mode = False  # Track if lower body is in teleop mode
+
+    def set_observation(self, observation):
+        # Upper body policy is open loop (just interpolation), so we don't need to set the observation
+        self.lower_body_policy.set_observation(observation)
+
+    def set_goal(self, goal):
+        """
+        Set the goal for both upper and lower body policies.
+
+        Args:
+            goal: Command from the planners
+            goal["target_upper_body_pose"]: Target pose for the upper body policy
+            goal["target_time"]: Target goal time
+            goal["interpolation_garbage_collection_time"]: Waypoints earlier than this time are removed
+            goal["navigate_cmd"]: Target navigation velocities for the lower body policy
+            goal["base_height_command"]: Target base height for both upper and lower body policies
+        """
+        # Update goal timestamp for timeout safety
+        self.last_goal_time = time_module.monotonic()
+
+        upper_body_goal = {}
+        lower_body_goal = {}
+
+        # Upper body goal keys
+        upper_body_keys = [
+            "target_upper_body_pose",
+            "base_height_command",
+            "target_time",
+            "interpolation_garbage_collection_time",
+            "navigate_cmd",
+        ]
+        for key in upper_body_keys:
+            if key in goal:
+                upper_body_goal[key] = goal[key]
+
+        # Always ensure navigate_cmd is present to prevent interpolation from old dangerous values
+        if "navigate_cmd" not in goal:
+            # Safety: Inject safe default navigate_cmd to ensure interpolation goes to stop
+            if "target_time" in goal and isinstance(goal["target_time"], list):
+                upper_body_goal["navigate_cmd"] = [np.array(DEFAULT_NAV_CMD)] * len(
+                    goal["target_time"]
+                )
+            else:
+                upper_body_goal["navigate_cmd"] = np.array(DEFAULT_NAV_CMD)
+
+        # Set teleop policy command flag
+        has_teleop_commands = ("navigate_cmd" in goal) or ("base_height_command" in goal)
+        self.is_in_teleop_mode = has_teleop_commands  # Track teleop state for timeout safety
+        self.lower_body_policy.set_use_teleop_policy_cmd(has_teleop_commands)
+
+        # Lower body goal keys
+        lower_body_keys = [
+            "toggle_stand_command",
+            "toggle_policy_action",
+        ]
+        for key in lower_body_keys:
+            if key in goal:
+                lower_body_goal[key] = goal[key]
+
+        self.upper_body_policy.set_goal(upper_body_goal)
+        self.lower_body_policy.set_goal(lower_body_goal)
+
+    def get_action(self, time: Optional[float] = None):
+        current_time = time if time is not None else time_module.monotonic()
+
+        # Safety timeout: Only apply when in teleop mode (communication loss dangerous)
+        # When in keyboard mode, no timeout needed (user controls directly)
+        if self.is_in_teleop_mode:
+            time_since_goal = current_time - self.last_goal_time
+            if time_since_goal > 1.0:  # 1 second timeout
+                print(
+                    f"SAFETY: Teleop mode timeout after {time_since_goal:.1f}s, injecting safe goal"
+                )
+                # Inject safe goal to trigger all safety mechanisms (gear_wbc reset + interpolation reset)
+                safe_goal = {
+                    "target_time": current_time + 0.1,
+                    "interpolation_garbage_collection_time": current_time - 1.0,
+                }
+                self.set_goal(
+                    safe_goal
+                )  # This will reset is_in_teleop_mode to False and trigger all safety
+
+        # Get indices for groups
+        lower_body_indices = self.robot_model.get_joint_group_indices("lower_body")
+        upper_body_indices = self.robot_model.get_joint_group_indices("upper_body")
+
+        # Initialize full configuration with zeros
+        q = np.zeros(self.robot_model.num_dofs)
+
+        upper_body_action = self.upper_body_policy.get_action(time)
+        q[upper_body_indices] = upper_body_action["target_upper_body_pose"]
+        q_arms = q[self.robot_model.get_joint_group_indices("arms")]
+        base_height_command = upper_body_action.get("base_height_command", None)
+        interpolated_navigate_cmd = upper_body_action.get("navigate_cmd", None)
+
+        # Compute torso orientation relative to waist, to pass to lower body policy
+        self.robot_model.cache_forward_kinematics(q, auto_clip=False)
+        torso_orientation = self.robot_model.frame_placement("torso_link").rotation
+        waist_orientation = self.robot_model.frame_placement("pelvis").rotation
+        # Extract yaw from rotation matrix and create a rotation with only yaw
+        # The rotation property is a 3x3 numpy array
+        waist_yaw = np.arctan2(waist_orientation[1, 0], waist_orientation[0, 0])
+        # Create a rotation matrix with only yaw using Pinocchio's rpy functions
+        waist_yaw_only_rotation = rpy.rpyToMatrix(0, 0, waist_yaw)
+        yaw_only_waist_from_torso = waist_yaw_only_rotation.T @ torso_orientation
+        torso_orientation_rpy = rpy.matrixToRpy(yaw_only_waist_from_torso)
+
+        lower_body_action = self.lower_body_policy.get_action(
+            time, q_arms, base_height_command, torso_orientation_rpy, interpolated_navigate_cmd
+        )
+
+        # If pelvis is both in upper and lower body, lower body policy takes preference
+        q[lower_body_indices] = lower_body_action["body_action"][0][
+            : len(lower_body_indices)
+        ]  # lower body (legs + waist)
+
+        self.last_action = {"q": q}
+
+        return {"q": q}
+
+    def handle_keyboard_button(self, key):
+        try:
+            self.lower_body_policy.locomotion_policy.handle_keyboard_button(key)
+        except AttributeError:
+            # Only catch AttributeError, let other exceptions propagate
+            self.lower_body_policy.handle_keyboard_button(key)
+
+    def activate_policy(self):
+        self.handle_keyboard_button("]")

GR00T-WholeBodyControl/decoupled_wbc/control/policy/g1_gear_wbc_policy.py

@@ -0,0 +1,295 @@
+import collections
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import numpy as np
+import onnxruntime as ort
+import torch
+
+from decoupled_wbc.control.base.policy import Policy
+from decoupled_wbc.control.utils.gear_wbc_utils import get_gravity_orientation, load_config
+
+
+class G1GearWbcPolicy(Policy):
+    """Simple G1 robot policy using OpenGearWbc trained neural network."""
+
+    def __init__(self, robot_model, config: str, model_path: str):
+        """Initialize G1GearWbcPolicy.
+
+        Args:
+            config_path: Path to gear_wbc YAML configuration file
+        """
+        self.config, self.LEGGED_GYM_ROOT_DIR = load_config(config)
+        self.robot_model = robot_model
+        self.use_teleop_policy_cmd = False
+
+        package_root = Path(__file__).resolve().parents[2]
+        self.sim2mujoco_root_dir = str(package_root / "sim2mujoco")
+        model_path_1, model_path_2 = model_path.split(",")
+
+        self.policy_1 = self.load_onnx_policy(
+            self.sim2mujoco_root_dir + "/resources/robots/g1/" + model_path_1
+        )
+        self.policy_2 = self.load_onnx_policy(
+            self.sim2mujoco_root_dir + "/resources/robots/g1/" + model_path_2
+        )
+
+        # Initialize observation history buffer
+        self.observation = None
+        self.obs_history = collections.deque(maxlen=self.config["obs_history_len"])
+        self.obs_buffer = np.zeros(self.config["num_obs"], dtype=np.float32)
+        self.counter = 0
+
+        # Initialize state variables
+        self.use_policy_action = False
+        self.action = np.zeros(self.config["num_actions"], dtype=np.float32)
+        self.target_dof_pos = self.config["default_angles"].copy()
+        self.cmd = self.config["cmd_init"].copy()
+        self.height_cmd = self.config["height_cmd"]
+        self.freq_cmd = self.config["freq_cmd"]
+        self.roll_cmd = self.config["rpy_cmd"][0]
+        self.pitch_cmd = self.config["rpy_cmd"][1]
+        self.yaw_cmd = self.config["rpy_cmd"][2]
+        self.gait_indices = torch.zeros((1), dtype=torch.float32)
+
+    def load_onnx_policy(self, model_path: str):
+        print(f"Loading ONNX policy from {model_path}")
+        model = ort.InferenceSession(model_path)
+
+        def run_inference(input_tensor):
+            ort_inputs = {model.get_inputs()[0].name: input_tensor.cpu().numpy()}
+            ort_outs = model.run(None, ort_inputs)
+            return torch.tensor(ort_outs[0], device="cpu")
+
+        print(f"Successfully loaded ONNX policy from {model_path}")
+
+        return run_inference
+
+    def compute_observation(self, observation: Dict[str, Any]) -> tuple[np.ndarray, int]:
+        """Compute the observation vector from current state"""
+        # Get body joint indices (excluding waist roll and pitch)
+        self.gait_indices = torch.remainder(self.gait_indices + 0.02 * self.freq_cmd, 1.0)
+        durations = torch.full_like(self.gait_indices, 0.5)
+        phases = 0.5
+        foot_indices = [
+            self.gait_indices + phases,  # FL
+            self.gait_indices,  # FR
+        ]
+        self.foot_indices = torch.remainder(
+            torch.cat([foot_indices[i].unsqueeze(1) for i in range(2)], dim=1), 1.0
+        )
+        for fi in foot_indices:
+            stance = fi < durations
+            swing = fi >= durations
+            fi[stance] = fi[stance] * (0.5 / durations[stance])
+            fi[swing] = 0.5 + (fi[swing] - durations[swing]) * (0.5 / (1 - durations[swing]))
+
+        self.clock_inputs = torch.stack([torch.sin(2 * np.pi * fi) for fi in foot_indices], dim=1)
+
+        body_indices = self.robot_model.get_joint_group_indices("body")
+        body_indices = [idx for idx in body_indices]
+
+        n_joints = len(body_indices)
+
+        # Extract joint data
+        qj = observation["q"][body_indices].copy()
+        dqj = observation["dq"][body_indices].copy()
+
+        # Extract floating base data
+        quat = observation["floating_base_pose"][3:7].copy()  # quaternion
+        omega = observation["floating_base_vel"][3:6].copy()  # angular velocity
+
+        # Handle default angles padding
+        if len(self.config["default_angles"]) < n_joints:
+            padded_defaults = np.zeros(n_joints, dtype=np.float32)
+            padded_defaults[: len(self.config["default_angles"])] = self.config["default_angles"]
+        else:
+            padded_defaults = self.config["default_angles"][:n_joints]
+
+        # Scale the values
+        qj_scaled = (qj - padded_defaults) * self.config["dof_pos_scale"]
+        dqj_scaled = dqj * self.config["dof_vel_scale"]
+        gravity_orientation = get_gravity_orientation(quat)
+        omega_scaled = omega * self.config["ang_vel_scale"]
+
+        # Calculate single observation dimension
+        single_obs_dim = 86  # 3 + 1 + 3 + 3 + 3 + n_joints + n_joints + 15, n_joints = 29
+
+        # Create single observation
+        single_obs = np.zeros(single_obs_dim, dtype=np.float32)
+        single_obs[0:3] = self.cmd[:3] * self.config["cmd_scale"]
+        single_obs[3:4] = np.array([self.height_cmd])
+        single_obs[4:7] = np.array([self.roll_cmd, self.pitch_cmd, self.yaw_cmd])
+        single_obs[7:10] = omega_scaled
+        single_obs[10:13] = gravity_orientation
+        # single_obs[14:17] = omega_scaled_torso
+        # single_obs[17:20] = gravity_torso
+        single_obs[13 : 13 + n_joints] = qj_scaled
+        single_obs[13 + n_joints : 13 + 2 * n_joints] = dqj_scaled
+        single_obs[13 + 2 * n_joints : 13 + 2 * n_joints + 15] = self.action
+        # single_obs[13 + 2 * n_joints + 15 : 13 + 2 * n_joints + 15 + 2] = (
+        #     processed_clock_inputs.detach().cpu().numpy()
+        # )
+        return single_obs, single_obs_dim
+
+    def set_observation(self, observation: Dict[str, Any]):
+        """Update the policy's current observation of the environment.
+
+        Args:
+            observation: Dictionary containing single observation from current state
+                        Should include 'obs' key with current single observation
+        """
+
+        # Extract the single observation
+        self.observation = observation
+        single_obs, single_obs_dim = self.compute_observation(observation)
+
+        # Update observation history every control_decimation steps
+        # if self.counter % self.config['control_decimation'] == 0:
+        # Add current observation to history
+        self.obs_history.append(single_obs)
+
+        # Fill history with zeros if not enough observations yet
+        while len(self.obs_history) < self.config["obs_history_len"]:
+            self.obs_history.appendleft(np.zeros_like(single_obs))
+
+        # Construct full observation with history
+        single_obs_dim = len(single_obs)
+        for i, hist_obs in enumerate(self.obs_history):
+            start_idx = i * single_obs_dim
+            end_idx = start_idx + single_obs_dim
+            self.obs_buffer[start_idx:end_idx] = hist_obs
+
+        # Convert to tensor for policy
+        self.obs_tensor = torch.from_numpy(self.obs_buffer).unsqueeze(0)
+        # self.counter += 1
+
+        assert self.obs_tensor.shape[1] == self.config["num_obs"]
+
+    def set_use_teleop_policy_cmd(self, use_teleop_policy_cmd: bool):
+        self.use_teleop_policy_cmd = use_teleop_policy_cmd
+        # Safety: When teleop is disabled, reset navigation to stop
+        if not use_teleop_policy_cmd:
+            self.nav_cmd = self.config["cmd_init"].copy()  # Reset to safe default
+
+    def set_goal(self, goal: Dict[str, Any]):
+        """Set the goal for the policy.
+
+        Args:
+            goal: Dictionary containing the goal for the policy
+        """
+
+        if "toggle_policy_action" in goal:
+            if goal["toggle_policy_action"]:
+                self.use_policy_action = not self.use_policy_action
+
+    def get_action(
+        self,
+        time: Optional[float] = None,
+        arms_target_pose: Optional[np.ndarray] = None,
+        base_height_command: Optional[np.ndarray] = None,
+        torso_orientation_rpy: Optional[np.ndarray] = None,
+        interpolated_navigate_cmd: Optional[np.ndarray] = None,
+    ) -> Dict[str, Any]:
+        """Compute and return the next action based on current observation.
+
+        Args:
+            time: Optional "monotonic time" for time-dependent policies (unused)
+
+        Returns:
+            Dictionary containing the action to be executed
+        """
+        if self.obs_tensor is None:
+            raise ValueError("No observation set. Call set_observation() first.")
+
+        if base_height_command is not None and self.use_teleop_policy_cmd:
+            self.height_cmd = (
+                base_height_command[0]
+                if isinstance(base_height_command, list)
+                else base_height_command
+            )
+
+        if interpolated_navigate_cmd is not None and self.use_teleop_policy_cmd:
+            self.cmd = interpolated_navigate_cmd
+
+        if torso_orientation_rpy is not None and self.use_teleop_policy_cmd:
+            self.roll_cmd = torso_orientation_rpy[0]
+            self.pitch_cmd = torso_orientation_rpy[1]
+            self.yaw_cmd = torso_orientation_rpy[2]
+
+        # Run policy inference
+        with torch.no_grad():
+            # Select appropriate policy based on command magnitude
+            if np.linalg.norm(self.cmd) < 0.05:
+                # Use standing policy for small commands
+                policy = self.policy_1
+            else:
+                # Use walking policy for movement commands
+                policy = self.policy_2
+
+            self.action = policy(self.obs_tensor).detach().numpy().squeeze()
+
+        # Transform action to target_dof_pos
+        if self.use_policy_action:
+            cmd_q = self.action * self.config["action_scale"] + self.config["default_angles"]
+        else:
+            cmd_q = self.observation["q"][self.robot_model.get_joint_group_indices("lower_body")]
+
+        cmd_dq = np.zeros(self.config["num_actions"])
+        cmd_tau = np.zeros(self.config["num_actions"])
+
+        return {"body_action": (cmd_q, cmd_dq, cmd_tau)}
+
+    def handle_keyboard_button(self, key):
+        if key == "]":
+            self.use_policy_action = True
+        elif key == "o":
+            self.use_policy_action = False
+        elif key == "w":
+            self.cmd[0] += 0.2
+        elif key == "s":
+            self.cmd[0] -= 0.2
+        elif key == "a":
+            self.cmd[1] += 0.2
+        elif key == "d":
+            self.cmd[1] -= 0.2
+        elif key == "q":
+            self.cmd[2] += 0.2
+        elif key == "e":
+            self.cmd[2] -= 0.2
+        elif key == "z":
+            self.cmd[0] = 0.0
+            self.cmd[1] = 0.0
+            self.cmd[2] = 0.0
+        elif key == "1":
+            self.height_cmd += 0.1
+        elif key == "2":
+            self.height_cmd -= 0.1
+        elif key == "n":
+            self.freq_cmd -= 0.1
+            self.freq_cmd = max(1.0, self.freq_cmd)
+        elif key == "m":
+            self.freq_cmd += 0.1
+            self.freq_cmd = min(2.0, self.freq_cmd)
+        elif key == "3":
+            self.roll_cmd -= np.deg2rad(10)
+        elif key == "4":
+            self.roll_cmd += np.deg2rad(10)
+        elif key == "5":
+            self.pitch_cmd -= np.deg2rad(10)
+        elif key == "6":
+            self.pitch_cmd += np.deg2rad(10)
+        elif key == "7":
+            self.yaw_cmd -= np.deg2rad(10)
+        elif key == "8":
+            self.yaw_cmd += np.deg2rad(10)
+
+        if key:
+            print("--------------------------------")
+            print(f"Linear velocity command: {self.cmd}")
+            print(f"Base height command: {self.height_cmd}")
+            print(f"Use policy action: {self.use_policy_action}")
+            print(f"roll deg angle: {np.rad2deg(self.roll_cmd)}")
+            print(f"pitch deg angle: {np.rad2deg(self.pitch_cmd)}")
+            print(f"yaw deg angle: {np.rad2deg(self.yaw_cmd)}")
+            print(f"Gait frequency: {self.freq_cmd}")

GR00T-WholeBodyControl/decoupled_wbc/control/policy/identity_policy.py

@@ -0,0 +1,25 @@
+from copy import deepcopy
+from typing import Optional
+
+import gymnasium as gym
+
+from decoupled_wbc.control.base.policy import Policy
+
+
+class IdentityPolicy(Policy):
+    def __init__(self):
+        self.reset()
+
+    def get_action(self, time: Optional[float] = None) -> dict[str, any]:
+        return self.goal
+
+    def set_goal(self, goal: dict[str, any]) -> None:
+        self.goal = deepcopy(goal)
+        self.goal.pop("interpolation_garbage_collection_time", None)
+        self.goal.pop("target_time", None)
+
+    def observation_space(self) -> gym.spaces.Dict:
+        return gym.spaces.Dict()
+
+    def action_space(self) -> gym.spaces.Dict:
+        return gym.spaces.Dict()

GR00T-WholeBodyControl/decoupled_wbc/control/policy/interpolation_policy.py

@@ -0,0 +1,297 @@
+import numbers
+import time as time_module
+from typing import Any, Dict, Optional, Union
+
+import gymnasium as gym
+import numpy as np
+import scipy.interpolate as si
+
+from decoupled_wbc.control.base.policy import Policy
+
+
+class InterpolationPolicy(Policy):
+    def __init__(
+        self,
+        init_time: float,
+        init_values: dict[str, np.ndarray],
+        max_change_rate: float,
+    ):
+        """
+        Args:
+            init_time: The time of recording the initial values.
+            init_values: The initial values of the features.
+                The keys are the names of the features, and the values
+                are the initial values of the features (1D array).
+            max_change_rate: The maximum change rate.
+        """
+        super().__init__()
+        self.last_action = init_values  # Vecs are 1D arrays
+        self.concat_order = sorted(init_values.keys())
+        self.concat_dims = []
+        for key in self.concat_order:
+            vec = np.array(init_values[key])
+            if vec.ndim == 2 and vec.shape[0] == 1:
+                vec = vec[0]
+                init_values[key] = vec
+            assert vec.ndim == 1, f"The shape of {key} should be (D,). Got {vec.shape}."
+            self.concat_dims.append(vec.shape[0])
+
+        self.init_values_concat = self._concat_vecs(init_values, 1)
+        self.max_change_rate = max_change_rate
+        self.reset(init_time)
+
+    def reset(self, init_time: float = time_module.monotonic()):
+        self.interp = PoseTrajectoryInterpolator(np.array([init_time]), self.init_values_concat)
+        self.last_waypoint_time = init_time
+        self.max_change_rate = self.max_change_rate
+
+    def _concat_vecs(self, values: dict[str, np.ndarray], length: int) -> np.ndarray:
+        """
+        Concatenate the vectors into a 2D array to be used for interpolation.
+        Args:
+            values: The values to concatenate.
+            length: The length of the concatenated vectors (time dimension).
+        Returns:
+            The concatenated vectors (T, D) arrays.
+        """
+        concat_vecs = []
+        for key in self.concat_order:
+            if key in values:
+                vec = np.array(values[key])
+                if vec.ndim == 1:
+                    # If the vector is 1D, tile it to the length of the time dimension
+                    vec = np.tile(vec, (length, 1))
+                assert vec.ndim == 2, f"The shape of {key} should be (T, D). Got {vec.shape}."
+                concat_vecs.append(vec)
+            else:
+                # If the vector is not in the values, use the last action
+                # Since the last action is 1D, we need to tile it to the length of the time dimension
+                concat_vecs.append(np.tile(self.last_action[key], (length, 1)))
+        return np.concatenate(concat_vecs, axis=1)  # Vecs are 2D (T, D) arrays
+
+    def _unconcat_vecs(self, concat_vec: np.ndarray) -> dict[str, np.ndarray]:
+        curr_idx = 0
+        action = {}
+        assert (
+            concat_vec.ndim == 1
+        ), f"The shape of the concatenated vectors should be (T, D). Got {concat_vec.shape}."
+        for key, dim in zip(self.concat_order, self.concat_dims):
+            action[key] = concat_vec[curr_idx : curr_idx + dim]
+            curr_idx += dim
+        return action  # Vecs are 1D arrays
+
+    def __call__(
+        self, observation: Dict[str, Any], goal: Dict[str, Any], time: float
+    ) -> Dict[str, np.ndarray]:
+        raise NotImplementedError(
+            "`InterpolationPolicy` accepts goal and provide action in two separate methods."
+        )
+
+    def set_goal(self, goal: Dict[str, Any]) -> None:
+        if "target_time" not in goal:
+            return
+        assert (
+            "interpolation_garbage_collection_time" in goal
+        ), "`interpolation_garbage_collection_time` is required."
+        target_time = goal.pop("target_time")
+        interpolation_garbage_collection_time = goal.pop("interpolation_garbage_collection_time")
+
+        if isinstance(target_time, list):
+            for key, vec in goal.items():
+                assert isinstance(vec, list)
+                assert len(vec) == len(target_time), (
+                    f"The length of {key} and `target_time` should be the same. "
+                    f"Got {len(vec)} and {len(target_time)}."
+                )
+        else:
+            target_time = [target_time]
+            for key in goal:
+                goal[key] = [goal[key]]
+
+        # Concatenate all vectors in goal
+        concat_vecs = self._concat_vecs(goal, len(target_time))
+        assert concat_vecs.shape[0] == len(target_time), (
+            f"The length of the concatenated goal and `target_time` should be the same. "
+            f"Got {concat_vecs.shape[0]} and {len(target_time)}."
+        )
+
+        for tt, vec in zip(target_time, concat_vecs):
+            if tt < interpolation_garbage_collection_time:
+                continue
+            self.interp = self.interp.schedule_waypoint(
+                pose=vec,
+                time=tt,
+                max_change_rate=self.max_change_rate,
+                interpolation_garbage_collection_time=interpolation_garbage_collection_time,
+                last_waypoint_time=self.last_waypoint_time,
+            )
+            self.last_waypoint_time = tt
+
+    def get_action(self, time: Optional[float] = None) -> dict[str, Any]:
+        """Get the next action based on the (current) monotonic time."""
+        if time is None:
+            time = time_module.monotonic()
+        concat_vec = self.interp(time)
+        self.last_action.update(self._unconcat_vecs(concat_vec))
+        return self.last_action
+
+    def observation_space(self) -> gym.spaces.Dict:
+        """Return the observation space."""
+        pass
+
+    def action_space(self) -> gym.spaces.Dict:
+        """Return the action space."""
+        pass
+
+    def close(self) -> None:
+        """Clean up resources."""
+        pass
+
+
+class PoseTrajectoryInterpolator:
+    def __init__(self, times: np.ndarray, poses: np.ndarray):
+        assert len(times) >= 1
+        assert len(poses) == len(times)
+
+        times = np.asarray(times)
+        poses = np.asarray(poses)
+
+        self.num_joint = len(poses[0])
+
+        if len(times) == 1:
+            # special treatment for single step interpolation
+            self.single_step = True
+            self._times = times
+            self._poses = poses
+        else:
+            self.single_step = False
+            assert np.all(times[1:] >= times[:-1])
+            self.pose_interp = si.interp1d(times, poses, axis=0, assume_sorted=True)
+
+    @property
+    def times(self) -> np.ndarray:
+        if self.single_step:
+            return self._times
+        else:
+            return self.pose_interp.x
+
+    @property
+    def poses(self) -> np.ndarray:
+        if self.single_step:
+            return self._poses
+        else:
+            return self.pose_interp.y
+
+    def trim(self, start_t: float, end_t: float) -> "PoseTrajectoryInterpolator":
+        assert start_t <= end_t
+        times = self.times
+        should_keep = (start_t < times) & (times < end_t)
+        keep_times = times[should_keep]
+        all_times = np.concatenate([[start_t], keep_times, [end_t]])
+        # remove duplicates, Slerp requires strictly increasing x
+        all_times = np.unique(all_times)
+        # interpolate
+        all_poses = self(all_times)
+        return PoseTrajectoryInterpolator(times=all_times, poses=all_poses)
+
+    def schedule_waypoint(
+        self,
+        pose,
+        time,
+        max_change_rate=np.inf,
+        interpolation_garbage_collection_time=None,
+        last_waypoint_time=None,
+    ) -> "PoseTrajectoryInterpolator":
+        if not isinstance(max_change_rate, np.ndarray):
+            max_change_rate = np.array([max_change_rate] * self.num_joint)
+
+        assert len(max_change_rate) == self.num_joint
+        assert np.max(max_change_rate) > 0
+
+        if last_waypoint_time is not None:
+            assert interpolation_garbage_collection_time is not None
+
+        # trim current interpolator to between interpolation_garbage_collection_time and last_waypoint_time
+        start_time = self.times[0]
+        end_time = self.times[-1]
+        assert start_time <= end_time
+        if interpolation_garbage_collection_time is not None:
+            if time <= interpolation_garbage_collection_time:
+                # if insert time is earlier than current time
+                # no effect should be done to the interpolator
+                return self
+            # now, interpolation_garbage_collection_time < time
+            start_time = max(interpolation_garbage_collection_time, start_time)
+
+            if last_waypoint_time is not None:
+                # if last_waypoint_time is earlier than start_time
+                # use start_time
+                if time <= last_waypoint_time:
+                    end_time = interpolation_garbage_collection_time
+                else:
+                    end_time = max(last_waypoint_time, interpolation_garbage_collection_time)
+            else:
+                end_time = interpolation_garbage_collection_time
+
+        end_time = min(end_time, time)
+        start_time = min(start_time, end_time)
+        # end time should be the latest of all times except time
+        # after this we can assume order (proven by zhenjia, due to the 2 min operations)
+
+        # Constraints:
+        # start_time <= end_time <= time (proven by zhenjia)
+        # interpolation_garbage_collection_time <= start_time (proven by zhenjia)
+        # interpolation_garbage_collection_time <= time (proven by zhenjia)
+
+        # time can't change
+        # last_waypoint_time can't change
+        # interpolation_garbage_collection_time can't change
+        assert start_time <= end_time
+        assert end_time <= time
+        if last_waypoint_time is not None:
+            if time <= last_waypoint_time:
+                assert end_time == interpolation_garbage_collection_time
+            else:
+                assert end_time == max(last_waypoint_time, interpolation_garbage_collection_time)
+
+        if interpolation_garbage_collection_time is not None:
+            assert interpolation_garbage_collection_time <= start_time
+            assert interpolation_garbage_collection_time <= time
+        trimmed_interp = self.trim(start_time, end_time)
+        # after this, all waypoints in trimmed_interp is within start_time and end_time
+        # and is earlier than time
+
+        # determine speed
+        duration = time - end_time
+        end_pose = trimmed_interp(end_time)
+        pose_min_duration = np.max(np.abs(end_pose - pose) / max_change_rate)
+        duration = max(duration, pose_min_duration)
+        assert duration >= 0
+        last_waypoint_time = end_time + duration
+
+        # insert new pose
+        times = np.append(trimmed_interp.times, [last_waypoint_time], axis=0)
+        poses = np.append(trimmed_interp.poses, [pose], axis=0)
+
+        # create new interpolator
+        final_interp = PoseTrajectoryInterpolator(times, poses)
+        return final_interp
+
+    def __call__(self, t: Union[numbers.Number, np.ndarray]) -> np.ndarray:
+        is_single = False
+        if isinstance(t, numbers.Number):
+            is_single = True
+            t = np.array([t])
+
+        pose = np.zeros((len(t), self.num_joint))
+        if self.single_step:
+            pose[:] = self._poses[0]
+        else:
+            start_time = self.times[0]
+            end_time = self.times[-1]
+            t = np.clip(t, start_time, end_time)
+            pose = self.pose_interp(t)
+
+        if is_single:
+            pose = pose[0]
+        return pose

GR00T-WholeBodyControl/decoupled_wbc/control/policy/keyboard_navigation_policy.py

@@ -0,0 +1,87 @@
+from typing import Any, Dict, Optional
+
+import numpy as np
+
+from decoupled_wbc.control.base.policy import Policy
+
+
+class KeyboardNavigationPolicy(Policy):
+    def __init__(
+        self,
+        max_linear_velocity: float = 0.5,
+        max_angular_velocity: float = 0.5,
+        verbose: bool = True,
+        **kwargs,
+    ):
+        """
+        Initialize the navigation policy.
+
+        Args:
+            max_linear_velocity: Maximum linear velocity in m/s (for x and y components)
+            max_angular_velocity: Maximum angular velocity in rad/s (for yaw component)
+            **kwargs: Additional arguments passed to the base Policy class
+        """
+        super().__init__(**kwargs)
+        self.max_linear_velocity = max_linear_velocity
+        self.max_angular_velocity = max_angular_velocity
+        self.verbose = verbose
+
+        # Initialize velocity commands
+        self.lin_vel_command = np.zeros(2, dtype=np.float32)  # [vx, vy]
+        self.ang_vel_command = np.zeros(1, dtype=np.float32)  # [wz]
+
+    def get_action(self, time: Optional[float] = None) -> Dict[str, Any]:
+        """
+        Get the action to execute based on current state.
+
+        Args:
+            time: Current time (optional)
+
+        Returns:
+            Dict containing the action to execute with:
+                - navigate_cmd: np.array([vx, vy, wz]) where:
+                    - vx: linear velocity in x direction (m/s)
+                    - vy: linear velocity in y direction (m/s)
+                    - wz: angular velocity around z axis (rad/s)
+        """
+        # Combine linear and angular velocities into a single command
+        # Ensure velocities are within limits
+        vx = np.clip(self.lin_vel_command[0], -self.max_linear_velocity, self.max_linear_velocity)
+        vy = np.clip(self.lin_vel_command[1], -self.max_linear_velocity, self.max_linear_velocity)
+        wz = np.clip(self.ang_vel_command[0], -self.max_angular_velocity, self.max_angular_velocity)
+
+        navigate_cmd = np.array([vx, vy, wz], dtype=np.float32)
+
+        action = {"navigate_cmd": navigate_cmd}
+        return action
+
+    def handle_keyboard_button(self, keycode: str):
+        """
+        Handle keyboard inputs for navigation control.
+
+        Args:
+            keycode: The key that was pressed
+        """
+        if keycode == "w":
+            self.lin_vel_command[0] += 0.1  # Increase forward velocity
+        elif keycode == "s":
+            self.lin_vel_command[0] -= 0.1  # Increase backward velocity
+        elif keycode == "a":
+            self.lin_vel_command[1] += 0.1  # Increase left velocity
+        elif keycode == "d":
+            self.lin_vel_command[1] -= 0.1  # Increase right velocity
+        elif keycode == "q":
+            self.ang_vel_command[0] += 0.1  # Increase counter-clockwise rotation
+        elif keycode == "e":
+            self.ang_vel_command[0] -= 0.1  # Increase clockwise rotation
+        elif keycode == "z":
+            # Reset all velocities
+            self.lin_vel_command[:] = 0.0
+            self.ang_vel_command[:] = 0.0
+            if self.verbose:
+                print("Navigation policy: Reset all velocity commands to zero")
+
+        # Print current velocities after any keyboard input
+        if self.verbose:
+            print(f"Nav lin vel: ({self.lin_vel_command[0]:.2f}, {self.lin_vel_command[1]:.2f})")
+            print(f"Nav ang vel: {self.ang_vel_command[0]:.2f}")

GR00T-WholeBodyControl/decoupled_wbc/control/policy/lerobot_replay_policy.py

@@ -0,0 +1,111 @@
+import time
+
+import pandas as pd
+
+from decoupled_wbc.control.base.policy import Policy
+from decoupled_wbc.control.main.constants import (
+    DEFAULT_BASE_HEIGHT,
+    DEFAULT_NAV_CMD,
+    DEFAULT_WRIST_POSE,
+)
+from decoupled_wbc.control.robot_model.robot_model import RobotModel
+from decoupled_wbc.data.viz.rerun_viz import RerunViz
+
+
+class LerobotReplayPolicy(Policy):
+    """Replay policy for Lerobot dataset, so we can replay the dataset
+    and just use the action from the dataset.
+
+    Args:
+        parquet_path: Path to the parquet file containing the dataset.
+    """
+
+    is_active = True  # by default, the replay policy is active
+
+    def __init__(self, robot_model: RobotModel, parquet_path: str, use_viz: bool = False):
+        # self.dataset = LerobotDataset(dataset_path)
+        self.parquet_path = parquet_path
+        self._ctr = 0
+        # read the parquet file
+        self.df = pd.read_parquet(self.parquet_path)
+        self._max_ctr = len(self.df)
+        # get the action from the dataframe
+        self.action = self.df.iloc[self._ctr]["action"]
+        self.use_viz = use_viz
+        if self.use_viz:
+            self.viz = RerunViz(
+                image_keys=["egoview_image"],
+                tensor_keys=[
+                    "left_arm_qpos",
+                    "left_hand_qpos",
+                    "right_arm_qpos",
+                    "right_hand_qpos",
+                ],
+                window_size=5.0,
+            )
+        self.robot_model = robot_model
+        self.upper_body_joint_indices = self.robot_model.get_joint_group_indices("upper_body")
+
+    def get_action(self) -> dict[str, any]:
+        # get the action from the dataframe
+        action = self.df.iloc[self._ctr]["action"]
+        wrist_pose = self.df.iloc[self._ctr]["action.eef"]
+        navigate_cmd = self.df.iloc[self._ctr].get("teleop.navigate_command", DEFAULT_NAV_CMD)
+        base_height_cmd = self.df.iloc[self._ctr].get(
+            "teleop.base_height_command", DEFAULT_BASE_HEIGHT
+        )
+
+        self._ctr += 1
+        if self._ctr >= self._max_ctr:
+            self._ctr = 0
+        # print(f"Replay {self._ctr} / {self._max_ctr}")
+        if self.use_viz:
+            self.viz.plot_tensors(
+                {
+                    "left_arm_qpos": action[self.robot_model.get_joint_group_indices("left_arm")]
+                    + 15,
+                    "left_hand_qpos": action[self.robot_model.get_joint_group_indices("left_hand")]
+                    + 15,
+                    "right_arm_qpos": action[self.robot_model.get_joint_group_indices("right_arm")]
+                    + 15,
+                    "right_hand_qpos": action[
+                        self.robot_model.get_joint_group_indices("right_hand")
+                    ]
+                    + 15,
+                },
+                time.monotonic(),
+            )
+
+        return {
+            "target_upper_body_pose": action[self.upper_body_joint_indices],
+            "wrist_pose": wrist_pose,
+            "navigate_cmd": navigate_cmd,
+            "base_height_cmd": base_height_cmd,
+            "timestamp": time.time(),
+        }
+
+    def action_to_cmd(self, action: dict[str, any]) -> dict[str, any]:
+        action["target_upper_body_pose"] = action["q"][
+            self.robot_model.get_joint_group_indices("upper_body")
+        ]
+        del action["q"]
+        return action
+
+    def set_observation(self, observation: dict[str, any]):
+        pass
+
+    def get_observation(self) -> dict[str, any]:
+        return {
+            "wrist_pose": self.df.iloc[self._ctr - 1].get(
+                "observation.eef_state", DEFAULT_WRIST_POSE
+            ),
+            "timestamp": time.time(),
+        }
+
+
+if __name__ == "__main__":
+    policy = LerobotReplayPolicy(
+        parquet_path="outputs/g1-open-hands-may7/data/chunk-000/episode_000000.parquet"
+    )
+    action = policy.get_action()
+    print(action)

GR00T-WholeBodyControl/decoupled_wbc/control/policy/teleop_policy.py

@@ -0,0 +1,207 @@
+from contextlib import contextmanager
+import time
+from typing import Optional
+
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+
+from decoupled_wbc.control.base.policy import Policy
+from decoupled_wbc.control.robot_model import RobotModel
+from decoupled_wbc.control.teleop.teleop_retargeting_ik import TeleopRetargetingIK
+from decoupled_wbc.control.teleop.teleop_streamer import TeleopStreamer
+
+
+class TeleopPolicy(Policy):
+    """
+    Robot-agnostic teleop policy.
+    Clean separation: IK processing vs command passing.
+    All robot-specific properties are abstracted through robot_model and hand_ik_solvers.
+    """
+
+    def __init__(
+        self,
+        body_control_device: str,
+        hand_control_device: str,
+        robot_model: RobotModel,
+        retargeting_ik: TeleopRetargetingIK,
+        body_streamer_ip: str = "192.168.?.?",
+        body_streamer_keyword: str = "shoulder",
+        enable_real_device: bool = True,
+        replay_data_path: Optional[str] = None,
+        replay_speed: float = 1.0,
+        wait_for_activation: int = 5,
+        activate_keyboard_listener: bool = True,
+    ):
+        if activate_keyboard_listener:
+            from decoupled_wbc.control.utils.keyboard_dispatcher import KeyboardListenerSubscriber
+
+            self.keyboard_listener = KeyboardListenerSubscriber()
+        else:
+            self.keyboard_listener = None
+
+        self.wait_for_activation = wait_for_activation
+
+        self.teleop_streamer = TeleopStreamer(
+            robot_model=robot_model,
+            body_control_device=body_control_device,
+            hand_control_device=hand_control_device,
+            enable_real_device=enable_real_device,
+            body_streamer_ip=body_streamer_ip,
+            body_streamer_keyword=body_streamer_keyword,
+            replay_data_path=replay_data_path,
+            replay_speed=replay_speed,
+        )
+        self.robot_model = robot_model
+        self.retargeting_ik = retargeting_ik
+        self.is_active = False
+
+        self.latest_left_wrist_data = np.eye(4)
+        self.latest_right_wrist_data = np.eye(4)
+        self.latest_left_fingers_data = {"position": np.zeros((25, 4, 4))}
+        self.latest_right_fingers_data = {"position": np.zeros((25, 4, 4))}
+
+    def set_goal(self, goal: dict[str, any]):
+        # The current teleop policy doesn't take higher level commands yet.
+        pass
+
+    def get_action(self) -> dict[str, any]:
+        # Get structured data
+        streamer_output = self.teleop_streamer.get_streamer_data()
+
+        # Handle activation using teleop_data commands
+        self.check_activation(
+            streamer_output.teleop_data, wait_for_activation=self.wait_for_activation
+        )
+
+        action = {}
+
+        # Process streamer data if active
+        if self.is_active and streamer_output.ik_data:
+            body_data = streamer_output.ik_data["body_data"]
+            left_hand_data = streamer_output.ik_data["left_hand_data"]
+            right_hand_data = streamer_output.ik_data["right_hand_data"]
+
+            left_wrist_name = self.robot_model.supplemental_info.hand_frame_names["left"]
+            right_wrist_name = self.robot_model.supplemental_info.hand_frame_names["right"]
+            self.latest_left_wrist_data = body_data[left_wrist_name]
+            self.latest_right_wrist_data = body_data[right_wrist_name]
+            self.latest_left_fingers_data = left_hand_data
+            self.latest_right_fingers_data = right_hand_data
+
+            # TODO: This stores the same data again
+            ik_data = {
+                "body_data": body_data,
+                "left_hand_data": left_hand_data,
+                "right_hand_data": right_hand_data,
+            }
+            action["ik_data"] = ik_data
+
+        # Wrist poses (pos and quat)
+        # TODO: This stores the same wrist poses in two different formats
+        left_wrist_matrix = self.latest_left_wrist_data
+        right_wrist_matrix = self.latest_right_wrist_data
+        left_wrist_pose = np.concatenate(
+            [
+                left_wrist_matrix[:3, 3],
+                R.from_matrix(left_wrist_matrix[:3, :3]).as_quat(scalar_first=True),
+            ]
+        )
+        right_wrist_pose = np.concatenate(
+            [
+                right_wrist_matrix[:3, 3],
+                R.from_matrix(right_wrist_matrix[:3, :3]).as_quat(scalar_first=True),
+            ]
+        )
+
+        # Combine IK results with control commands (no teleop_data commands)
+        action.update(
+            {
+                "left_wrist": self.latest_left_wrist_data,
+                "right_wrist": self.latest_right_wrist_data,
+                "left_fingers": self.latest_left_fingers_data,
+                "right_fingers": self.latest_right_fingers_data,
+                "wrist_pose": np.concatenate([left_wrist_pose, right_wrist_pose]),
+                **streamer_output.control_data,  # Only control & data collection commands pass through
+                **streamer_output.data_collection_data,
+            }
+        )
+
+        # Run retargeting IK
+        if "ik_data" in action:
+            self.retargeting_ik.set_goal(action["ik_data"])
+        action["target_upper_body_pose"] = self.retargeting_ik.get_action()
+
+        return action
+
+    def close(self) -> bool:
+        self.teleop_streamer.stop_streaming()
+        return True
+
+    def check_activation(self, teleop_data: dict, wait_for_activation: int = 5):
+        """Activation logic only looks at teleop data commands"""
+        key = self.keyboard_listener.read_msg() if self.keyboard_listener else ""
+        toggle_activation_by_keyboard = key == "l"
+        reset_teleop_policy_by_keyboard = key == "k"
+        toggle_activation_by_teleop = teleop_data.get("toggle_activation", False)
+
+        if reset_teleop_policy_by_keyboard:
+            print("Resetting teleop policy")
+            self.reset()
+
+        if toggle_activation_by_keyboard or toggle_activation_by_teleop:
+            self.is_active = not self.is_active
+            if self.is_active:
+                print("Starting teleop policy")
+
+                if wait_for_activation > 0 and toggle_activation_by_keyboard:
+                    print(f"Sleeping for {wait_for_activation} seconds before starting teleop...")
+                    for i in range(wait_for_activation, 0, -1):
+                        print(f"Starting in {i}...")
+                        time.sleep(1)
+
+                # dda: calibration logic should use current IK data
+                self.teleop_streamer.calibrate()
+                print("Teleop policy calibrated")
+            else:
+                print("Stopping teleop policy")
+
+    @contextmanager
+    def activate(self):
+        try:
+            yield self
+        finally:
+            self.close()
+
+    def handle_keyboard_button(self, keycode):
+        """
+        Handle keyboard input with proper state toggle.
+        """
+        if keycode == "l":
+            # Toggle start state
+            self.is_active = not self.is_active
+            # Reset initialization when stopping
+            if not self.is_active:
+                self._initialized = False
+        if keycode == "k":
+            print("Resetting teleop policy")
+            self.reset()
+
+    def activate_policy(self, wait_for_activation: int = 5):
+        """activate the teleop policy"""
+        self.is_active = False
+        self.check_activation(
+            teleop_data={"toggle_activation": True}, wait_for_activation=wait_for_activation
+        )
+
+    def reset(self, wait_for_activation: int = 5, auto_activate: bool = False):
+        """Reset the teleop policy to the initial state, and re-activate it."""
+        self.teleop_streamer.reset()
+        self.retargeting_ik.reset()
+        self.is_active = False
+        self.latest_left_wrist_data = np.eye(4)
+        self.latest_right_wrist_data = np.eye(4)
+        self.latest_left_fingers_data = {"position": np.zeros((25, 4, 4))}
+        self.latest_right_fingers_data = {"position": np.zeros((25, 4, 4))}
+
+        if auto_activate:
+            self.activate_policy(wait_for_activation)

GR00T-WholeBodyControl/decoupled_wbc/control/policy/wbc_policy_factory.py

@@ -0,0 +1,65 @@
+import os
+from pathlib import Path
+import time
+
+import numpy as np
+
+import decoupled_wbc
+from decoupled_wbc.control.main.constants import DEFAULT_BASE_HEIGHT, DEFAULT_NAV_CMD
+from decoupled_wbc.control.policy.g1_gear_wbc_policy import G1GearWbcPolicy
+from decoupled_wbc.control.policy.identity_policy import IdentityPolicy
+from decoupled_wbc.control.policy.interpolation_policy import InterpolationPolicy
+
+from .g1_decoupled_whole_body_policy import G1DecoupledWholeBodyPolicy
+
+WBC_VERSIONS = ["gear_wbc"]
+
+
+def get_wbc_policy(
+    robot_type,
+    robot_model,
+    wbc_config,
+    init_time=time.monotonic(),
+):
+    current_upper_body_pose = robot_model.get_initial_upper_body_pose()
+
+    if robot_type == "g1":
+        upper_body_policy_type = wbc_config.get("upper_body_policy_type", "interpolation")
+        if upper_body_policy_type == "identity":
+            upper_body_policy = IdentityPolicy()
+        else:
+            upper_body_policy = InterpolationPolicy(
+                init_time=init_time,
+                init_values={
+                    "target_upper_body_pose": current_upper_body_pose,
+                    "base_height_command": np.array([DEFAULT_BASE_HEIGHT]),
+                    "navigate_cmd": np.array([DEFAULT_NAV_CMD]),
+                },
+                max_change_rate=wbc_config["upper_body_max_joint_speed"],
+            )
+
+        lower_body_policy_type = wbc_config.get("VERSION", "gear_wbc")
+        if lower_body_policy_type not in ["gear_wbc"]:
+            raise ValueError(
+                f"Invalid lower body policy version: {lower_body_policy_type}. "
+                f"Only 'gear_wbc' is supported."
+            )
+
+        # Get the base path to decoupled_wbc and convert to Path object
+        package_path = Path(os.path.dirname(decoupled_wbc.__file__))
+        gear_wbc_config = str(package_path / ".." / wbc_config["GEAR_WBC_CONFIG"])
+        if lower_body_policy_type == "gear_wbc":
+            lower_body_policy = G1GearWbcPolicy(
+                robot_model=robot_model,
+                config=gear_wbc_config,
+                model_path=wbc_config["model_path"],
+            )
+
+        wbc_policy = G1DecoupledWholeBodyPolicy(
+            robot_model=robot_model,
+            upper_body_policy=upper_body_policy,
+            lower_body_policy=lower_body_policy,
+        )
+    else:
+        raise ValueError(f"Invalid robot type: {robot_type}")
+    return wbc_policy

GR00T-WholeBodyControl/decoupled_wbc/control/robot_model/__init__.py

@@ -0,0 +1,3 @@
+from .robot_model import ReducedRobotModel, RobotModel
+
+__all__ = ["RobotModel", "ReducedRobotModel"]

GR00T-WholeBodyControl/decoupled_wbc/control/robot_model/robot_model.py

@@ -0,0 +1,772 @@
+from typing import List, Optional, Set, Union
+
+import numpy as np
+import pinocchio as pin
+
+from decoupled_wbc.control.robot_model.supplemental_info import RobotSupplementalInfo
+
+
+class RobotModel:
+    def __init__(
+        self,
+        urdf_path,
+        asset_path,
+        set_floating_base=False,
+        supplemental_info: Optional[RobotSupplementalInfo] = None,
+    ):
+        self.pinocchio_wrapper = pin.RobotWrapper.BuildFromURDF(
+            filename=urdf_path,
+            package_dirs=[asset_path],
+            root_joint=pin.JointModelFreeFlyer() if set_floating_base else None,
+        )
+        self.is_floating_base_model = set_floating_base
+
+        self.joint_to_dof_index = {}
+        # Assume we only have single-dof joints
+        # First two names correspond to universe and floating base joints
+        names = (
+            self.pinocchio_wrapper.model.names[2:]
+            if set_floating_base
+            else self.pinocchio_wrapper.model.names[1:]
+        )
+        for name in names:
+            j_id = self.pinocchio_wrapper.model.getJointId(name)
+            jmodel = self.pinocchio_wrapper.model.joints[j_id]
+            self.joint_to_dof_index[name] = jmodel.idx_q
+
+        # Store joint limits only for actual joints (excluding floating base)
+        # if set floating base is true and the robot can move in the world
+        # then we don't want to impose joint limits for the 7 dofs corresponding
+        # to the floating base dofs.
+        root_nq = 7 if set_floating_base else 0
+        self.upper_joint_limits = self.pinocchio_wrapper.model.upperPositionLimit[root_nq:].copy()
+        self.lower_joint_limits = self.pinocchio_wrapper.model.lowerPositionLimit[root_nq:].copy()
+
+        # Set up supplemental info if provided
+        self.supplemental_info = supplemental_info
+        if self.supplemental_info is not None:
+            # Cache indices for body and hand actuated joints separately
+            self._body_actuated_joint_indices = [
+                self.dof_index(name) for name in self.supplemental_info.body_actuated_joints
+            ]
+            self._left_hand_actuated_joint_indices = [
+                self.dof_index(name) for name in self.supplemental_info.left_hand_actuated_joints
+            ]
+            self._right_hand_actuated_joint_indices = [
+                self.dof_index(name) for name in self.supplemental_info.right_hand_actuated_joints
+            ]
+            self._hand_actuated_joint_indices = (
+                self._left_hand_actuated_joint_indices + self._right_hand_actuated_joint_indices
+            )
+
+            # Cache indices for joint groups, handling nested groups
+            self._joint_group_indices = {}
+            for group_name, group_info in self.supplemental_info.joint_groups.items():
+                indices = []
+                # Add indices for direct joints
+                indices.extend([self.dof_index(name) for name in group_info["joints"]])
+                # Add indices from subgroups
+                for subgroup_name in group_info["groups"]:
+                    indices.extend(self.get_joint_group_indices(subgroup_name))
+                self._joint_group_indices[group_name] = sorted(set(indices))
+
+            # Update joint limits from supplemental info if available
+            if (
+                hasattr(self.supplemental_info, "joint_limits")
+                and self.supplemental_info.joint_limits
+            ):
+                for joint_name, limits in self.supplemental_info.joint_limits.items():
+                    if joint_name in self.joint_to_dof_index:
+                        idx = self.joint_to_dof_index[joint_name] - root_nq
+                        self.lower_joint_limits[idx] = limits[0]
+                        self.upper_joint_limits[idx] = limits[1]
+
+        # Initialize default body pose
+        self.default_body_pose = self.q_zero.copy()
+
+        # Update with supplemental info if available
+        if self.supplemental_info is not None:
+            default_joint_q = self.supplemental_info.default_joint_q
+            for joint, joint_values in default_joint_q.items():
+                # Get the joint name mapping for this type
+                joint_mapping = self.supplemental_info.joint_name_mapping[joint]
+
+                # Handle both single joint names and left/right mappings
+                if isinstance(joint_mapping, str):
+                    # Single joint (e.g., waist joints)
+                    if joint_mapping in self.joint_to_dof_index:
+                        joint_idx = self.dof_index(joint_mapping)
+                        self.default_body_pose[joint_idx] = (
+                            joint_values  # joint_values is the value for single joints
+                        )
+                else:
+                    # Left/right mapping (e.g., arm joints)
+                    for side, value in joint_values.items():
+                        if side in joint_mapping and joint_mapping[side] in self.joint_to_dof_index:
+                            joint_idx = self.dof_index(joint_mapping[side])
+                            self.default_body_pose[joint_idx] = value
+
+        # Initialize initial body pose
+        self.initial_body_pose = self.default_body_pose.copy()
+
+    @property
+    def num_dofs(self) -> int:
+        """Get the number of degrees of freedom of the robot (floating base pose + joints)."""
+        return self.pinocchio_wrapper.model.nq
+
+    @property
+    def q_zero(self) -> np.ndarray:
+        """Get the zero pose of the robot."""
+        return self.pinocchio_wrapper.q0
+
+    @property
+    def joint_names(self) -> List[str]:
+        """Get the names of the joints of the robot."""
+        return list(self.joint_to_dof_index.keys())
+
+    @property
+    def num_joints(self) -> int:
+        """Get the number of joints of the robot."""
+        return len(self.joint_to_dof_index)
+
+    def dof_index(self, joint_name: str) -> int:
+        """
+        Get the index in the degrees of freedom vector corresponding
+        to the single-DoF joint with name `joint_name`.
+        """
+        if joint_name not in self.joint_to_dof_index:
+            raise ValueError(
+                f"Unknown joint name: '{joint_name}'. "
+                f"Available joints: {list(self.joint_to_dof_index.keys())}"
+            )
+        return self.joint_to_dof_index[joint_name]
+
+    def get_body_actuated_joint_indices(self) -> List[int]:
+        """
+        Get the indices of body actuated joints in the full configuration.
+        Ordering is that of the actuated joints as defined in the supplemental info.
+        Requires supplemental_info to be provided.
+        """
+        if self.supplemental_info is None:
+            raise ValueError("supplemental_info must be provided to use this method")
+        return self._body_actuated_joint_indices
+
+    def get_hand_actuated_joint_indices(self, side: str = "both") -> List[int]:
+        """
+        Get the indices of hand actuated joints in the full configuration.
+        Ordering is that of the actuated joints as defined in the supplemental info.
+        Requires supplemental_info to be provided.
+
+        Args:
+            side: String specifying which hand to get indices for ('left', 'right', or 'both')
+        """
+        if self.supplemental_info is None:
+            raise ValueError("supplemental_info must be provided to use this method")
+
+        if side.lower() == "both":
+            return self._hand_actuated_joint_indices
+        elif side.lower() == "left":
+            return self._left_hand_actuated_joint_indices
+        elif side.lower() == "right":
+            return self._right_hand_actuated_joint_indices
+        else:
+            raise ValueError("side must be 'left', 'right', or 'both'")
+
+    def get_joint_group_indices(self, group_names: Union[str, Set[str]]) -> List[int]:
+        """
+        Get the indices of joints in one or more groups in the full configuration.
+        Requires supplemental_info to be provided.
+        The returned indices are sorted in ascending order, so that the joint ordering
+        of the full model is preserved.
+
+        Args:
+            group_names: Either a single group name (str) or a set of group names (Set[str])
+
+        Returns:
+            List of joint indices in sorted order with no duplicates
+        """
+        if self.supplemental_info is None:
+            raise ValueError("supplemental_info must be provided to use this method")
+
+        # Convert single string to set for uniform handling
+        if isinstance(group_names, str):
+            group_names = {group_names}
+
+        # Collect indices from all groups
+        all_indices = set()
+        for group_name in group_names:
+            if group_name not in self._joint_group_indices:
+                raise ValueError(f"Unknown joint group: {group_name}")
+            all_indices.update(self._joint_group_indices[group_name])
+
+        return sorted(all_indices)
+
+    def cache_forward_kinematics(self, q: np.ndarray, auto_clip=True) -> None:
+        """
+        Perform forward kinematics to update the pose of every joint and frame
+        in the Pinocchio data structures for the given configuration `q`.
+
+        :param q: A numpy array of shape (num_dofs,) representing the robot configuration.
+        """
+        if q.shape[0] != self.num_dofs:
+            raise ValueError(f"Expected q of length {self.num_dofs}, got {q.shape[0]} instead.")
+
+        # Apply auto-clip if enabled
+        if auto_clip:
+            q = self.clip_configuration(q)
+
+        pin.framesForwardKinematics(self.pinocchio_wrapper.model, self.pinocchio_wrapper.data, q)
+
+    def compute_gravity_compensation_torques(
+        self, q: np.ndarray, joint_groups: Union[str, List[str], Set[str]] = None, auto_clip=True
+    ) -> np.ndarray:
+        """
+        Compute gravity compensation torques for specified joint groups using pinocchio.
+
+        :param q: Robot configuration (joint positions)
+        :param joint_groups: Joint groups to compensate (e.g., "arms", ["left_arm", "waist"],
+             {"left_arm", "waist"}). If None, compensates all joints
+        :param auto_clip: Whether to automatically clip joint values to limits
+        :return: Array of gravity compensation torques for all DOFs (zero for non-compensated joints)
+        """
+        if q.shape[0] != self.num_dofs:
+            raise ValueError(f"Expected q of length {self.num_dofs}, got {q.shape[0]} instead.")
+
+        # Apply auto-clip if enabled
+        if auto_clip:
+            q = self.clip_configuration(q)
+
+        try:
+            # Cache forward kinematics for the current configuration
+            self.cache_forward_kinematics(q, auto_clip=False)  # Already clipped if needed
+
+            # Compute gravity vector using RNEA with zero velocity and acceleration
+            v = np.zeros(self.num_dofs)
+            a = np.zeros(self.num_dofs)
+
+            gravity_torques_full = pin.rnea(
+                self.pinocchio_wrapper.model, self.pinocchio_wrapper.data, q, v, a
+            )
+
+            # If no joint groups specified, return full gravity torques
+            if joint_groups is None:
+                return gravity_torques_full
+
+            # Convert list to set for get_joint_group_indices compatibility
+            if isinstance(joint_groups, list):
+                joint_groups = set(joint_groups)
+
+            # Get joint indices for specified groups - get_joint_group_indices handles str and Set[str]
+            try:
+                compensated_joint_indices = self.get_joint_group_indices(joint_groups)
+            except ValueError as e:
+                raise ValueError(f"Error resolving joint groups {joint_groups}: {e}")
+
+            # Create mask for joints that should receive gravity compensation
+            compensation_mask = np.zeros(self.num_dofs, dtype=bool)
+            for joint_idx in compensated_joint_indices:
+                if 0 <= joint_idx < len(compensation_mask):
+                    compensation_mask[joint_idx] = True
+
+            # Apply mask to only compensate specified joints
+            compensated_torques = np.zeros_like(gravity_torques_full)
+            compensated_torques[compensation_mask] = gravity_torques_full[compensation_mask]
+
+            return compensated_torques
+
+        except Exception as e:
+            raise RuntimeError(f"Error computing gravity compensation: {e}")
+
+    def clip_configuration(self, q: np.ndarray, margin: float = 1e-6) -> np.ndarray:
+        """
+        Clip the configuration to stay within joint limits with a small tolerance.
+
+        :param q: Configuration to clip
+        :param margin: Tolerance to keep away from joint limits
+        :return: Clipped configuration
+        """
+        q_clipped = q.copy()
+
+        # Only clip joint positions, not floating base
+        root_nq = 7 if self.is_floating_base_model else 0
+        q_clipped[root_nq:] = np.clip(
+            q[root_nq:], self.lower_joint_limits + margin, self.upper_joint_limits - margin
+        )
+
+        return q_clipped
+
+    def frame_placement(self, frame_name: str) -> pin.SE3:
+        """
+        Returns the SE3 transform of the specified frame in the world coordinate system.
+        Note: make sure cache_forward_kinematics() has been previously called.
+
+        :param frame_name: Name of the frame, e.g. "link_elbow_frame", "hand_imu_frame", etc.
+        :return: A pin.SE3 object representing the pose of the frame.
+        """
+        model = self.pinocchio_wrapper.model
+        data = self.pinocchio_wrapper.data
+
+        frame_id = model.getFrameId(frame_name)
+        if frame_id < 0 or frame_id >= len(model.frames):
+            valid_frames = [f.name for f in model.frames]
+            raise ValueError(f"Unknown frame '{frame_name}'. Valid frames: {valid_frames}")
+
+        # Pinocchio's data.oMf[frame_id] is a pin.SE3.
+        return data.oMf[frame_id].copy()
+
+    def get_body_actuated_joints(self, q: np.ndarray) -> np.ndarray:
+        """
+        Get the configuration of body actuated joints from a full configuration.
+
+        :param q: Configuration in full space
+        :return: Configuration of body actuated joints
+        """
+        indices = self.get_body_actuated_joint_indices()
+
+        return q[indices]
+
+    def get_hand_actuated_joints(self, q: np.ndarray, side: str = "both") -> np.ndarray:
+        """
+        Get the configuration of hand actuated joints from a full configuration.
+
+        Args:
+            q: Configuration in full space
+            side: String specifying which hand to get joints for ('left', 'right', or 'both')
+        """
+        indices = self.get_hand_actuated_joint_indices(side)
+        return q[indices]
+
+    def get_configuration_from_actuated_joints(
+        self,
+        body_actuated_joint_values: np.ndarray,
+        hand_actuated_joint_values: Optional[np.ndarray] = None,
+        left_hand_actuated_joint_values: Optional[np.ndarray] = None,
+        right_hand_actuated_joint_values: Optional[np.ndarray] = None,
+    ) -> np.ndarray:
+        """
+        Get the full configuration from the body and hand actuated joint configurations.
+        Can specify either both hands together or left and right hands separately.
+
+        Args:
+            body_actuated_joint_values: Configuration of body actuated joints
+            hand_actuated_joint_values: Configuration of both hands' actuated joints (optional)
+            left_hand_actuated_joint_values: Configuration of left hand actuated joints (optional)
+            right_hand_actuated_joint_values: Configuration of right hand actuated joints (optional)
+
+        Returns:
+            Full configuration including body and hand joints
+        """
+        q = self.pinocchio_wrapper.q0.copy()
+        q[self.get_body_actuated_joint_indices()] = body_actuated_joint_values
+
+        # Handle hand configurations
+        if hand_actuated_joint_values is not None:
+            # Use combined hand configuration
+            q[self.get_hand_actuated_joint_indices("both")] = hand_actuated_joint_values
+        else:
+            # Use separate hand configurations
+            if left_hand_actuated_joint_values is not None:
+                q[self.get_hand_actuated_joint_indices("left")] = left_hand_actuated_joint_values
+            if right_hand_actuated_joint_values is not None:
+                q[self.get_hand_actuated_joint_indices("right")] = right_hand_actuated_joint_values
+
+        return q
+
+    def reset_forward_kinematics(self) -> None:
+        """
+        Reset the forward kinematics to the initial configuration.
+        """
+        self.cache_forward_kinematics(self.q_zero)
+
+    def get_initial_upper_body_pose(self) -> np.ndarray:
+        """
+        Get the initial upper body pose of the robot.
+        """
+        return self.initial_body_pose[self.get_joint_group_indices("upper_body")]
+
+    def get_default_body_pose(self) -> np.ndarray:
+        """
+        Get the default body pose of the robot.
+        """
+        return self.default_body_pose
+
+    def set_initial_body_pose(self, q: np.ndarray, q_idx=None) -> None:
+        """
+        Set the initial body pose of the robot.
+        """
+        if q_idx is None:
+            self.initial_body_pose = q
+        else:
+            self.initial_body_pose[q_idx] = q
+
+
+class ReducedRobotModel(RobotModel):
+    """
+    A class that creates a reduced order robot model by fixing certain joints.
+    This class maintains a mapping between the reduced state space and the full state space.
+    """
+
+    def __init__(
+        self,
+        full_robot_model: RobotModel,
+        fixed_joints: List[str],
+        fixed_values: Optional[List[float]] = None,
+    ):
+        """
+        Create a reduced order robot model by fixing specified joints.
+
+        :param full_robot_model: The original robot model
+        :param fixed_joints: List of joint names to fix
+        :param fixed_values: Optional list of values to fix the joints to. If None, uses the initial
+                            joint positions (q0) from the full robot model.
+        """
+        self.full_robot = full_robot_model
+        self.supplemental_info = full_robot_model.supplemental_info
+
+        # If fixed_values is None, use q0 from the full robot model
+        if fixed_values is None:
+            fixed_values = []
+            for joint_name in fixed_joints:
+                full_idx = full_robot_model.dof_index(joint_name)
+                fixed_values.append(full_robot_model.pinocchio_wrapper.q0[full_idx])
+        elif len(fixed_joints) != len(fixed_values):
+            raise ValueError("fixed_joints and fixed_values must have the same length")
+
+        # Store fixed joints and their values
+        self.fixed_joints = fixed_joints
+        self.fixed_values = fixed_values
+
+        # Create mapping between reduced and full state spaces
+        self.reduced_to_full = []
+        self.full_to_reduced = {}
+
+        # Initialize with floating base indices if present
+        if full_robot_model.is_floating_base_model:
+            self.reduced_to_full.extend(range(7))  # Floating base indices
+            for i in range(7):
+                self.full_to_reduced[i] = i
+
+        # Add active joint indices
+        for joint_name in full_robot_model.joint_names:
+            if joint_name not in fixed_joints:
+                full_idx = full_robot_model.dof_index(joint_name)
+                reduced_idx = len(self.reduced_to_full)
+                self.reduced_to_full.append(full_idx)
+                self.full_to_reduced[full_idx] = reduced_idx
+
+        # Create a reduced Pinocchio model using buildReducedModel
+        # First, get the list of joint IDs to lock
+        locked_joint_ids = []
+        for joint_name in fixed_joints:
+            joint_id = full_robot_model.pinocchio_wrapper.model.getJointId(joint_name)
+            if (full_robot_model.is_floating_base_model and joint_id > 1) or (
+                not full_robot_model.is_floating_base_model and joint_id > 0
+            ):
+                locked_joint_ids.append(joint_id)
+
+        # First build the reduced kinematic model
+        reduced_model = pin.buildReducedModel(
+            full_robot_model.pinocchio_wrapper.model,
+            locked_joint_ids,
+            full_robot_model.pinocchio_wrapper.q0,
+        )
+
+        # Then build the reduced geometry models using the reduced kinematic model
+        self.pinocchio_wrapper = pin.RobotWrapper(
+            model=reduced_model,
+        )
+
+        # Create joint to dof index mapping
+        self.joint_to_dof_index = {}
+        # Assume we only have single-dof joints
+        # First two names correspond to universe and floating base joints
+        names = (
+            self.pinocchio_wrapper.model.names[2:]
+            if self.full_robot.is_floating_base_model
+            else self.pinocchio_wrapper.model.names[1:]
+        )
+        for name in names:
+            j_id = self.pinocchio_wrapper.model.getJointId(name)
+            jmodel = self.pinocchio_wrapper.model.joints[j_id]
+            self.joint_to_dof_index[name] = jmodel.idx_q
+
+        # Initialize joint limits
+        root_nq = 7 if self.full_robot.is_floating_base_model else 0
+        self.lower_joint_limits = self.pinocchio_wrapper.model.lowerPositionLimit[root_nq:].copy()
+        self.upper_joint_limits = self.pinocchio_wrapper.model.upperPositionLimit[root_nq:].copy()
+
+        # Update joint limits from supplemental info if available
+        if self.supplemental_info is not None:
+            if (
+                hasattr(self.supplemental_info, "joint_limits")
+                and self.supplemental_info.joint_limits
+            ):
+                for joint_name, limits in self.supplemental_info.joint_limits.items():
+                    if joint_name in self.joint_to_dof_index:
+                        idx = self.joint_to_dof_index[joint_name] - root_nq
+                        self.lower_joint_limits[idx] = limits[0]
+                        self.upper_joint_limits[idx] = limits[1]
+
+            # Get full indices for body and hand actuated joints
+            full_body_indices = full_robot_model.get_body_actuated_joint_indices()
+            full_hand_indices = full_robot_model.get_hand_actuated_joint_indices("both")
+            full_left_hand_indices = full_robot_model.get_hand_actuated_joint_indices("left")
+            full_right_hand_indices = full_robot_model.get_hand_actuated_joint_indices("right")
+
+            # Map to reduced indices
+            self._body_actuated_joint_indices = []
+            for idx in full_body_indices:
+                if idx in self.full_to_reduced:
+                    self._body_actuated_joint_indices.append(self.full_to_reduced[idx])
+
+            self._hand_actuated_joint_indices = []
+            for idx in full_hand_indices:
+                if idx in self.full_to_reduced:
+                    self._hand_actuated_joint_indices.append(self.full_to_reduced[idx])
+
+            self._left_hand_actuated_joint_indices = []
+            for idx in full_left_hand_indices:
+                if idx in self.full_to_reduced:
+                    self._left_hand_actuated_joint_indices.append(self.full_to_reduced[idx])
+
+            self._right_hand_actuated_joint_indices = []
+            for idx in full_right_hand_indices:
+                if idx in self.full_to_reduced:
+                    self._right_hand_actuated_joint_indices.append(self.full_to_reduced[idx])
+
+            # Cache indices for joint groups in reduced space
+            self._joint_group_indices = {}
+            for group_name in self.supplemental_info.joint_groups:
+                full_indices = full_robot_model.get_joint_group_indices(group_name)
+                reduced_indices = []
+                for idx in full_indices:
+                    if idx in self.full_to_reduced:
+                        reduced_indices.append(self.full_to_reduced[idx])
+                self._joint_group_indices[group_name] = sorted(set(reduced_indices))
+
+        # Initialize default body pose in reduced space
+        self.default_body_pose = self.full_to_reduced_configuration(
+            full_robot_model.default_body_pose
+        )
+
+        # Initialize initial body pose in reduced space
+        self.initial_body_pose = self.full_to_reduced_configuration(
+            full_robot_model.initial_body_pose
+        )
+
+    @property
+    def num_joints(self) -> int:
+        """Get the number of active joints in the reduced model."""
+        return len(self.joint_names)
+
+    @property
+    def joint_names(self) -> List[str]:
+        """Get the names of the active joints in the reduced model."""
+        return [name for name in self.full_robot.joint_names if name not in self.fixed_joints]
+
+    @classmethod
+    def from_fixed_groups(
+        cls,
+        full_robot_model: RobotModel,
+        fixed_group_names: List[str],
+        fixed_values: Optional[List[float]] = None,
+    ) -> "ReducedRobotModel":
+        """
+        Create a reduced order robot model by fixing all joints in specified groups.
+
+        :param full_robot_model: The original robot model
+        :param fixed_group_names: List of joint group names to fix
+        :param fixed_values: Optional list of values to fix the joints to. If None, uses the initial
+                            joint positions (q0) from the full robot model.
+        :return: A ReducedRobotModel instance
+        """
+        if full_robot_model.supplemental_info is None:
+            raise ValueError("supplemental_info must be provided to use this method")
+
+        # Get all joints in the groups, including those from subgroups
+        fixed_joints = set()  # Use a set to avoid duplicates
+
+        for group_name in fixed_group_names:
+            if group_name not in full_robot_model.supplemental_info.joint_groups:
+                raise ValueError(f"Unknown joint group: {group_name}")
+
+            group_info = full_robot_model.supplemental_info.joint_groups[group_name]
+
+            # Add direct joints
+            fixed_joints.update(group_info["joints"])
+
+            # Add joints from subgroups
+            for subgroup_name in group_info["groups"]:
+                subgroup_joints = full_robot_model.get_joint_group_indices(subgroup_name)
+                fixed_joints.update([full_robot_model.joint_names[idx] for idx in subgroup_joints])
+
+        # Convert set back to list for compatibility with the original constructor
+        return cls(full_robot_model, list(fixed_joints), fixed_values)
+
+    @classmethod
+    def from_fixed_group(
+        cls,
+        full_robot_model: RobotModel,
+        fixed_group_name: str,
+        fixed_values: Optional[List[float]] = None,
+    ) -> "ReducedRobotModel":
+        """
+        Create a reduced order robot model by fixing all joints in a specified group.
+        This is a convenience method that calls from_fixed_groups with a single group.
+
+        :param full_robot_model: The original robot model
+        :param fixed_group_name: Name of the joint group to fix
+        :param fixed_values: Optional list of values to fix the joints to. If None, uses the initial
+                            joint positions (q0) from the full robot model.
+        :return: A ReducedRobotModel instance
+        """
+        return cls.from_fixed_groups(full_robot_model, [fixed_group_name], fixed_values)
+
+    @classmethod
+    def from_active_group(
+        cls,
+        full_robot_model: RobotModel,
+        active_group_name: str,
+        fixed_values: Optional[List[float]] = None,
+    ) -> "ReducedRobotModel":
+        """
+        Create a reduced order robot model by fixing all joints EXCEPT those in the specified group.
+        This is a convenience method that calls from_active_groups with a single group.
+
+        :param full_robot_model: The original robot model
+        :param active_group_name: Name of the joint group to keep active (all other joints will be fixed)
+        :param fixed_values: Optional list of values to fix the joints to. If None, uses the initial
+                            joint positions (q0) from the full robot model.
+        :return: A ReducedRobotModel instance
+        """
+        return cls.from_active_groups(full_robot_model, [active_group_name], fixed_values)
+
+    @classmethod
+    def from_active_groups(
+        cls,
+        full_robot_model: RobotModel,
+        active_group_names: List[str],
+        fixed_values: Optional[List[float]] = None,
+    ) -> "ReducedRobotModel":
+        """
+        Create a reduced order robot model by fixing all joints EXCEPT those in the specified groups.
+        This is useful when you want to keep multiple groups active and fix everything else.
+
+        :param full_robot_model: The original robot model
+        :param active_group_names: List of joint group names to keep active (all other joints will be fixed)
+        :param fixed_values: Optional list of values to fix the joints to. If None, uses the initial
+                            joint positions (q0) from the full robot model.
+        :return: A ReducedRobotModel instance
+        """
+        if full_robot_model.supplemental_info is None:
+            raise ValueError("supplemental_info must be provided to use this method")
+
+        # Get all joints in the active groups, including those from subgroups
+        active_joints = set()
+
+        def add_group_joints(group_name: str):
+            if group_name not in full_robot_model.supplemental_info.joint_groups:
+                raise ValueError(f"Unknown joint group: {group_name}")
+
+            group_info = full_robot_model.supplemental_info.joint_groups[group_name]
+
+            # Add direct joints
+            if "joints" in group_info:
+                active_joints.update(group_info["joints"])
+
+            # Add joints from subgroups
+            if "groups" in group_info:
+                for subgroup_name in group_info["groups"]:
+                    add_group_joints(subgroup_name)
+
+        for group_name in active_group_names:
+            add_group_joints(group_name)
+
+        # Get all joints from the model
+        all_joints = set(full_robot_model.joint_names)
+
+        # The fixed joints are all joints minus the active joints
+        fixed_joints = list(all_joints - active_joints)
+
+        return cls(full_robot_model, fixed_joints, fixed_values)
+
+    def reduced_to_full_configuration(self, q_reduced: np.ndarray) -> np.ndarray:
+        """
+        Convert a reduced configuration to the full configuration space.
+
+        :param q_reduced: Configuration in reduced space
+        :return: Configuration in full space with fixed joints set to their fixed values
+        """
+        if q_reduced.shape[0] != self.num_dofs:
+            raise ValueError(
+                f"Expected q_reduced of length {self.num_dofs}, got {q_reduced.shape[0]} instead"
+            )
+
+        q_full = np.zeros(self.full_robot.num_dofs)
+
+        # Set active joints
+        for reduced_idx, full_idx in enumerate(self.reduced_to_full):
+            q_full[full_idx] = q_reduced[reduced_idx]
+
+        # Set fixed joints
+        for joint_name, value in zip(self.fixed_joints, self.fixed_values):
+            full_idx = self.full_robot.dof_index(joint_name)
+            q_full[full_idx] = value
+
+        return q_full
+
+    def full_to_reduced_configuration(self, q_full: np.ndarray) -> np.ndarray:
+        """
+        Convert a full configuration to the reduced configuration space.
+
+        :param q_full: Configuration in full space
+        :return: Configuration in reduced space
+        """
+        if q_full.shape[0] != self.full_robot.num_dofs:
+            raise ValueError(
+                f"Expected q_full of length {self.full_robot.num_dofs}, got {q_full.shape[0]} instead"
+            )
+
+        q_reduced = np.zeros(self.num_dofs)
+
+        # Copy active joints
+        for reduced_idx, full_idx in enumerate(self.reduced_to_full):
+            q_reduced[reduced_idx] = q_full[full_idx]
+
+        return q_reduced
+
+    def cache_forward_kinematics(self, q_reduced: np.ndarray, auto_clip=True) -> None:
+        """
+        Perform forward kinematics using the reduced configuration.
+
+        :param q_reduced: Configuration in reduced space
+        """
+        # First update the full robot's forward kinematics
+        q_full = self.reduced_to_full_configuration(q_reduced)
+        self.full_robot.cache_forward_kinematics(q_full, auto_clip)
+
+        # Then update the reduced model's forward kinematics
+        pin.framesForwardKinematics(
+            self.pinocchio_wrapper.model, self.pinocchio_wrapper.data, q_reduced
+        )
+
+    def clip_configuration(self, q_reduced: np.ndarray, margin: float = 1e-6) -> np.ndarray:
+        """
+        Clip the reduced configuration to stay within joint limits with a small tolerance.
+
+        :param q_reduced: Configuration to clip
+        :param margin: Tolerance to keep away from joint limits
+        :return: Clipped configuration
+        """
+        q_full = self.reduced_to_full_configuration(q_reduced)
+        q_full_clipped = self.full_robot.clip_configuration(q_full, margin)
+        return self.full_to_reduced_configuration(q_full_clipped)
+
+    def reset_forward_kinematics(self):
+        """
+        Reset the forward kinematics to the initial configuration.
+        """
+        # Reset full robot's forward kinematics
+        self.full_robot.reset_forward_kinematics()
+        # Reset reduced model's forward kinematics
+        self.cache_forward_kinematics(self.q_zero)

GR00T-WholeBodyControl/decoupled_wbc/control/sensor/composed_camera.py

@@ -0,0 +1,440 @@
+from collections import deque
+from dataclasses import dataclass
+import queue
+import threading
+import time
+from typing import Any, Dict, Optional
+
+# we need to import these first in this order to avoid TSL segmentation fault
+# caused by zed and oak libraries
+try:
+    import cv2  # noqa
+    import depthai as dai  # noqa
+    import pyzed.sl as sl  # noqa
+except ImportError:
+    print(
+        """
+    Some of the camera specific dependencies are not installed. If you are
+    not running this on the robot, having these libraries is optional.
+    """
+    )
+
+import numpy as np  # noqa
+
+from decoupled_wbc.control.base.sensor import Sensor
+from decoupled_wbc.control.sensor.sensor_server import (
+    ImageMessageSchema,
+    SensorClient,
+    SensorServer,
+    CameraMountPosition,
+)
+
+
+def read_qr_code(data):
+    current_time = time.monotonic()
+    detector = cv2.QRCodeDetector()
+    for key, img in data["images"].items():
+        decoded_time, bbox, _ = detector.detectAndDecode(img)
+        if bbox is not None and decoded_time:
+            print(f"{key} latency: {(current_time - float(decoded_time)) * 1e3:.1f} ms")
+        else:
+            print(f"{key} QR code not detected.")
+
+
+@dataclass
+class ComposedCameraConfig:
+    """Camera configuration for composed camera"""
+
+    ego_view_camera: Optional[str] = "oak"
+    """Camera type for ego view: oak, realsense, zed, or None"""
+
+    ego_view_device_id: Optional[str] = None
+    """Device ID for ego view camera (optional, used for OAK cameras)"""
+
+    head_camera: Optional[str] = None
+    """Camera type for head view: oak, oak_mono, realsense, zed or None"""
+
+    head_device_id: Optional[str] = None
+    """Device ID for head camera (optional, used for OAK cameras)"""
+
+    left_wrist_camera: Optional[str] = None
+    """Camera type for left wrist view: oak, realsense, zed or None"""
+
+    left_wrist_device_id: Optional[str] = None
+    """Device ID for left wrist camera (optional, used for OAK cameras)"""
+
+    right_wrist_camera: Optional[str] = None
+    """Camera type for right wrist view: oak, realsense, zed or None"""
+
+    right_wrist_device_id: Optional[str] = None
+    """Device ID for right wrist camera (optional, used for OAK cameras)"""
+
+    fps: int = 30
+    """Rate at which the composed camera will publish the images. Since composed camera
+    can read from multiple cameras, it will publish all the images.
+    Note that OAK can only run at 30 FPS. 20 FPS will cause large latency.
+    """
+
+    # Server configuration
+    run_as_server: bool = True
+    """Whether to run as server or client"""
+
+    server: bool = True
+    """Whether to run the camera as a server"""
+
+    port: int = 5555
+    """Port number for server/client communication"""
+
+    test_latency: bool = False
+    """Whether to test latency"""
+
+    # Queue configuration
+    queue_size: int = 3
+    """Size of each camera's image queue"""
+
+    def __post_init__(self):
+        # runyu: Note that this is a hack to make the config work with G1 camera server in orin
+        # we should not use this hack in the future
+        self.run_as_server: bool = self.server
+
+
+class ComposedCameraSensor(Sensor, SensorServer):
+
+    def __init__(self, config: ComposedCameraConfig):
+        self.config = config
+        self.camera_queues: Dict[str, queue.Queue] = {}
+        self.camera_threads: Dict[str, threading.Thread] = {}
+        self.shutdown_events: Dict[str, threading.Event] = {}
+        self.error_events: Dict[str, threading.Event] = {}
+        self.error_messages: Dict[str, str] = {}
+        self._observation_spaces: Dict[str, Any] = {}
+
+        camera_configs = self._get_camera_configs()
+
+        # Then create worker threads
+        for mount_position, camera_config in camera_configs.items():
+            # Create queue and shutdown event for this camera
+            camera_queue = queue.Queue(maxsize=config.queue_size)
+            shutdown_event = threading.Event()
+            error_event = threading.Event()
+
+            self.camera_queues[mount_position] = camera_queue
+            self.shutdown_events[mount_position] = shutdown_event
+            self.error_events[mount_position] = error_event
+
+            # Start camera thread
+            thread = threading.Thread(
+                target=self._camera_worker_wrapper,
+                args=(
+                    mount_position,
+                    camera_config["camera_type"],
+                    camera_config["device_id"],
+                    camera_queue,
+                    shutdown_event,
+                    error_event,
+                ),
+            )
+            thread.start()
+            self.camera_threads[mount_position] = thread
+
+        if config.run_as_server:
+            self.start_server(config.port)
+
+    def _get_camera_configs(self) -> Dict[str, str]:
+        """Get camera configurations as mount_position -> camera_type mapping"""
+        camera_configs = {}
+
+        if self.config.ego_view_camera is not None:
+            camera_configs[CameraMountPosition.EGO_VIEW.value] = {
+                "camera_type": self.config.ego_view_camera,
+                "device_id": self.config.ego_view_device_id,
+            }
+
+        if self.config.head_camera is not None:
+            camera_configs[CameraMountPosition.HEAD.value] = {
+                "camera_type": self.config.head_camera,
+                "device_id": self.config.head_device_id,
+            }
+
+        if self.config.left_wrist_camera is not None:
+            camera_configs[CameraMountPosition.LEFT_WRIST.value] = {
+                "camera_type": self.config.left_wrist_camera,
+                "device_id": self.config.left_wrist_device_id,
+            }
+
+        if self.config.right_wrist_camera is not None:
+            camera_configs[CameraMountPosition.RIGHT_WRIST.value] = {
+                "camera_type": self.config.right_wrist_camera,
+                "device_id": self.config.right_wrist_device_id,
+            }
+
+        return camera_configs
+
+    def _camera_worker_wrapper(
+        self,
+        mount_position: str,
+        camera_type: str,
+        device_id: Optional[str],
+        image_queue: queue.Queue,
+        shutdown_event: threading.Event,
+        error_event: threading.Event,
+    ):
+        """Worker thread that continuously captures from a single camera"""
+        try:
+            camera = self._instantiate_camera(mount_position, camera_type, device_id)
+            self._observation_spaces[mount_position] = camera.observation_space()
+
+            consecutive_failures = 0
+            max_consecutive_failures = 5
+
+            while not shutdown_event.is_set():
+                frame = camera.read()
+                if frame:
+                    consecutive_failures = 0  # Reset on successful read
+                    # Non-blocking queue put with frame dropping
+                    try:
+                        image_queue.put_nowait(frame)
+                    except queue.Full:
+                        # Remove oldest frame and add new one
+                        try:
+                            image_queue.get_nowait()
+                            image_queue.put_nowait(frame)
+                        except queue.Empty:
+                            pass
+                else:
+                    consecutive_failures += 1
+                    if consecutive_failures >= max_consecutive_failures:
+                        error_msg = (
+                            f"Camera {mount_position} ({camera_type}) dropped: "
+                            f"failed to read {consecutive_failures} consecutive frames"
+                        )
+                        print(f"[ERROR] {error_msg}")
+                        self.error_messages[mount_position] = error_msg
+                        error_event.set()
+                        break
+
+            camera.close()
+
+        except Exception as e:
+            error_msg = f"Camera {mount_position} ({camera_type}) error: {str(e)}"
+            print(f"[ERROR] {error_msg}")
+            self.error_messages[mount_position] = error_msg
+            error_event.set()
+
+    def _instantiate_camera(
+        self, mount_position: str, camera_type: str, device_id: Optional[str] = None
+    ) -> Sensor:
+        """
+        Instantiate a camera sensor based on the camera type.
+
+        Args:
+            camera_type: Type of camera ("oak", "oak_mono", "realsense", "zed")
+            device_id: Optional device ID for the camera (used for OAK cameras)
+
+        Returns:
+            Sensor instance for the specified camera type
+        """
+        if camera_type in ("oak", "oak_mono"):
+            from decoupled_wbc.control.sensor.oak import OAKConfig, OAKSensor
+
+            oak_config = OAKConfig()
+            if camera_type == "oak_mono":
+                oak_config.enable_mono_cameras = True
+            print("Initializing OAK sensor for camera type: ", camera_type)
+            return OAKSensor(config=oak_config, mount_position=mount_position, device_id=device_id)
+        elif camera_type == "realsense":
+            from decoupled_wbc.control.sensor.realsense import RealSenseSensor
+
+            print("Initializing RealSense sensor for camera type: ", camera_type)
+            return RealSenseSensor(mount_position=mount_position)
+        elif camera_type == "zed":
+            from decoupled_wbc.control.sensor.zed import ZEDSensor
+
+            print("Initializing ZED sensor for camera type: ", camera_type)
+            return ZEDSensor(mount_position=mount_position)
+        elif camera_type.endswith(".mp4"):
+            from decoupled_wbc.control.sensor.dummy import ReplayDummySensor
+
+            print("Initializing Replay Dummy Sensor for camera type: ", camera_type)
+            return ReplayDummySensor(video_path=camera_type)
+        else:
+            raise ValueError(f"Unsupported camera type: {camera_type}")
+
+    def _check_for_errors(self):
+        """Check if any camera thread has encountered an error and raise exception if so."""
+        for mount_position, error_event in self.error_events.items():
+            if error_event.is_set():
+                error_msg = self.error_messages.get(
+                    mount_position, f"Camera {mount_position} encountered an unknown error"
+                )
+                raise RuntimeError(error_msg)
+
+    def read(self):
+        """Read frames from all cameras."""
+        # Check for errors from camera threads
+        self._check_for_errors()
+
+        message = {}
+        for mount_position, camera_queue in self.camera_queues.items():
+            frame = self._get_latest_from_queue(camera_queue)
+            if frame is not None:
+                message[mount_position] = frame
+        return message
+
+    def _get_latest_from_queue(self, camera_queue: queue.Queue) -> Optional[Dict[str, Any]]:
+        """Get most recent frame, discard older ones"""
+        latest = None
+        try:
+            while True:
+                latest = camera_queue.get_nowait()
+        except queue.Empty:
+            pass
+        return latest
+
+    def close(self):
+        """Close all cameras."""
+        # Signal all worker threads to shutdown
+        for shutdown_event in self.shutdown_events.values():
+            shutdown_event.set()
+
+        # Wait for all threads to finish
+        for thread in self.camera_threads.values():
+            thread.join(timeout=5.0)
+
+        # Clear queues
+        for camera_queue in self.camera_queues.values():
+            try:
+                while True:
+                    camera_queue.get_nowait()
+            except queue.Empty:
+                pass
+
+        # Stop server if running
+        if self.config.run_as_server:
+            self.stop_server()
+
+    def serialize_message(self, message: Dict[str, Any]) -> Dict[str, Any]:
+        """Merge all camera data into a single ImageMessageSchema."""
+        all_timestamps = {}
+        all_images = {}
+
+        for _, camera_data in message.items():
+            all_timestamps.update(camera_data.get("timestamps", {}))
+            all_images.update(camera_data.get("images", {}))
+
+        # Create a single ImageMessageSchema with all data
+        img_schema = ImageMessageSchema(timestamps=all_timestamps, images=all_images)
+        return img_schema.serialize()
+
+    def run_server(self):
+        """Run the server."""
+        idx = 0
+        server_start_time = time.monotonic()
+        fps_print_time = time.monotonic()
+        frame_interval = 1.0 / self.config.fps
+
+        while True:
+            # Calculate when this frame should ideally complete
+            target_time = server_start_time + (idx + 1) * frame_interval
+
+            message = self.read()
+            if message:
+                if self.config.test_latency:
+                    read_qr_code(message)
+
+                serialized_message = self.serialize_message(message)
+                self.send_message(serialized_message)
+                idx += 1
+
+                if idx % 10 == 0:
+                    print(f"Image sending FPS: {10 / (time.monotonic() - fps_print_time):.2f}")
+                    fps_print_time = time.monotonic()
+
+            # Sleep to maintain precise timing
+            current_time = time.monotonic()
+            sleep_time = target_time - current_time
+            if sleep_time > 0:
+                time.sleep(sleep_time)
+            else:
+                # If we're behind, increment idx to stay on schedule
+                if not message:
+                    idx += 1
+
+    def observation_space(self):
+        """Return the observation space."""
+        import gymnasium as gym
+
+        return gym.spaces.Dict(self._observation_spaces)
+
+
+class ComposedCameraClientSensor(Sensor, SensorClient):
+    """Class that serves as client for multiple different cameras."""
+
+    def __init__(self, server_ip: str = "localhost", port: int = 5555):
+        self.start_client(server_ip, port)
+
+        # Initialize tracking variables
+        self._latest_message = {}
+        self._avg_time_per_frame = deque(maxlen=20)
+        self._msg_received_time = 0
+        self._start_time = 0.0  # Initialize _start_time
+        self.idx = 0
+
+        print("Initialized composed camera client sensor")
+
+    def read(self, **kwargs) -> Optional[Dict[str, Any]]:
+        self._start_time = time.time()
+        message = self.receive_message()
+        if not message:
+            return None
+        self.idx += 1
+
+        self._latest_message = ImageMessageSchema.deserialize(message).asdict()
+
+        # if self.idx % 10 == 0:
+        #     for image_key, image_time in self._latest_message["timestamps"].items():
+        #         image_latency = (time.time() - image_time) * 1000
+        # print(f"Image latency for {image_key}: {image_latency:.2f} ms")
+
+        self._msg_received_time = time.time()
+        self._avg_time_per_frame.append(self._msg_received_time - self._start_time)
+
+        return self._latest_message
+
+    def close(self):
+        """Close the client connection."""
+        self.stop_client()
+
+    def fps(self) -> float:
+        """Get the current FPS of the client."""
+        if len(self._avg_time_per_frame) == 0:
+            return 0.0
+        return float(1 / np.mean(self._avg_time_per_frame))
+
+
+if __name__ == "__main__":
+    """Test function for ComposedCamera."""
+    import tyro
+
+    config = tyro.cli(ComposedCameraConfig)
+
+    if config.run_as_server:
+        composed_camera = ComposedCameraSensor(config)
+        print("Running composed camera server...")
+        composed_camera.run_server()
+
+    else:
+        # Client mode
+        composed_client = ComposedCameraClientSensor(server_ip="localhost", port=config.port)
+
+        try:
+            while True:
+                data = composed_client.read()
+                if data is not None:
+                    print(f"FPS: {composed_client.fps():.2f}")
+                    if "timestamp" in data:
+                        print(f"Timestamp: {data['timestamp']}")
+                time.sleep(0.1)
+        except KeyboardInterrupt:
+            print("Stopping client...")
+            composed_client.close()

GR00T-WholeBodyControl/decoupled_wbc/control/sensor/oak.py

@@ -0,0 +1,324 @@
+import time
+from typing import Any, Dict, Optional, Tuple
+
+import cv2
+import depthai as dai
+import gymnasium as gym
+import numpy as np
+
+from decoupled_wbc.control.base.sensor import Sensor
+from decoupled_wbc.control.sensor.sensor_server import (
+    CameraMountPosition,
+    ImageMessageSchema,
+    SensorServer,
+)
+
+
+class OAKConfig:
+    """Configuration for the OAK camera."""
+
+    color_image_dim: Tuple[int, int] = (640, 480)  # RGB camera resolution
+    monochrome_image_dim: Tuple[int, int] = (640, 480)  # Monochrome camera resolution
+    fps: int = 30
+    enable_color: bool = True  # Enable CAM_A (RGB)
+    enable_mono_cameras: bool = False  # Enable CAM_B & CAM_C (Monochrome stereo pair)
+    mount_position: str = CameraMountPosition.EGO_VIEW.value
+
+
+class OAKSensor(Sensor, SensorServer):
+    """Sensor for the OAK camera family."""
+
+    def __init__(
+        self,
+        run_as_server: bool = False,
+        port: int = 5555,
+        config: OAKConfig = OAKConfig(),
+        device_id: Optional[str] = None,
+        mount_position: str = CameraMountPosition.EGO_VIEW.value,
+    ):
+        """Initialize the OAK camera."""
+        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)
+                    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]}")
+
+        # Create pipeline (without context manager to persist across method calls)
+        self.pipeline = dai.Pipeline(self.device)
+        self.output_queues = {}
+
+        # Configure 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)
+            # Create RGB output queue
+            self.output_queues["color"] = self.cam_rgb.requestOutput(
+                config.color_image_dim,
+                fps=config.fps,
+            ).createOutputQueue()
+            print("Enabled CAM_A (RGB)")
+
+        # Configure 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)
+                # Create mono left output queue
+                self.output_queues["mono_left"] = self.cam_mono_left.requestOutput(
+                    config.monochrome_image_dim,
+                    fps=config.fps,
+                ).createOutputQueue()
+                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)
+                # Create mono right output queue
+                self.output_queues["mono_right"] = self.cam_mono_right.requestOutput(
+                    config.monochrome_image_dim,
+                    fps=config.fps,
+                ).createOutputQueue()
+                print("Enabled CAM_C (Monochrome Right)")
+
+        assert len(self.output_queues) > 0, "No output queues enabled"
+        # auto exposure compensation, for CoRL demo
+        # cam_q_in = self.cam_rgb.inputControl.createInputQueue()
+        # ctrl = dai.CameraControl()
+        # ctrl.setAutoExposureEnable()
+        # ctrl.setAutoExposureCompensation(-2)
+        # cam_q_in.send(ctrl)
+
+        # Start pipeline on device
+        self.pipeline.start()
+
+        if run_as_server:
+            self.start_server(port)
+
+    def read(self) -> Optional[Dict[str, Any]]:
+        """Read images from the camera."""
+        if not self.pipeline.isRunning():
+            print(f"[ERROR] OAK pipeline stopped for {self.mount_position}")
+            return None
+
+        # Check if device is still connected
+        if not self.device.isPipelineRunning():
+            print(f"[ERROR] OAK device disconnected for {self.mount_position}")
+            return None
+
+        timestamps = {}
+        images = {}
+        rgb_frame_time = None
+
+        # Get color frame if enabled
+        if "color" in self.output_queues:
+            try:
+                rgb_frame = self.output_queues["color"].get()
+                rgb_frame_time = rgb_frame.getTimestamp()
+                if rgb_frame is not None:
+                    images[self.mount_position] = rgb_frame.getCvFrame()[..., ::-1]  # BGR to RGB
+                    timestamps[self.mount_position] = (
+                        rgb_frame_time - dai.Clock.now()
+                    ).total_seconds() + time.time()
+            except Exception as e:
+                print(f"[ERROR] Failed to read color frame from {self.mount_position}: {e}")
+                return None
+
+        # Get mono frames if enabled
+        if "mono_left" in self.output_queues:
+            try:
+                mono_left_frame = self.output_queues["mono_left"].get()
+                mono_left_frame_time = mono_left_frame.getTimestamp()
+                if mono_left_frame is not None:
+                    key = f"{self.mount_position}_left_mono"
+                    images[key] = mono_left_frame.getCvFrame()
+                    timestamps[key] = (
+                        mono_left_frame_time - dai.Clock.now()
+                    ).total_seconds() + time.time()
+            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 = self.output_queues["mono_right"].get()
+                mono_right_frame_time = mono_right_frame.getTimestamp()
+                if mono_right_frame is not None:
+                    key = f"{self.mount_position}_right_mono"
+                    images[key] = mono_right_frame.getCvFrame()
+                    timestamps[key] = (
+                        mono_right_frame_time - dai.Clock.now()
+                    ).total_seconds() + time.time()
+            except Exception as e:
+                print(f"[ERROR] Failed to read mono_right frame from {self.mount_position}: {e}")
+                return None
+
+        if (
+            rgb_frame_time is not None
+            and (rgb_frame_time - dai.Clock.now()).total_seconds() <= -0.2
+        ):
+            print(
+                f"[{self.mount_position}] OAK latency too large: "
+                f"{(dai.Clock.now() - rgb_frame_time).total_seconds() * 1000}ms"
+            )
+
+        return {
+            "timestamps": timestamps,
+            "images": images,
+        }
+
+    def serialize(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Serialize data using ImageMessageSchema."""
+        serialized_msg = ImageMessageSchema(timestamps=data["timestamps"], images=data["images"])
+        return serialized_msg.serialize()
+
+    def observation_space(self) -> gym.Space:
+        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):
+        """Close the camera connection."""
+        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):
+        """Run the server."""
+        if not self._run_as_server:
+            raise ValueError("This function is only available when run_as_server is True")
+
+        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__":
+    """Test function for OAK camera."""
+
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--server", action="store_true", help="Run as server")
+    parser.add_argument("--client", action="store_true", help="Run as client")
+    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")
+    args = parser.parse_args()
+
+    oak_config = OAKConfig()
+    if args.enable_mono:
+        oak_config.enable_mono_cameras = True
+
+    if args.server:
+        # Run as 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:
+        # Run standalone
+        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 "color_image" in frame:
+                print(f"Color image shape: {frame['color_image'].shape}")
+            if "mono_left_image" in frame:
+                print(f"Mono left image shape: {frame['mono_left_image'].shape}")
+            if "mono_right_image" in frame:
+                print(f"Mono right image shape: {frame['mono_right_image'].shape}")
+            if "depth_image" in frame:
+                print(f"Depth image shape: {frame['depth_image'].shape}")
+
+            if args.show_image:
+                if "color_image" in frame:
+                    cv2.imshow("Color Image", frame["color_image"])
+
+                if "mono_left_image" in frame:
+                    cv2.imshow("Mono Left", frame["mono_left_image"])
+                if "mono_right_image" in frame:
+                    cv2.imshow("Mono Right", frame["mono_right_image"])
+
+                if "depth_image" in frame:
+                    depth_colormap = cv2.applyColorMap(
+                        cv2.convertScaleAbs(frame["depth_image"], alpha=0.03), cv2.COLORMAP_JET
+                    )
+                    cv2.imshow("Depth Image", depth_colormap)
+
+                if cv2.waitKey(1) == ord("q"):
+                    break
+
+            time.sleep(0.01)
+
+        cv2.destroyAllWindows()
+        oak.close()

GR00T-WholeBodyControl/decoupled_wbc/control/sensor/sensor_server.py

@@ -0,0 +1,128 @@
+import base64
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Dict
+
+import cv2
+import msgpack
+import msgpack_numpy as m
+import numpy as np
+import zmq
+
+
+@dataclass
+class ImageMessageSchema:
+    """
+    This is a standardized message schema for image data.
+    Any camera should use this schema to serialize (send to queue) and
+    deserialize (receive from queue) the image data.
+
+    """
+
+    timestamps: Dict[str, float]
+    """Dictionary of timestamps, keyed by image identifier (e.g., {"ego_view": 123.45})"""
+    images: Dict[str, np.ndarray]
+    """Dictionary of images, keyed by image identifier (e.g., {"ego_view": array, "ego_view_left_mono": array})"""
+
+    def serialize(self) -> Dict[str, Any]:
+        """Serialize the message for transmission."""
+        serialized_msg = {"timestamps": self.timestamps, "images": {}}
+        for key, image in self.images.items():
+            serialized_msg["images"][key] = ImageUtils.encode_image(image)
+        return serialized_msg
+
+    @staticmethod
+    def deserialize(data: Dict[str, Any]) -> "ImageMessageSchema":
+        """Deserialize received message data."""
+        timestamps = data.get("timestamps", {})
+        images = {}
+        for key, value in data.get("images", {}).items():
+            if isinstance(value, str):
+                images[key] = ImageUtils.decode_image(value)
+            else:
+                images[key] = value
+        return ImageMessageSchema(timestamps=timestamps, images=images)
+
+    def asdict(self) -> Dict[str, Any]:
+        """Convert to dictionary format."""
+        return {"timestamps": self.timestamps, "images": self.images}
+
+
+class SensorServer:
+    def start_server(self, port: int):
+        self.context = zmq.Context()
+        self.socket = self.context.socket(zmq.PUB)
+        self.socket.setsockopt(zmq.SNDHWM, 20)  # high water mark
+        self.socket.setsockopt(zmq.LINGER, 0)
+        self.socket.bind(f"tcp://*:{port}")
+        print(f"Sensor server running at tcp://*:{port}")
+
+        self.message_sent = 0
+        self.message_dropped = 0
+
+    def stop_server(self):
+        self.socket.close()
+        self.context.term()
+
+    def send_message(self, data: Dict[str, Any]):
+        try:
+            packed = msgpack.packb(data, use_bin_type=True)
+            self.socket.send(packed, flags=zmq.NOBLOCK)
+        except zmq.Again:
+            self.message_dropped += 1
+            print(f"[Warning] message dropped: {self.message_dropped}")
+        self.message_sent += 1
+
+        if self.message_sent % 100 == 0:
+            print(
+                f"[Sensor server] Message sent: {self.message_sent}, message dropped: {self.message_dropped}"
+            )
+
+
+class SensorClient:
+    def start_client(self, server_ip: str, port: int):
+        self.context = zmq.Context()
+        self.socket = self.context.socket(zmq.SUB)
+        self.socket.setsockopt_string(zmq.SUBSCRIBE, "")
+        self.socket.setsockopt(zmq.CONFLATE, True)  # last msg only.
+        self.socket.setsockopt(zmq.RCVHWM, 3)  # queue size 3 for receive buffer
+        self.socket.connect(f"tcp://{server_ip}:{port}")
+
+    def stop_client(self):
+        self.socket.close()
+        self.context.term()
+
+    def receive_message(self):
+        packed = self.socket.recv()
+        return msgpack.unpackb(packed, object_hook=m.decode)
+
+
+class CameraMountPosition(Enum):
+    EGO_VIEW = "ego_view"
+    HEAD = "head"
+    LEFT_WRIST = "left_wrist"
+    RIGHT_WRIST = "right_wrist"
+
+
+class ImageUtils:
+    @staticmethod
+    def encode_image(image: np.ndarray) -> str:
+        _, color_buffer = cv2.imencode(".jpg", image, [int(cv2.IMWRITE_JPEG_QUALITY), 80])
+        return base64.b64encode(color_buffer).decode("utf-8")
+
+    @staticmethod
+    def encode_depth_image(image: np.ndarray) -> str:
+        depth_compressed = cv2.imencode(".png", image)[1].tobytes()
+        return base64.b64encode(depth_compressed).decode("utf-8")
+
+    @staticmethod
+    def decode_image(image: str) -> np.ndarray:
+        color_data = base64.b64decode(image)
+        color_array = np.frombuffer(color_data, dtype=np.uint8)
+        return cv2.imdecode(color_array, cv2.IMREAD_COLOR)
+
+    @staticmethod
+    def decode_depth_image(image: str) -> np.ndarray:
+        depth_data = base64.b64decode(image)
+        depth_array = np.frombuffer(depth_data, dtype=np.uint8)
+        return cv2.imdecode(depth_array, cv2.IMREAD_UNCHANGED)

GR00T-WholeBodyControl/decoupled_wbc/control/teleop/teleop_retargeting_ik.py

@@ -0,0 +1,148 @@
+import time
+from typing import List, Optional
+
+import numpy as np
+
+from decoupled_wbc.control.base.policy import Policy
+from decoupled_wbc.control.robot_model.robot_model import ReducedRobotModel, RobotModel
+from decoupled_wbc.control.teleop.solver.body.body_ik_solver import BodyIKSolver
+from decoupled_wbc.control.teleop.solver.body.body_ik_solver_settings import BodyIKSolverSettings
+from decoupled_wbc.control.teleop.solver.solver import Solver
+from decoupled_wbc.control.visualization.humanoid_visualizer import RobotVisualizer
+
+
+class TeleopRetargetingIK(Policy):
+    """
+    Robot-agnostic teleop retargeting inverse kinematics code.
+    Focus only on IK processing, ignore commands.
+    """
+
+    def __init__(
+        self,
+        robot_model: RobotModel,
+        left_hand_ik_solver: Solver,
+        right_hand_ik_solver: Solver,
+        enable_visualization=False,
+        body_active_joint_groups: Optional[List[str]] = None,
+        body_ik_solver_settings_type: str = "default",
+    ):
+        # initialize the body
+        if body_active_joint_groups is not None:
+            self.body = ReducedRobotModel.from_active_groups(robot_model, body_active_joint_groups)
+            self.full_robot = self.body.full_robot
+            self.using_reduced_robot_model = True
+        else:
+            self.body = robot_model
+            self.full_robot = self.body
+            self.using_reduced_robot_model = False
+        if body_ik_solver_settings_type == "default":
+            body_ik_solver_settings = BodyIKSolverSettings()
+        else:
+            raise ValueError(
+                f"Unknown body_ik_solver_settings_type: {body_ik_solver_settings_type}"
+            )
+        self.body_ik_solver = BodyIKSolver(body_ik_solver_settings)
+
+        # We register the specific robot model to the robot-agnostic body IK solver class
+        self.body_ik_solver.register_robot(self.body)
+
+        # Hand IK solvers are hand specific, so we pass them in the constructor
+        self.left_hand_ik_solver = left_hand_ik_solver
+        self.right_hand_ik_solver = right_hand_ik_solver
+
+        # enable visualizer
+        self.enable_visualization = enable_visualization
+        if self.enable_visualization:
+            self.visualizer = RobotVisualizer(self.full_robot)
+            self.visualizer.visualize(self.full_robot.q_zero)
+            time.sleep(1)  # wait for the visualizer to start
+
+        self.in_warmup = True
+        self._most_recent_ik_data = None
+        self._most_recent_q = self.full_robot.default_body_pose.copy()
+
+    def compute_joint_positions(
+        self, body_data: dict, left_hand_data: dict, right_hand_data: dict
+    ) -> np.ndarray:
+        """Process only IK-related data, return joint positions"""
+        if self.in_warmup:
+            # TODO: Warmup is not necessary if we start IK from the current robot qpos, rather than the zero qpos
+            for _ in range(50):
+                target_robot_joints = self._inverse_kinematics(
+                    body_data, left_hand_data, right_hand_data
+                )
+            self.in_warmup = False
+        else:
+            target_robot_joints = self._inverse_kinematics(
+                body_data, left_hand_data, right_hand_data
+            )
+
+        return target_robot_joints
+
+    def _inverse_kinematics(
+        self,
+        body_target_pose,
+        left_hand_target_pose,
+        right_hand_target_pose,
+    ):
+        """
+        Solve the inverse kinematics problem for the given target poses.
+        Args:
+            body_target_pose: Dictionary of link names and their corresponding target pose.
+            left_hand_target_pose: Dictionary with key "position" mapping to a (25, 4, 4) np.ndarray from AVP data
+            right_hand_target_pose: Dictionary with key "position" mapping to a (25, 4, 4) np.ndarray from AVP data
+            q: Initial configuration vector.
+        Returns:
+            Configuration vector that achieves the target poses.
+        """
+        if body_target_pose:
+            if self.using_reduced_robot_model:
+                body_q = self.body.reduced_to_full_configuration(
+                    self.body_ik_solver(body_target_pose)
+                )
+            else:
+                body_q = self.body_ik_solver(body_target_pose)
+        else:
+            # If no body target pose is provided, set the body to the default pose
+            body_q = self.full_robot.default_body_pose.copy()
+
+        if left_hand_target_pose is not None:
+            left_hand_actuated_q = self.left_hand_ik_solver(left_hand_target_pose)
+            body_q[self.full_robot.get_hand_actuated_joint_indices(side="left")] = (
+                left_hand_actuated_q
+            )
+
+        if right_hand_target_pose is not None:
+            right_hand_actuated_q = self.right_hand_ik_solver(right_hand_target_pose)
+            body_q[self.full_robot.get_hand_actuated_joint_indices(side="right")] = (
+                right_hand_actuated_q
+            )
+
+        if self.enable_visualization:
+            self.visualizer.visualize(np.array(body_q))
+
+        return body_q
+
+    def reset(self):
+        """Reset the robot model and IK solvers to the initial state, and re-activate the warmup procedure."""
+        self.body.reset_forward_kinematics()  # self.body is the same one as self.body_ik_solver.robot
+        self.full_robot.reset_forward_kinematics()
+        self.body_ik_solver.initialize()
+        # If in the future, the hand IK solver has initialize method, call it
+        self._most_recent_ik_data = None
+        self._most_recent_q = self.full_robot.default_body_pose.copy()
+        self.in_warmup = True
+
+    def set_goal(self, ik_data: dict):
+        self._most_recent_ik_data = ik_data
+
+    def get_action(self) -> dict[str, any]:
+        # Process IK if active
+        if self._most_recent_ik_data is not None:
+            body_data = self._most_recent_ik_data["body_data"]
+            left_hand_data = self._most_recent_ik_data["left_hand_data"]
+            right_hand_data = self._most_recent_ik_data["right_hand_data"]
+            target_joints = self.compute_joint_positions(body_data, left_hand_data, right_hand_data)
+            self._most_recent_q = target_joints
+
+        return self._most_recent_q[self.full_robot.get_joint_group_indices("upper_body")]

GR00T-WholeBodyControl/decoupled_wbc/control/teleop/teleop_streamer.py

@@ -0,0 +1,240 @@
+from math import floor
+import pickle
+from typing import Optional
+
+from decoupled_wbc.control.robot_model.robot_model import RobotModel
+from decoupled_wbc.control.teleop.pre_processor.fingers.fingers import FingersPreProcessor
+from decoupled_wbc.control.teleop.pre_processor.wrists.wrists import WristsPreProcessor
+from decoupled_wbc.control.teleop.streamers.base_streamer import StreamerOutput
+
+
+class TeleopStreamer:
+    def __init__(
+        self,
+        robot_model: RobotModel,
+        body_control_device: Optional[str] = None,
+        hand_control_device: Optional[str] = None,
+        enable_real_device=True,
+        body_streamer_ip="",
+        body_streamer_keyword="",
+        replay_data_path: Optional[str] = None,
+        replay_speed: float = 1.0,
+    ):
+        # initialize the body
+        self.body = robot_model
+
+        self.body_control_device = body_control_device
+        self.hand_control_device = hand_control_device
+        self.body_streamer_ip = body_streamer_ip
+        self.body_streamer_keyword = body_streamer_keyword
+        self.replay_speed = replay_speed
+
+        # enable real robot and devices
+        self.enable_real_device = enable_real_device
+        if self.enable_real_device:
+            if body_control_device == "vive":
+                from decoupled_wbc.control.teleop.streamers.vive_streamer import ViveStreamer
+
+                self.body_streamer = ViveStreamer(
+                    ip=self.body_streamer_ip, keyword=self.body_streamer_keyword
+                )
+                self.body_streamer.start_streaming()
+            elif body_control_device == "iphone":
+                from decoupled_wbc.control.teleop.streamers.iphone_streamer import IphoneStreamer
+
+                self.body_streamer = IphoneStreamer()
+                self.body_streamer.start_streaming()
+            elif body_control_device == "leapmotion":
+                from decoupled_wbc.control.teleop.streamers.leapmotion_streamer import (
+                    LeapMotionStreamer,
+                )
+
+                self.body_streamer = LeapMotionStreamer()
+                self.body_streamer.start_streaming()
+            elif body_control_device == "joycon":
+                from decoupled_wbc.control.teleop.streamers.joycon_streamer import JoyconStreamer
+
+                self.body_streamer = JoyconStreamer()
+                self.body_streamer.start_streaming()
+
+            elif body_control_device == "pico":
+                from decoupled_wbc.control.teleop.streamers.pico_streamer import PicoStreamer
+
+                self.body_streamer = PicoStreamer()
+                self.body_streamer.start_streaming()
+            elif body_control_device == "dummy":
+                from decoupled_wbc.control.teleop.streamers.dummy_streamer import DummyStreamer
+
+                self.body_streamer = DummyStreamer()
+                self.body_streamer.start_streaming()
+            else:
+                self.body_streamer = None
+
+            if hand_control_device and hand_control_device != body_control_device:
+                if hand_control_device == "manus":
+                    from decoupled_wbc.control.teleop.streamers.manus_streamer import ManusStreamer
+
+                    self.hand_streamer = ManusStreamer()
+                    self.hand_streamer.start_streaming()
+                elif hand_control_device == "joycon":
+                    from decoupled_wbc.control.teleop.streamers.joycon_streamer import JoyconStreamer
+
+                    self.hand_streamer = JoyconStreamer()
+                    self.hand_streamer.start_streaming()
+                elif hand_control_device == "iphone":
+                    from decoupled_wbc.control.teleop.streamers.iphone_streamer import IphoneStreamer
+
+                    self.hand_streamer = IphoneStreamer()
+                    self.hand_streamer.start_streaming()
+                elif hand_control_device == "pico":
+                    from decoupled_wbc.control.teleop.streamers.pico_streamer import PicoStreamer
+
+                    self.hand_streamer = PicoStreamer()
+                    self.hand_streamer.start_streaming()
+                else:
+                    self.hand_streamer = None
+            else:
+                self.hand_streamer = None
+        else:
+            self.body_streamer = None
+            self.hand_streamer = None
+
+        self.raw_replay_data = None
+        self.replay_calibration_data = None
+        self.replay_mode = False
+        if replay_data_path and not self.enable_real_device:
+            with open(replay_data_path, "rb") as f:
+                data_ = pickle.load(f)
+            self.raw_replay_data = data_["replay_data"]
+            self.replay_calibration_data = data_["calibration_data"]
+            print("Found teleop replay data in file: ", replay_data_path)
+            self.replay_idx = 0
+            self.replay_mode = True
+
+        # initialize pre_processors
+        self.body_control_device = body_control_device
+        if body_control_device or self.replay_mode:
+            self.body_pre_processor = WristsPreProcessor(
+                motion_scale=robot_model.supplemental_info.teleop_upper_body_motion_scale
+            )
+            self.body_pre_processor.register(self.body)
+        else:
+            self.body_pre_processor = None
+
+        # initialize hand pre-processors and post-processors
+        self.hand_control_device = hand_control_device
+        if hand_control_device or self.replay_mode:
+            self.left_hand_pre_processor = FingersPreProcessor(side="left")
+            self.right_hand_pre_processor = FingersPreProcessor(side="right")
+
+        else:
+            self.left_hand_pre_processor = None
+            self.right_hand_pre_processor = None
+
+        self.is_calibrated = False
+
+    def _get_replay_data(self) -> StreamerOutput:
+        streamer_data = StreamerOutput()
+
+        if self.replay_idx < len(self.raw_replay_data):
+            streamer_data.ik_data.update(
+                self.raw_replay_data[floor(self.replay_idx / self.replay_speed)]
+            )
+            self.replay_idx += 1
+
+        return streamer_data
+
+    def _get_live_data(self) -> StreamerOutput:
+        """Get structured data instead of raw dict"""
+        if self.body_streamer:
+            streamer_data = self.body_streamer.get()
+        else:
+            streamer_data = StreamerOutput()
+
+        if self.hand_streamer and self.hand_streamer != self.body_streamer:
+            hand_data = self.hand_streamer.get()
+
+            # Merge hand data into body data (hand data takes precedence)
+            streamer_data.ik_data.update(hand_data.ik_data)
+            streamer_data.control_data.update(hand_data.control_data)
+            streamer_data.teleop_data.update(hand_data.teleop_data)
+            streamer_data.data_collection_data.update(hand_data.data_collection_data)
+
+        return streamer_data
+
+    def get_streamer_data(self) -> StreamerOutput:
+        if self.enable_real_device:
+            streamer_data = self._get_live_data()
+        elif self.replay_mode:
+            streamer_data = self._get_replay_data()
+        else:
+            streamer_data = StreamerOutput()
+
+        if self.is_calibrated and streamer_data.ik_data:
+            body_data, left_hand_data, right_hand_data = self.pre_process(streamer_data.ik_data)
+            streamer_data.ik_data = {
+                "body_data": body_data,
+                "left_hand_data": left_hand_data,
+                "right_hand_data": right_hand_data,
+            }
+        elif not self.is_calibrated:
+            streamer_data.ik_data = {}
+
+        return streamer_data
+
+    def calibrate(self):
+        """Calibrate the pre-processors using only IK data."""
+        if self.replay_mode:
+            ik_data = self.replay_calibration_data
+        else:
+            streamer_data = self._get_live_data()
+            ik_data = streamer_data.ik_data
+
+        if self.body_pre_processor:
+            self.body_pre_processor.calibrate(ik_data, self.body_control_device)
+        if self.left_hand_pre_processor:
+            self.left_hand_pre_processor.calibrate(ik_data, self.hand_control_device)
+        if self.right_hand_pre_processor:
+            self.right_hand_pre_processor.calibrate(ik_data, self.hand_control_device)
+
+        self.is_calibrated = True
+
+    def pre_process(self, raw_data):
+        """Pre-process the raw data."""
+        assert (
+            self.body_pre_processor or self.left_hand_pre_processor or self.right_hand_pre_processor
+        ), "Pre-processors are not initialized."
+
+        # Check if finger data is present in raw_data
+        has_finger_data = "left_fingers" in raw_data and "right_fingers" in raw_data
+
+        if self.body_pre_processor:
+            body_data = self.body_pre_processor(raw_data)
+            # Only process hand data if finger keys are present and preprocessors are available
+            if has_finger_data and self.left_hand_pre_processor and self.right_hand_pre_processor:
+                left_hand_data = self.left_hand_pre_processor(raw_data)
+                right_hand_data = self.right_hand_pre_processor(raw_data)
+                return body_data, left_hand_data, right_hand_data
+            else:
+                return body_data, None, None
+        else:  # only hands
+            if has_finger_data and self.left_hand_pre_processor and self.right_hand_pre_processor:
+                left_hand_data = self.left_hand_pre_processor(raw_data)
+                right_hand_data = self.right_hand_pre_processor(raw_data)
+                return None, left_hand_data, right_hand_data
+            else:
+                # No finger data available, return None for hand data
+                return None, None, None
+
+    def reset(self):
+        if self.body_streamer is not None:
+            self.body_streamer.reset_status()
+        if self.hand_streamer is not None:
+            self.hand_streamer.reset_status()
+
+    def stop_streaming(self):
+        if self.body_streamer:
+            self.body_streamer.stop_streaming()
+        # Only stop hand_streamer if it's a different instance than body_streamer
+        if self.hand_streamer and self.hand_streamer is not self.body_streamer:
+            self.hand_streamer.stop_streaming()

GR00T-WholeBodyControl/decoupled_wbc/control/utils/cv_bridge.py

@@ -0,0 +1,396 @@
+# Software License Agreement (BSD License)
+#
+# Copyright (c) 2011, Willow Garage, Inc.
+# Copyright (c) 2016, Tal Regev.
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above
+#    copyright notice, this list of conditions and the following
+#    disclaimer in the documentation and/or other materials provided
+#    with the distribution.
+#  * Neither the name of Willow Garage, Inc. nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
+# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
+# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+
+import sys
+
+import cv2
+import sensor_msgs.msg
+
+
+def CV_MAT_CNWrap(flags):
+    return (((flags) & ((63) << 3)) >> 3) + 1
+
+
+def CV_MAT_DEPTHWrap(flags):
+    return (flags) & 7
+
+
+_CV_CONVERSIONS = {
+    ("mono8", "rgb8"): cv2.COLOR_GRAY2RGB,
+    ("mono8", "bgr8"): cv2.COLOR_GRAY2BGR,
+    ("mono8", "rgba8"): cv2.COLOR_GRAY2RGBA,
+    ("mono8", "bgra8"): cv2.COLOR_GRAY2BGRA,
+    ("rgb8", "mono8"): cv2.COLOR_RGB2GRAY,
+    ("rgb8", "bgr8"): cv2.COLOR_RGB2BGR,
+    ("rgb8", "rgba8"): cv2.COLOR_RGB2RGBA,
+    ("rgb8", "bgra8"): cv2.COLOR_RGB2BGRA,
+    ("bgr8", "mono8"): cv2.COLOR_BGR2GRAY,
+    ("bgr8", "rgb8"): cv2.COLOR_BGR2RGB,
+    ("bgr8", "rgba8"): cv2.COLOR_BGR2RGBA,
+    ("bgr8", "bgra8"): cv2.COLOR_BGR2BGRA,
+    ("rgba8", "mono8"): cv2.COLOR_RGBA2GRAY,
+    ("rgba8", "rgb8"): cv2.COLOR_RGBA2RGB,
+    ("rgba8", "bgr8"): cv2.COLOR_RGBA2BGR,
+    ("rgba8", "bgra8"): cv2.COLOR_RGBA2BGRA,
+    ("bgra8", "mono8"): cv2.COLOR_BGRA2GRAY,
+    ("bgra8", "rgb8"): cv2.COLOR_BGRA2RGB,
+    ("bgra8", "bgr8"): cv2.COLOR_BGRA2BGR,
+    ("bgra8", "rgba8"): cv2.COLOR_BGRA2RGBA,
+    ("yuv422", "mono8"): cv2.COLOR_YUV2GRAY_UYVY,
+    ("yuv422", "rgb8"): cv2.COLOR_YUV2RGB_UYVY,
+    ("yuv422", "bgr8"): cv2.COLOR_YUV2BGR_UYVY,
+    ("yuv422", "rgba8"): cv2.COLOR_YUV2RGBA_UYVY,
+    ("yuv422", "bgra8"): cv2.COLOR_YUV2BGRA_UYVY,
+    ("bayer_rggb8", "mono8"): cv2.COLOR_BayerBG2GRAY,
+    ("bayer_rggb8", "rgb8"): cv2.COLOR_BayerBG2RGB,
+    ("bayer_rggb8", "bgr8"): cv2.COLOR_BayerBG2BGR,
+    ("bayer_bggr8", "mono8"): cv2.COLOR_BayerRG2GRAY,
+    ("bayer_bggr8", "rgb8"): cv2.COLOR_BayerRG2RGB,
+    ("bayer_bggr8", "bgr8"): cv2.COLOR_BayerRG2BGR,
+    ("bayer_gbrg8", "mono8"): cv2.COLOR_BayerGR2GRAY,
+    ("bayer_gbrg8", "rgb8"): cv2.COLOR_BayerGR2RGB,
+    ("bayer_gbrg8", "bgr8"): cv2.COLOR_BayerGR2BGR,
+    ("bayer_grbg", "mono8"): cv2.COLOR_BayerGB2GRAY,
+    ("bayer_grbg", "rgb8"): cv2.COLOR_BayerGB2RGB,
+    ("bayer_grbg", "bgr8"): cv2.COLOR_BayerGB2BGR,
+}
+
+_CV_TYPES = {
+    "rgb8": cv2.CV_8UC3,
+    "rgba8": cv2.CV_8UC4,
+    "rgb16": cv2.CV_16UC3,
+    "rgba16": cv2.CV_16UC4,
+    "bgr8": cv2.CV_8UC3,
+    "bgra8": cv2.CV_8UC4,
+    "bgr16": cv2.CV_16UC3,
+    "bgra16": cv2.CV_16UC4,
+    "mono8": cv2.CV_8UC1,
+    "mono16": cv2.CV_16UC1,
+    "8UC1": cv2.CV_8UC1,
+    "8UC2": cv2.CV_8UC2,
+    "8UC3": cv2.CV_8UC3,
+    "8UC4": cv2.CV_8UC4,
+    "8SC1": cv2.CV_8SC1,
+    "8SC2": cv2.CV_8SC2,
+    "8SC3": cv2.CV_8SC3,
+    "8SC4": cv2.CV_8SC4,
+    "16UC1": cv2.CV_8UC1,
+    "16UC2": cv2.CV_8UC2,
+    "16UC3": cv2.CV_8UC3,
+    "16UC4": cv2.CV_8UC4,
+    "16SC1": cv2.CV_16SC1,
+    "16SC2": cv2.CV_16SC2,
+    "16SC3": cv2.CV_16SC3,
+    "16SC4": cv2.CV_16SC4,
+    "32SC1": cv2.CV_32SC1,
+    "32SC2": cv2.CV_32SC2,
+    "32SC3": cv2.CV_32SC3,
+    "32SC4": cv2.CV_32SC4,
+    "32FC1": cv2.CV_32FC1,
+    "32FC2": cv2.CV_32FC2,
+    "32FC3": cv2.CV_32FC3,
+    "32FC4": cv2.CV_32FC4,
+    "64FC1": cv2.CV_64FC1,
+    "64FC2": cv2.CV_64FC2,
+    "64FC3": cv2.CV_64FC3,
+    "64FC4": cv2.CV_64FC4,
+    "bayer_rggb8": cv2.CV_8UC1,
+    "bayer_bggr8": cv2.CV_8UC1,
+    "bayer_gbrg8": cv2.CV_8UC1,
+    "bayer_grbg8": cv2.CV_8UC1,
+    "bayer_rggb16": cv2.CV_16UC1,
+    "bayer_bggr16": cv2.CV_16UC1,
+    "bayer_gbrg16": cv2.CV_16UC1,
+    "bayer_grbg16": cv2.CV_16UC1,
+}
+
+
+def cvtColor2(img, encoding_in, encoding_out):
+    if encoding_in == encoding_out:
+        return img
+
+    conversion = _CV_CONVERSIONS[(encoding_in, encoding_out)]
+    # depth conversion is not yet implemented
+    return cv2.cvtColor(img, conversion)
+
+
+def getCvType(encoding):
+    return _CV_TYPES[encoding]
+
+
+class CvBridgeError(TypeError):
+    """
+    This is the error raised by :class:`cv_bridge.CvBridge` methods when they fail.
+    """
+
+    pass
+
+
+class CvBridge(object):
+    """
+    The CvBridge is an object that converts between OpenCV Images and ROS Image messages.
+
+       .. doctest::
+           :options: -ELLIPSIS, +NORMALIZE_WHITESPACE
+
+           >>> import cv2
+           >>> import numpy as np
+           >>> from cv_bridge import CvBridge
+           >>> br = CvBridge()
+           >>> dtype, n_channels = br.encoding_as_cvtype2('8UC3')
+           >>> im = np.ndarray(shape=(480, 640, n_channels), dtype=dtype)
+           >>> msg = br.cv2_to_imgmsg(im)  # Convert the image to a message
+           >>> im2 = br.imgmsg_to_cv2(msg) # Convert the message to a new image
+           >>> cmprsmsg = br.cv2_to_compressed_imgmsg(im)  # Convert the image to a compress message
+           >>> im22 = br.compressed_imgmsg_to_cv2(msg) # Convert the compress message to a new image
+           >>> cv2.imwrite("this_was_a_message_briefly.png", im2)
+
+    """
+
+    def __init__(self):
+        import cv2
+
+        self.cvtype_to_name = {}
+        self.cvdepth_to_numpy_depth = {
+            cv2.CV_8U: "uint8",
+            cv2.CV_8S: "int8",
+            cv2.CV_16U: "uint16",
+            cv2.CV_16S: "int16",
+            cv2.CV_32S: "int32",
+            cv2.CV_32F: "float32",
+            cv2.CV_64F: "float64",
+        }
+
+        for t in ["8U", "8S", "16U", "16S", "32S", "32F", "64F"]:
+            for c in [1, 2, 3, 4]:
+                nm = "%sC%d" % (t, c)
+                self.cvtype_to_name[getattr(cv2, "CV_%s" % nm)] = nm
+
+        self.numpy_type_to_cvtype = {
+            "uint8": "8U",
+            "int8": "8S",
+            "uint16": "16U",
+            "int16": "16S",
+            "int32": "32S",
+            "float32": "32F",
+            "float64": "64F",
+        }
+        self.numpy_type_to_cvtype.update(
+            dict((v, k) for (k, v) in self.numpy_type_to_cvtype.items())
+        )
+
+    def dtype_with_channels_to_cvtype2(self, dtype, n_channels):
+        return "%sC%d" % (self.numpy_type_to_cvtype[dtype.name], n_channels)
+
+    def cvtype2_to_dtype_with_channels(self, cvtype):
+        return self.cvdepth_to_numpy_depth[CV_MAT_DEPTHWrap(cvtype)], CV_MAT_CNWrap(cvtype)
+
+    def encoding_to_cvtype2(self, encoding):
+        try:
+            return getCvType(encoding)
+        except RuntimeError as e:
+            raise CvBridgeError(e)
+
+    def encoding_to_dtype_with_channels(self, encoding):
+        return self.cvtype2_to_dtype_with_channels(self.encoding_to_cvtype2(encoding))
+
+    def compressed_imgmsg_to_cv2(self, cmprs_img_msg, desired_encoding="passthrough"):
+        """
+        Convert a sensor_msgs::CompressedImage message to an OpenCV :cpp:type:`cv::Mat`.
+
+        :param cmprs_img_msg:   A :cpp:type:`sensor_msgs::CompressedImage` message
+        :param desired_encoding:  The encoding of the image data, one of the following strings:
+
+           * ``"passthrough"``
+           * one of the standard strings in sensor_msgs/image_encodings.h
+
+        :rtype: :cpp:type:`cv::Mat`
+        :raises CvBridgeError: when conversion is not possible.
+
+        If desired_encoding is ``"passthrough"``, then the returned image has the same format as img_msg.
+        Otherwise desired_encoding must be one of the standard image encodings
+
+        This function returns an OpenCV :cpp:type:`cv::Mat` message on success, or raises
+        :exc:`cv_bridge.CvBridgeError` on failure.
+
+        If the image only has one channel, the shape has size 2 (width and height)
+        """
+        import cv2
+        import numpy as np
+
+        str_msg = cmprs_img_msg.data
+        buf = np.ndarray(shape=(1, len(str_msg)), dtype=np.uint8, buffer=cmprs_img_msg.data)
+        im = cv2.imdecode(buf, cv2.IMREAD_ANYCOLOR)
+
+        if desired_encoding == "passthrough":
+            return im
+
+        try:
+            res = cvtColor2(im, "bgr8", desired_encoding)
+        except RuntimeError as e:
+            raise CvBridgeError(e)
+
+        return res
+
+    def imgmsg_to_cv2(self, img_msg, desired_encoding="passthrough"):
+        """
+        Convert a sensor_msgs::Image message to an OpenCV :cpp:type:`cv::Mat`.
+
+        :param img_msg:   A :cpp:type:`sensor_msgs::Image` message
+        :param desired_encoding:  The encoding of the image data, one of the following strings:
+
+           * ``"passthrough"``
+           * one of the standard strings in sensor_msgs/image_encodings.h
+
+        :rtype: :cpp:type:`cv::Mat`
+        :raises CvBridgeError: when conversion is not possible.
+
+        If desired_encoding is ``"passthrough"``, then the returned image has the same format as img_msg.
+        Otherwise desired_encoding must be one of the standard image encodings
+
+        This function returns an OpenCV :cpp:type:`cv::Mat` message on success, or raises
+        :exc:`cv_bridge.CvBridgeError` on failure.
+
+        If the image only has one channel, the shape has size 2 (width and height)
+        """
+        import numpy as np
+
+        dtype, n_channels = self.encoding_to_dtype_with_channels(img_msg.encoding)
+        dtype = np.dtype(dtype)
+        dtype = dtype.newbyteorder(">" if img_msg.is_bigendian else "<")
+        if n_channels == 1:
+            im = np.ndarray(shape=(img_msg.height, img_msg.width), dtype=dtype, buffer=img_msg.data)
+        else:
+            im = np.ndarray(
+                shape=(img_msg.height, img_msg.width, n_channels), dtype=dtype, buffer=img_msg.data
+            )
+        # If the byt order is different between the message and the system.
+        if img_msg.is_bigendian == (sys.byteorder == "little"):
+            im = im.byteswap().newbyteorder()
+
+        if desired_encoding == "passthrough":
+            return im
+
+        try:
+            res = cvtColor2(im, img_msg.encoding, desired_encoding)
+        except RuntimeError as e:
+            raise CvBridgeError(e)
+
+        return res
+
+    def cv2_to_compressed_imgmsg(self, cvim, dst_format="jpg"):
+        """
+        Convert an OpenCV :cpp:type:`cv::Mat` type to a ROS sensor_msgs::CompressedImage message.
+
+        :param cvim:      An OpenCV :cpp:type:`cv::Mat`
+        :param dst_format:  The format of the image data, one of the following strings:
+
+           * from http://docs.opencv.org/2.4/modules/highgui/doc/reading_and_writing_images_and_video.html
+           * from http://docs.opencv.org/2.4/modules/highgui/doc/reading_and_writing_images_and_video.html#Mat
+           imread(const string& filename, int flags)
+           * bmp, dib
+           * jpeg, jpg, jpe
+           * jp2
+           * png
+           * pbm, pgm, ppm
+           * sr, ras
+           * tiff, tif
+
+        :rtype:           A sensor_msgs.msg.CompressedImage message
+        :raises CvBridgeError: when the ``cvim`` has a type that is incompatible with ``format``
+
+
+        This function returns a sensor_msgs::Image message on success, or raises
+        :exc:`cv_bridge.CvBridgeError` on failure.
+        """
+        import cv2
+        import numpy as np
+
+        if not isinstance(cvim, (np.ndarray, np.generic)):
+            raise TypeError("Your input type is not a numpy array")
+        cmprs_img_msg = sensor_msgs.msg.CompressedImage()
+        cmprs_img_msg.format = dst_format
+        ext_format = "." + dst_format
+        try:
+            cmprs_img_msg.data = np.array(cv2.imencode(ext_format, cvim)[1]).tostring()
+        except RuntimeError as e:
+            raise CvBridgeError(e)
+
+        return cmprs_img_msg
+
+    def cv2_to_imgmsg(self, cvim, encoding="passthrough"):
+        """
+        Convert an OpenCV :cpp:type:`cv::Mat` type to a ROS sensor_msgs::Image message.
+
+        :param cvim:      An OpenCV :cpp:type:`cv::Mat`
+        :param encoding:  The encoding of the image data, one of the following strings:
+
+           * ``"passthrough"``
+           * one of the standard strings in sensor_msgs/image_encodings.h
+
+        :rtype:           A sensor_msgs.msg.Image message
+        :raises CvBridgeError: when the ``cvim`` has a type that is incompatible with ``encoding``
+
+        If encoding is ``"passthrough"``, then the message has the same encoding as the image's OpenCV type.
+        Otherwise desired_encoding must be one of the standard image encodings
+
+        This function returns a sensor_msgs::Image message on success, or raises
+        :exc:`cv_bridge.CvBridgeError`on failure.
+        """
+        import numpy as np
+
+        if not isinstance(cvim, (np.ndarray, np.generic)):
+            raise TypeError("Your input type is not a numpy array")
+        img_msg = sensor_msgs.msg.Image()
+        img_msg.height = cvim.shape[0]
+        img_msg.width = cvim.shape[1]
+        if len(cvim.shape) < 3:
+            cv_type = self.dtype_with_channels_to_cvtype2(cvim.dtype, 1)
+        else:
+            cv_type = self.dtype_with_channels_to_cvtype2(cvim.dtype, cvim.shape[2])
+        if encoding == "passthrough":
+            img_msg.encoding = cv_type
+        else:
+            img_msg.encoding = encoding
+            # # Verify that the supplied encoding is compatible with the type of the OpenCV image
+            # if self.cvtype_to_name[self.encoding_to_cvtype2(encoding)] != cv_type:
+            #     raise CvBridgeError(
+            #         "encoding specified as %s, but image has incompatible type %s"
+            #         % (encoding, cv_type)
+            #     )
+        if cvim.dtype.byteorder == ">":
+            img_msg.is_bigendian = True
+        img_msg.data = cvim.tostring()
+        img_msg.step = len(img_msg.data) // img_msg.height
+
+        return img_msg

GR00T-WholeBodyControl/decoupled_wbc/control/utils/episode_state.py

@@ -0,0 +1,32 @@
+class EpisodeState:
+    """Episode state controller for data collection.
+
+    Manages the state transitions for episode recording:
+    - IDLE: Not recording
+    - RECORDING: Currently recording data
+    - NEED_TO_SAVE: Recording stopped, waiting to save
+    """
+
+    def __init__(self):
+        self.RECORDING = "recording"
+        self.IDLE = "idle"
+        self.NEED_TO_SAVE = "need_to_save"
+
+        self.state = self.IDLE
+
+    def change_state(self):
+        """Cycle through states: IDLE -> RECORDING -> NEED_TO_SAVE -> IDLE."""
+        if self.state == self.IDLE:
+            self.state = self.RECORDING
+        elif self.state == self.RECORDING:
+            self.state = self.NEED_TO_SAVE
+        elif self.state == self.NEED_TO_SAVE:
+            self.state = self.IDLE
+
+    def reset_state(self):
+        """Reset to IDLE state."""
+        self.state = self.IDLE
+
+    def get_state(self):
+        """Get current state."""
+        return self.state

GR00T-WholeBodyControl/decoupled_wbc/control/utils/gear_wbc_utils.py

@@ -0,0 +1,100 @@
+import os
+
+import numpy as np
+import yaml
+
+
+def load_config(config_path):
+    """Load and process the YAML configuration file"""
+    with open(config_path, "r") as f:
+        config = yaml.safe_load(f)
+
+    # Set the path to the LEGGED_GYM_ROOT_DIR using relative path
+    current_file_dir = os.path.dirname(os.path.abspath(config_path))
+    LEGGED_GYM_ROOT_DIR = os.path.join(current_file_dir, "..", "GearWbcRL", "legged_gym")
+    LEGGED_GYM_ROOT_DIR = os.path.abspath(LEGGED_GYM_ROOT_DIR)
+
+    # Process paths with LEGGED_GYM_ROOT_DIR
+    for path_key in ["policy_path", "xml_path", "onnx_policy_path"]:
+        if path_key in config:
+            config[path_key] = config[path_key].format(LEGGED_GYM_ROOT_DIR=LEGGED_GYM_ROOT_DIR)
+
+    # Convert lists to numpy arrays where needed
+    array_keys = ["kps", "kds", "default_angles", "cmd_scale", "cmd_init"]
+    for key in array_keys:
+        if key in config:
+            config[key] = np.array(config[key], dtype=np.float32)
+
+    return config, LEGGED_GYM_ROOT_DIR
+
+
+def pd_control(target_q, q, kp, target_dq, dq, kd):
+    """Calculates torques from position commands"""
+    return (target_q - q) * kp + (target_dq - dq) * kd
+
+
+def quat_rotate_inverse(q, v):
+    """Rotate vector v by the inverse of quaternion q"""
+    w = q[..., 0]
+    x = q[..., 1]
+    y = q[..., 2]
+    z = q[..., 3]
+
+    q_conj = np.array([w, -x, -y, -z])
+
+    return np.array(
+        [
+            v[0] * (q_conj[0] ** 2 + q_conj[1] ** 2 - q_conj[2] ** 2 - q_conj[3] ** 2)
+            + v[1] * 2 * (q_conj[1] * q_conj[2] - q_conj[0] * q_conj[3])
+            + v[2] * 2 * (q_conj[1] * q_conj[3] + q_conj[0] * q_conj[2]),
+            v[0] * 2 * (q_conj[1] * q_conj[2] + q_conj[0] * q_conj[3])
+            + v[1] * (q_conj[0] ** 2 - q_conj[1] ** 2 + q_conj[2] ** 2 - q_conj[3] ** 2)
+            + v[2] * 2 * (q_conj[2] * q_conj[3] - q_conj[0] * q_conj[1]),
+            v[0] * 2 * (q_conj[1] * q_conj[3] - q_conj[0] * q_conj[2])
+            + v[1] * 2 * (q_conj[2] * q_conj[3] + q_conj[0] * q_conj[1])
+            + v[2] * (q_conj[0] ** 2 - q_conj[1] ** 2 - q_conj[2] ** 2 + q_conj[3] ** 2),
+        ]
+    )
+
+
+def get_gravity_orientation(quat):
+    """Get gravity vector in body frame"""
+    gravity_vec = np.array([0.0, 0.0, -1.0])
+    return quat_rotate_inverse(quat, gravity_vec)
+
+
+def compute_observation(d, config, action, cmd, height_cmd, n_joints):
+    """Compute the observation vector from current state"""
+    # Get state from MuJoCo
+    qj = d.qpos[7 : 7 + n_joints].copy()
+    dqj = d.qvel[6 : 6 + n_joints].copy()
+    quat = d.qpos[3:7].copy()
+    omega = d.qvel[3:6].copy()
+
+    # Handle default angles padding
+    if len(config["default_angles"]) < n_joints:
+        padded_defaults = np.zeros(n_joints, dtype=np.float32)
+        padded_defaults[: len(config["default_angles"])] = config["default_angles"]
+    else:
+        padded_defaults = config["default_angles"][:n_joints]
+
+    # Scale the values
+    qj_scaled = (qj - padded_defaults) * config["dof_pos_scale"]
+    dqj_scaled = dqj * config["dof_vel_scale"]
+    gravity_orientation = get_gravity_orientation(quat)
+    omega_scaled = omega * config["ang_vel_scale"]
+
+    # Calculate single observation dimension
+    single_obs_dim = 3 + 1 + 3 + 3 + n_joints + n_joints + 12
+
+    # Create single observation
+    single_obs = np.zeros(single_obs_dim, dtype=np.float32)
+    single_obs[0:3] = cmd[:3] * config["cmd_scale"]
+    single_obs[3:4] = np.array([height_cmd])
+    single_obs[4:7] = omega_scaled
+    single_obs[7:10] = gravity_orientation
+    single_obs[10 : 10 + n_joints] = qj_scaled
+    single_obs[10 + n_joints : 10 + 2 * n_joints] = dqj_scaled
+    single_obs[10 + 2 * n_joints : 10 + 2 * n_joints + 12] = action
+
+    return single_obs, single_obs_dim

GR00T-WholeBodyControl/decoupled_wbc/control/utils/keyboard_dispatcher.py

@@ -0,0 +1,255 @@
+import os
+import subprocess
+import sys
+import threading
+
+import rclpy
+from sshkeyboard import listen_keyboard, stop_listening
+from std_msgs.msg import String as RosStringMsg
+
+from decoupled_wbc.control.main.constants import KEYBOARD_INPUT_TOPIC
+
+# Global variable to store original terminal attributes
+_original_terminal_attrs = None
+
+
+def save_terminal_state():
+    """Save the current terminal state."""
+    global _original_terminal_attrs
+    try:
+        import termios
+
+        fd = sys.stdin.fileno()
+        _original_terminal_attrs = termios.tcgetattr(fd)
+    except (ImportError, OSError, termios.error):
+        _original_terminal_attrs = None
+
+
+def restore_terminal():
+    """Restore terminal to original state."""
+    global _original_terminal_attrs
+    try:
+        import termios
+
+        if _original_terminal_attrs is not None:
+            fd = sys.stdin.fileno()
+            termios.tcsetattr(fd, termios.TCSANOW, _original_terminal_attrs)
+            return
+    except (ImportError, OSError, termios.error):
+        pass
+
+    # Fallback for non-Unix systems or if termios fails
+    try:
+        if os.name == "posix":
+            os.system("stty sane")
+    except OSError:
+        pass
+
+
+class ROSKeyboardDispatcher:
+    """ROS-based keyboard dispatcher that receives keyboard events via ROS topics."""
+
+    def __init__(self):
+        self.listeners = []
+        self._active = False
+        assert rclpy.ok(), "Expected ROS2 to be initialized in this process..."
+        executor = rclpy.get_global_executor()
+        self.node = executor.get_nodes()[0]
+        print("creating keyboard input subscriber...")
+        self.subscription = self.node.create_subscription(
+            RosStringMsg, KEYBOARD_INPUT_TOPIC, self._callback, 10
+        )
+
+    def register(self, listener):
+        if not hasattr(listener, "handle_keyboard_button"):
+            raise NotImplementedError("handle_keyboard_button is not implemented")
+        self.listeners.append(listener)
+
+    def start(self):
+        """Start the ROS keyboard dispatcher."""
+        self._active = True
+        print("ROS keyboard dispatcher started")
+
+    def stop(self):
+        """Stop the ROS keyboard dispatcher and cleanup."""
+        if self._active:
+            self._active = False
+            # Clean up subscription
+            if hasattr(self, "subscription"):
+                self.node.destroy_subscription(self.subscription)
+            print("ROS keyboard dispatcher stopped")
+
+    def _callback(self, msg: RosStringMsg):
+        if self._active:
+            for listener in self.listeners:
+                listener.handle_keyboard_button(msg.data)
+
+    def __del__(self):
+        """Cleanup when object is destroyed."""
+        self.stop()
+
+
+class KeyboardDispatcher:
+    def __init__(self):
+        self.listeners = []
+        self._listening_thread = None
+        self._stop_event = threading.Event()
+        self._key = None
+
+    def register(self, listener):
+        # raise if handle_keyboard_button is not implemented
+        # TODO(YL): let listener be a Callable instead of a class
+        if not hasattr(listener, "handle_keyboard_button"):
+            raise NotImplementedError("handle_keyboard_button is not implemented")
+        self.listeners.append(listener)
+
+    def handle_key(self, key):
+        # Check if we should stop
+        if self._stop_event.is_set():
+            stop_listening()
+            return
+
+        for listener in self.listeners:
+            listener.handle_keyboard_button(key)
+
+    def start_listening(self):
+        try:
+            save_terminal_state()  # Save original terminal state before listening
+            listen_keyboard(
+                on_press=self.handle_key,
+                delay_second_char=0.1,
+                delay_other_chars=0.05,
+                sleep=0.01,
+            )
+        except Exception as e:
+            print(f"Keyboard listener stopped: {e}")
+        finally:
+            # Ensure terminal is restored even if an exception occurs
+            self._restore_terminal()
+
+    def start(self):
+        self._listening_thread = threading.Thread(target=self.start_listening, daemon=True)
+        self._listening_thread.start()
+
+    def stop(self):
+        """Stop the keyboard listener and restore terminal settings."""
+        if self._listening_thread and self._listening_thread.is_alive():
+            self._stop_event.set()
+            # Force stop_listening to be called
+            try:
+                stop_listening()
+            except Exception:
+                pass
+            # Wait a bit for the thread to finish
+            self._listening_thread.join(timeout=0.5)
+            # Restore terminal settings
+            self._restore_terminal()
+
+    def _restore_terminal(self):
+        """Restore terminal to a sane state."""
+        restore_terminal()
+
+    def __del__(self):
+        """Cleanup when object is destroyed."""
+        self.stop()
+
+
+KEYBOARD_LISTENER_TOPIC_NAME = "/Gr00tKeyboardListener"
+
+
+class KeyboardListener:
+    def __init__(self):
+        self.key = None
+
+    def handle_keyboard_button(self, key):
+        self.key = key
+
+    def pop_key(self):
+        key = self.key
+        self.key = None
+        return key
+
+
+class KeyboardListenerPublisher:
+    def __init__(self, topic_name: str = KEYBOARD_LISTENER_TOPIC_NAME):
+        """
+        Initialize keyboard listener for remote teleop with simplified interface.
+
+        Args:
+            remote_system: RemoteSystem instance
+            control_channel_name: Name of the control channel
+        """
+        assert rclpy.ok(), "Expected ROS2 to be initialized in this process..."
+        executor = rclpy.get_global_executor()
+        self.node = executor.get_nodes()[0]
+        self.publisher = self.node.create_publisher(RosStringMsg, topic_name, 1)
+
+    def handle_keyboard_button(self, key):
+        self.publisher.publish(RosStringMsg(data=key))
+
+
+class KeyboardListenerSubscriber:
+    def __init__(
+        self,
+        topic_name: str = KEYBOARD_LISTENER_TOPIC_NAME,
+        node_name: str = "keyboard_listener_subscriber",
+    ):
+        assert rclpy.ok(), "Expected ROS2 to be initialized in this process..."
+        executor = rclpy.get_global_executor()
+        nodes = executor.get_nodes()
+        if nodes:
+            self.node = nodes[0]
+            self._create_node = False
+        else:
+            self.node = rclpy.create_node("KeyboardListenerSubscriber")
+            executor.add_node(self.node)
+            self._create_node = True
+        self.subscriber = self.node.create_subscription(RosStringMsg, topic_name, self._callback, 1)
+        self._data = None
+
+    def _callback(self, msg: RosStringMsg):
+        self._data = msg.data
+
+    def read_msg(self):
+        data = self._data
+        self._data = None
+        return data
+
+
+class KeyboardEStop:
+    def __init__(self):
+        """Initialize KeyboardEStop with automatic tmux cleanup detection."""
+        # Automatically create tmux cleanup if in deployment mode
+        self.cleanup_callback = self._create_tmux_cleanup_callback()
+
+    def _create_tmux_cleanup_callback(self):
+        """Create a cleanup callback that kills the tmux session if running in deployment mode."""
+        tmux_session = os.environ.get("DECOUPLED_WBC_TMUX_SESSION")
+
+        def cleanup_callback():
+            if tmux_session:
+                print(f"Emergency stop: Killing tmux session '{tmux_session}'...")
+                try:
+                    subprocess.run(["tmux", "kill-session", "-t", tmux_session], timeout=5)
+                    print("Tmux session terminated successfully.")
+                except subprocess.TimeoutExpired:
+                    print("Warning: Tmux session termination timed out, forcing kill...")
+                    try:
+                        subprocess.run(["tmux", "kill-session", "-t", tmux_session, "-9"])
+                    except Exception:
+                        pass
+                except Exception as e:
+                    print(f"Warning: Error during tmux cleanup: {e}")
+                    # If tmux cleanup fails, fallback to immediate exit
+                    restore_terminal()
+                    os._exit(1)
+            else:
+                print("Emergency stop: No tmux session, exiting normally...")
+                sys.exit(1)
+
+        return cleanup_callback
+
+    def handle_keyboard_button(self, key):
+        if key == "`":
+            print("Emergency stop triggered - running cleanup...")
+            self.cleanup_callback()

GR00T-WholeBodyControl/decoupled_wbc/control/utils/network_utils.py

@@ -0,0 +1,137 @@
+#!/usr/bin/env python3
+"""
+Simple network interface utilities
+"""
+
+import platform
+import re
+import subprocess
+
+
+def get_network_interfaces():
+    """Get network interfaces with their IP addresses"""
+    try:
+        result = subprocess.run(
+            ["/sbin/ip", "addr", "show"], capture_output=True, text=True, check=True
+        )
+        return _parse_ip_output(result.stdout)
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        try:
+            result = subprocess.run(["ifconfig"], capture_output=True, text=True, check=True)
+            return _parse_ifconfig_output(result.stdout)
+        except (subprocess.CalledProcessError, FileNotFoundError):
+            return {}
+
+
+def _parse_ip_output(output):
+    """Parse 'ip addr' command output"""
+    interfaces = {}
+    current_interface = None
+
+    for line in output.split("\n"):
+        interface_match = re.match(r"^\d+:\s+(\w+):", line)
+        if interface_match:
+            current_interface = interface_match.group(1)
+            interfaces[current_interface] = []
+
+        ip_match = re.search(r"inet\s+(\d+\.\d+\.\d+\.\d+)", line)
+        if ip_match and current_interface:
+            interfaces[current_interface].append(ip_match.group(1))
+
+    return interfaces
+
+
+def _parse_ifconfig_output(output):
+    """Parse 'ifconfig' command output"""
+    interfaces = {}
+    current_interface = None
+
+    for line in output.split("\n"):
+        interface_match = re.match(r"^(\w+):", line)
+        if interface_match:
+            current_interface = interface_match.group(1)
+            interfaces[current_interface] = []
+
+        ip_match = re.search(r"inet\s+(\d+\.\d+\.\d+\.\d+)", line)
+        if ip_match and current_interface:
+            interfaces[current_interface].append(ip_match.group(1))
+
+    return interfaces
+
+
+def find_interface_by_ip(target_ip):
+    """Find interface name for given IP address"""
+    interfaces = get_network_interfaces()
+    for interface, ip_list in interfaces.items():
+        if target_ip in ip_list:
+            return interface
+    return None
+
+
+def resolve_interface(interface: str) -> tuple[str, str]:
+    """
+    Resolve interface parameter to actual network interface name and environment type
+
+    Args:
+        interface: "sim", "real", or direct interface name or IP address
+
+    Returns:
+        tuple: (interface_name, env_type) where env_type is "sim" or "real"
+    """
+    # Check if interface is an IP address
+    if re.match(r"^\d+\.\d+\.\d+\.\d+$", interface):
+        if interface == "127.0.0.1":
+            return interface, "sim"
+        else:
+            return interface, "real"
+
+    if interface == "sim":
+        lo_interface = find_interface_by_ip("127.0.0.1")
+        if lo_interface:
+            # macOS uses lo0 instead of lo
+            if platform.system() == "Darwin" and lo_interface == "lo":
+                return "lo0", "sim"
+            return lo_interface, "sim"
+        return ("lo0" if platform.system() == "Darwin" else "lo"), "sim"
+
+    elif interface == "real":
+        interfaces = get_network_interfaces()
+        for iface, ip_list in interfaces.items():
+            for ip in ip_list:
+                if ip.startswith("192.168.123."):
+                    return iface, "real"
+        return interface, "real"  # fallback
+
+    else:
+        # Direct interface name - check if it has 127.0.0.1 to determine env_type
+        interfaces = get_network_interfaces()
+        if interface in interfaces:
+            for ip in interfaces[interface]:
+                if ip == "127.0.0.1":
+                    return interface, "sim"
+
+        # macOS lo interface handling
+        if platform.system() == "Darwin" and interface == "lo":
+            return "lo0", "sim"
+
+        # Default to real for unknown interfaces
+        return interface, "real"
+
+
+if __name__ == "__main__":
+    interfaces = get_network_interfaces()
+
+    if not interfaces:
+        print("No network interfaces found")
+        exit(1)
+
+    # Show all interfaces
+    print("Network interfaces:")
+    for interface, ip_list in interfaces.items():
+        print(f"  {interface}: {', '.join(ip_list)}")
+
+    # Test resolve_interface function
+    print("\nTesting resolve_interface:")
+    for test_interface in ["sim", "real", "lo", "127.0.0.1"]:
+        interface_name, env_type = resolve_interface(test_interface)
+        print(f"  {test_interface} -> {interface_name} ({env_type})")

GR00T-WholeBodyControl/decoupled_wbc/control/utils/ros_utils.py

@@ -0,0 +1,201 @@
+import base64
+import signal
+import threading
+from typing import Optional
+
+import msgpack
+import msgpack_numpy as mnp
+import rclpy
+from rclpy.executors import SingleThreadedExecutor
+from rclpy.node import Node
+from sensor_msgs.msg import Image
+from std_msgs.msg import ByteMultiArray
+from std_srvs.srv import Trigger
+
+_signal_registered = False
+
+
+def register_keyboard_interrupt_handler():
+    """
+    Register a signal handler for SIGINT (Ctrl+C) and SIGTERM that raises KeyboardInterrupt.
+    This ensures consistent exception handling across different termination signals.
+    """
+    global _signal_registered
+    if not _signal_registered:
+
+        def signal_handler(signum, frame):
+            raise KeyboardInterrupt
+
+        signal.signal(signal.SIGINT, signal_handler)
+        signal.signal(signal.SIGTERM, signal_handler)
+        _signal_registered = True
+
+
+class ROSManager:
+    """
+    Manages the ROS2 node and executor.
+
+    Usage example:
+    ```python
+    def main():
+        ros_manager = ROSManager()
+        node = ros_manager.node
+
+        try:
+            while ros_manager.ok():
+                time.sleep(0.1)
+        except ros_manager.exceptions() as e:
+            print(f"ROSManager interrupted by user: {e}")
+        finally:
+            ros_manager.shutdown()
+    ```
+    """
+
+    def __init__(self, node_name: str = "ros_manager"):
+        if not rclpy.ok():
+            rclpy.init()
+            self.node = rclpy.create_node(node_name)
+            self.thread = threading.Thread(target=rclpy.spin, args=(self.node,), daemon=True)
+            self.thread.start()
+        else:
+            executor = rclpy.get_global_executor()
+            if len(executor.get_nodes()) > 0:
+                self.node = executor.get_nodes()[0]
+            else:
+                self.node = rclpy.create_node(node_name)
+
+        register_keyboard_interrupt_handler()
+
+    @staticmethod
+    def ok():
+        return rclpy.ok()
+
+    @staticmethod
+    def shutdown():
+        if rclpy.ok():
+            rclpy.shutdown()
+
+    @staticmethod
+    def exceptions():
+        return (rclpy.exceptions.ROSInterruptException, KeyboardInterrupt)
+
+
+class ROSMsgPublisher:
+    """
+    Publishes any serializable dict to a topic.
+    """
+
+    def __init__(self, topic_name: str):
+        ros_manager = ROSManager()
+        self.node = ros_manager.node
+        self.publisher = self.node.create_publisher(ByteMultiArray, topic_name, 1)
+
+    def publish(self, msg: dict):
+        payload = msgpack.packb(msg, default=mnp.encode)
+        payload = tuple(bytes([a]) for a in payload)
+        msg = ByteMultiArray()
+        msg.data = payload
+        self.publisher.publish(msg)
+
+
+class ROSMsgSubscriber:
+    """
+    Subscribes to any topics published by a ROSMsgPublisher.
+    """
+
+    def __init__(self, topic_name: str):
+        ros_manager = ROSManager()
+        self.node = ros_manager.node
+        self._msg = None
+        self.subscription = self.node.create_subscription(
+            ByteMultiArray, topic_name, self._callback, 1
+        )
+
+    def _callback(self, msg: ByteMultiArray):
+        self._msg = msg
+
+    def get_msg(self) -> Optional[dict]:
+        msg = self._msg
+        if msg is None:
+            return None
+        self._msg = None
+        return msgpack.unpackb(bytes([ab for a in msg.data for ab in a]), object_hook=mnp.decode)
+
+
+class ROSImgMsgSubscriber:
+    """
+    Subscribes to an `Image` topic and returns the image as a numpy array and timestamp.
+    """
+
+    def __init__(self, topic_name: str):
+        ros_manager = ROSManager()
+        self.node = ros_manager.node
+        self._msg = None
+        self.subscription = self.node.create_subscription(Image, topic_name, self._callback, 1)
+
+        from decoupled_wbc.control.utils.cv_bridge import CvBridge
+
+        self.bridge = CvBridge()
+
+    def _callback(self, msg: Image):
+        self._msg = msg
+
+    def get_image(self) -> Optional[dict]:
+        """
+        Returns the image as a numpy array and the timestamp.
+        """
+
+        msg = self._msg
+        if msg is None:
+            return None
+        return {
+            "image": self.bridge.imgmsg_to_cv2(msg),
+            "timestamp": msg.header.stamp.sec + msg.header.stamp.nanosec * 1e-9,
+        }
+
+
+class ROSServiceServer:
+    """
+    Generic ROS2 Service server that stores and serves a config dict.
+    """
+
+    def __init__(self, service_name: str, config: dict):
+        ros_manager = ROSManager()
+        self.node = ros_manager.node
+        packed = msgpack.packb(config, default=mnp.encode)
+        self.message = base64.b64encode(packed).decode("ascii")
+        self.server = self.node.create_service(Trigger, service_name, self._callback)
+
+    def _callback(self, request, response):
+        try:
+            response.success = True
+            response.message = self.message
+            print("Sending encoded message of length:", len(response.message))
+        except Exception as e:
+            response.success = False
+            response.message = str(e)
+        return response
+
+
+class ROSServiceClient(Node):
+
+    def __init__(self, service_name: str, node_name: str = "service_client"):
+        super().__init__(node_name)
+        self.cli = self.create_client(Trigger, service_name)
+        while not self.cli.wait_for_service(timeout_sec=1.0):
+            self.get_logger().info("service not available, waiting again...")
+        self.req = Trigger.Request()
+
+    def get_config(self):
+        future = self.cli.call_async(self.req)
+        executor = SingleThreadedExecutor()
+        executor.add_node(self)
+        executor.spin_until_future_complete(future, timeout_sec=1.0)
+        executor.remove_node(self)
+        executor.shutdown()
+        result = future.result()
+        if result.success:
+            decoded = base64.b64decode(result.message.encode("ascii"))
+            return msgpack.unpackb(decoded, object_hook=mnp.decode)
+        else:
+            raise RuntimeError(f"Service call failed: {result.message}")

GR00T-WholeBodyControl/decoupled_wbc/control/utils/run_real_checklist.py

@@ -0,0 +1,121 @@
+#!/usr/bin/env python3
+
+import sys
+
+
+def check_real_deployment(extra_args):
+    """Check if this is a real robot deployment."""
+    is_real_deployment = False
+
+    # Check if interface argument is provided and not 'lo' or 'lo0'
+    for i, arg in enumerate(extra_args):
+        if arg == "--interface":
+            # Get the next argument (interface value)
+            if i + 1 < len(extra_args):
+                interface_value = extra_args[i + 1]
+                if interface_value not in ["lo", "lo0"]:
+                    is_real_deployment = True
+                    print(f"Real deployment detected: interface = {interface_value}")
+                    break
+                else:
+                    print(f"Simulation deployment detected: interface = {interface_value}")
+
+    # If no interface specified, assume simulation (default is 'lo' in deploy_g1.py)
+    if not is_real_deployment:
+        print("No interface specified - assuming simulation (default interface = lo)")
+
+    return is_real_deployment
+
+
+def show_deployment_checklist():
+    """Show deployment checklist and get confirmation."""
+    checklist_content = """═══════════════════════════════════════════════════════════════════════════════
+                          G1 ROBOT DEPLOYMENT CHECKLIST
+═══════════════════════════════════════════════════════════════════════════════
+
+⚠️  SAFETY VERIFICATION - Complete ALL checks before deployment
+
+PRE-DEPLOYMENT CHECKLIST:
+
+□  Sim2Sim Verification
+   Test in simulation first with interface set to 'sim' before real deployment
+
+□  Camera System Check
+   Test real camera with simulation environment before full deployment
+
+□  State Reading Validation
+   • Disable action queue
+   • Verify sensor readings (IMU, joints, fingers)
+   • Use rerun for visualization
+   • Contact: Dennis Da (xda@nvidia.com) for assistance
+
+□  Low Gain Test
+   • Start with low kp values (2-5x lower than normal)
+   • Keep kd values unchanged
+
+□  Clear Workspace
+   • Remove obstacles and avoid tables
+   • Ensure adequate clearance in all directions
+
+□  Emergency Stop Ready
+   Ensure access to at least one:
+   • Keyboard e-stop
+   • Joycon controller
+   • External power cutoff
+
+═══════════════════════════════════════════════════════════════════════════════
+🚨 EMERGENCY: Press ` at any time to stop all processes
+📹 RECORDING: Connect a webcam to your computer to record the experiment
+═══════════════════════════════════════════════════════════════════════════════
+
+Usages:
+
+- hit `      to stop all processes
+- hit Ctrl+C to stop single process
+- hit Ctrl+\ to quit the tmux
+"""
+
+    print("")
+    print("🚨 REAL ROBOT DEPLOYMENT DETECTED 🚨")
+    print("")
+    print(checklist_content)
+    print("")
+
+    # Get user confirmation
+    while True:
+        user_input = input("Continue with deployment? [Y/n]: ").strip()
+
+        # Default to Y if empty input
+        if not user_input:
+            user_input = "Y"
+
+        user_input_upper = user_input.upper()
+
+        if user_input_upper in ["Y", "YES"]:
+            print("")
+            print("✅ Deployment confirmed. Proceeding with robot deployment...")
+            print("")
+            return True
+        elif user_input_upper in ["N", "NO"]:
+            print("")
+            print("❌ Deployment aborted by user.")
+            print("")
+            return False
+        else:
+            print(
+                "❌ Invalid input. Please enter 'Y' for yes, 'N' for no, or press Enter for default (Y)."
+            )
+
+
+def main():
+    """Main function."""
+    # Always show the checklist
+    if not show_deployment_checklist():
+        print("Deployment cancelled.")
+        sys.exit(1)
+
+    return 0
+
+
+if __name__ == "__main__":
+    main()

GR00T-WholeBodyControl/decoupled_wbc/control/utils/service.py

@@ -0,0 +1,182 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from io import BytesIO
+from typing import Any, Callable, Dict
+
+import torch
+import zmq
+
+
+class TorchSerializer:
+    @staticmethod
+    def to_bytes(data: dict) -> bytes:
+        buffer = BytesIO()
+        torch.save(data, buffer)
+        return buffer.getvalue()
+
+    @staticmethod
+    def from_bytes(data: bytes) -> dict:
+        buffer = BytesIO(data)
+        obj = torch.load(buffer, weights_only=False)
+        return obj
+
+
+@dataclass
+class EndpointHandler:
+    handler: Callable
+    requires_input: bool = True
+
+
+class BaseInferenceServer:
+    """
+    An inference server that spin up a ZeroMQ socket and listen for incoming requests.
+    Can add custom endpoints by calling `register_endpoint`.
+    """
+
+    def __init__(self, host: str = "*", port: int = 5555):
+        self.running = True
+        self.context = zmq.Context()
+        self.socket = self.context.socket(zmq.REP)
+        self.socket.bind(f"tcp://{host}:{port}")
+        self._endpoints: dict[str, EndpointHandler] = {}
+
+        # Register the ping endpoint by default
+        self.register_endpoint("ping", self._handle_ping, requires_input=False)
+        self.register_endpoint("kill", self._kill_server, requires_input=False)
+
+    def _kill_server(self):
+        """
+        Kill the server.
+        """
+        self.running = False
+
+    def _handle_ping(self) -> dict:
+        """
+        Simple ping handler that returns a success message.
+        """
+        return {"status": "ok", "message": "Server is running"}
+
+    def register_endpoint(self, name: str, handler: Callable, requires_input: bool = True):
+        """
+        Register a new endpoint to the server.
+
+        Args:
+            name: The name of the endpoint.
+            handler: The handler function that will be called when the endpoint is hit.
+            requires_input: Whether the handler requires input data.
+        """
+        self._endpoints[name] = EndpointHandler(handler, requires_input)
+
+    def run(self):
+        addr = self.socket.getsockopt_string(zmq.LAST_ENDPOINT)
+        print(f"Server is ready and listening on {addr}")
+        while self.running:
+            try:
+                message = self.socket.recv()
+                request = TorchSerializer.from_bytes(message)
+                endpoint = request.get("endpoint", "get_action")
+
+                if endpoint not in self._endpoints:
+                    raise ValueError(f"Unknown endpoint: {endpoint}")
+
+                handler = self._endpoints[endpoint]
+                result = (
+                    handler.handler(request.get("data", {}))
+                    if handler.requires_input
+                    else handler.handler()
+                )
+                self.socket.send(TorchSerializer.to_bytes(result))
+            except Exception as e:
+                print(f"Error in server: {e}")
+                import traceback
+
+                print(traceback.format_exc())
+                self.socket.send(b"ERROR")
+
+
+class BaseInferenceClient:
+    def __init__(self, host: str = "localhost", port: int = 5555, timeout_ms: int = 15000):
+        self.context = zmq.Context()
+        self.host = host
+        self.port = port
+        self.timeout_ms = timeout_ms
+        self._init_socket()
+
+    def _init_socket(self):
+        """Initialize or reinitialize the socket with current settings"""
+        self.socket = self.context.socket(zmq.REQ)
+        self.socket.connect(f"tcp://{self.host}:{self.port}")
+
+    def ping(self) -> bool:
+        try:
+            self.call_endpoint("ping", requires_input=False)
+            return True
+        except zmq.error.ZMQError:
+            self._init_socket()  # Recreate socket for next attempt
+            return False
+
+    def kill_server(self):
+        """
+        Kill the server.
+        """
+        self.call_endpoint("kill", requires_input=False)
+
+    def call_endpoint(
+        self, endpoint: str, data: dict | None = None, requires_input: bool = True
+    ) -> dict:
+        """
+        Call an endpoint on the server.
+
+        Args:
+            endpoint: The name of the endpoint.
+            data: The input data for the endpoint.
+            requires_input: Whether the endpoint requires input data.
+        """
+        request: dict = {"endpoint": endpoint}
+        if requires_input:
+            request["data"] = data
+
+        self.socket.send(TorchSerializer.to_bytes(request))
+        message = self.socket.recv()
+        if message == b"ERROR":
+            raise RuntimeError("Server error")
+        return TorchSerializer.from_bytes(message)
+
+    def __del__(self):
+        """Cleanup resources on destruction"""
+        self.socket.close()
+        self.context.term()
+
+
+class ExternalRobotInferenceClient(BaseInferenceClient):
+    """
+    Client for communicating with the RealRobotServer
+    """
+
+    def set_observation(self, observation: dict[str, Any]):
+        self.call_endpoint("set_observation", data=observation)
+
+    def get_action(self, time: float | None = None) -> Dict[str, Any]:
+        """
+        Get the action from the server.
+        The exact definition of the observations is defined
+        by the policy, which contains the modalities configuration.
+        """
+        return self.call_endpoint("get_action", data={"time": time})
+
+    def get_modality_config(self) -> dict[str, Any]:
+        return self.call_endpoint("get_modality_config")

GR00T-WholeBodyControl/decoupled_wbc/control/utils/term_color_constants.py

@@ -0,0 +1,19 @@
+GREEN_BOLD = "\033[1;32m"
+RED_BOLD = "\033[1;31m"
+YELLOW_BOLD = "\033[1;33m"
+BLUE_BOLD = "\033[1;34m"
+MAGENTA_BOLD = "\033[1;35m"
+CYAN_BOLD = "\033[1;36m"
+WHITE_BOLD = "\033[1;37m"
+GREY_BOLD = "\033[1;90m"
+
+GREEN = "\033[32m"
+RED = "\033[31m"
+YELLOW = "\033[33m"
+BLUE = "\033[34m"
+MAGENTA = "\033[35m"
+CYAN = "\033[36m"
+WHITE = "\033[37m"
+GREY = "\033[90m"
+
+RESET = "\033[0m"

GR00T-WholeBodyControl/decoupled_wbc/control/utils/text_to_speech.py

@@ -0,0 +1,28 @@
+# text to speech
+import pyttsx3
+
+
+class TextToSpeech:
+    def __init__(self, rate: int = 150, volume: float = 1.0):
+        try:
+            self.engine = pyttsx3.init(driverName="espeak")
+            self.engine.setProperty("rate", rate)
+            self.engine.setProperty("volume", volume)
+        except Exception as e:
+            print(f"[Text To Speech] Initialization failed: {e}")
+            self.engine = None
+
+    def say(self, message: str):
+        """Speak the message if engine is available."""
+        if self.engine:
+            try:
+                self.engine.say(message)
+                self.engine.runAndWait()
+            except Exception as e:
+                print(f"[Text To Speech] Failed to say message: {e}")
+
+    def print_and_say(self, message: str, say: bool = True):
+        """Print message and optionally speak it using Text To Speech."""
+        print(message)
+        if say and self.engine is not None:
+            self.say(message)

GR00T-WholeBodyControl/decoupled_wbc/control/visualization/humanoid_visualizer.py

@@ -0,0 +1,52 @@
+import time
+
+import meshcat_shapes
+import numpy as np
+from pinocchio.visualize import MeshcatVisualizer
+
+from decoupled_wbc.control.robot_model import RobotModel
+from decoupled_wbc.control.robot_model.instantiation.g1 import instantiate_g1_robot_model
+
+
+class RobotVisualizer:
+    def __init__(self, robot: RobotModel):
+        self.robot = robot
+        self.viz = MeshcatVisualizer(
+            self.robot.pinocchio_wrapper.model,
+            self.robot.pinocchio_wrapper.collision_model,
+            self.robot.pinocchio_wrapper.visual_model,
+        )
+        try:
+            self.viz.initViewer(open=True)
+
+        except ImportError as err:
+            print("Error while initializing the viewer. It seems you should install Python meshcat")
+            print(err)
+            exit(0)
+
+        self.viz.loadViewerModel()
+        self.viz.display(self.robot.q_zero)
+
+        # Visualize frames
+        self.viz_frames = [self.robot.supplemental_info.root_frame_name]
+        for side in ["left", "right"]:
+            self.viz_frames.append(self.robot.supplemental_info.hand_frame_names[side])
+        for frame in self.viz_frames:
+            meshcat_shapes.frame(self.viz.viewer[frame], opacity=1.0)
+
+    def visualize(self, robot_state: np.ndarray):
+        # visualize robot state
+        if robot_state is not None:
+            self.robot.cache_forward_kinematics(robot_state, auto_clip=False)
+            self.viz.display(robot_state)
+            for frame_name in self.viz_frames:
+                self.viz.viewer[frame_name].set_transform(self.robot.frame_placement(frame_name).np)
+
+
+if __name__ == "__main__":
+    # robot_model = instantiate_gr1_robot_model()
+    robot_model = instantiate_g1_robot_model()
+    visualizer = RobotVisualizer(robot_model)
+    while True:
+        visualizer.visualize(robot_model.q_zero)
+        time.sleep(0.01)

GR00T-WholeBodyControl/decoupled_wbc/control/visualization/meshcat_visualizer_env.py

@@ -0,0 +1,77 @@
+from contextlib import contextmanager
+import time
+
+import gymnasium as gym
+import numpy as np
+from pinocchio.visualize import MeshcatVisualizer
+
+from decoupled_wbc.control.base.env import Env
+from decoupled_wbc.control.robot_model import RobotModel
+
+
+class MeshcatVisualizerEnv(Env):
+    def __init__(self, robot_model: RobotModel):
+        self.robot_model = robot_model
+        self.viz = MeshcatVisualizer(
+            self.robot_model.pinocchio_wrapper.model,
+            self.robot_model.pinocchio_wrapper.collision_model,
+            self.robot_model.pinocchio_wrapper.visual_model,
+        )
+        try:
+            self.viz.initViewer(open=True)
+
+        except ImportError as err:
+            print("Error while initializing the viewer. It seems you should install Python meshcat")
+            print(err)
+            exit(0)
+
+        self.viz.loadViewerModel()
+        self.visualize(self.robot_model.pinocchio_wrapper.q0)
+        time.sleep(1.0)
+
+        self._observation_space = gym.spaces.Dict(
+            {
+                "q": gym.spaces.Box(
+                    low=-2 * np.pi, high=2 * np.pi, shape=(self.robot_model.num_dofs,)
+                )
+            }
+        )
+        self._action_space = gym.spaces.Dict(
+            {
+                "q": gym.spaces.Box(
+                    low=-2 * np.pi, high=2 * np.pi, shape=(self.robot_model.num_dofs,)
+                )
+            }
+        )
+
+    def visualize(self, robot_state: np.ndarray):
+        # visualize robot state
+        if robot_state is not None:
+            self.viz.display(robot_state)
+
+    def observe(self):
+        # Dummy observation
+        return {"q": self.robot_model.pinocchio_wrapper.q0}
+
+    def queue_action(self, action: dict[str, np.ndarray]):
+        self.visualize(action["q"])
+
+    def reset(self, **kwargs):
+        self.visualize(self.robot_model.pinocchio_wrapper.q0)
+        return {"q": self.robot_model.pinocchio_wrapper.q0}
+
+    def sensors(self) -> dict[str, any]:
+        return {}
+
+    def observation_space(self) -> gym.Space:
+        return self._observation_space
+
+    def action_space(self) -> gym.Space:
+        return self._action_space
+
+    def close(self):
+        return
+
+    @contextmanager
+    def activate(self):
+        yield

GR00T-WholeBodyControl/decoupled_wbc/data/constants.py

@@ -0,0 +1,5 @@
+# This will be used for both sim and real data collection
+RS_VIEW_CAMERA_HEIGHT = 480
+RS_VIEW_CAMERA_WIDTH = 640
+
+BUCKET_BASE_PATH = "GearRawDataLeRobotV0"

GR00T-WholeBodyControl/decoupled_wbc/data/exporter.py

@@ -0,0 +1,514 @@
+import copy
+from dataclasses import dataclass
+from functools import partial
+import json
+import os
+from pathlib import Path
+import shutil
+from typing import Optional
+
+import datasets
+from datasets import load_dataset
+from datasets.utils import disable_progress_bars
+from huggingface_hub.errors import RepositoryNotFoundError
+from lerobot.common.datasets.lerobot_dataset import (
+    LeRobotDataset,
+    LeRobotDatasetMetadata,
+    compute_episode_stats,
+)
+from lerobot.common.datasets.utils import (
+    check_timestamps_sync,
+    get_episode_data_index,
+    validate_episode_buffer,
+    validate_frame,
+)
+import numpy as np
+from PIL import Image as PILImage
+import torch
+from torchvision import transforms
+
+from decoupled_wbc.control.main.config_template import ArgsConfig
+from decoupled_wbc.data.video_writer import VideoWriter
+
+disable_progress_bars()  # Disable HuggingFace progress bars
+
+
+@dataclass
+class DataCollectionInfo:
+    """
+    This dataclass stores additional information that is relevant to the data collection process.
+    """
+
+    lower_body_policy: Optional[str] = None
+    wbc_model_path: Optional[str] = None
+    teleoperator_username: Optional[str] = None
+    support_operator_username: Optional[str] = None
+    robot_type: Optional[str] = None
+    robot_id: Optional[str] = None
+
+    def to_dict(self) -> dict:
+        """Convert the dataclass to a dictionary for JSON serialization."""
+        return {
+            "lower_body_policy": self.lower_body_policy,
+            "wbc_model_path": self.wbc_model_path,
+            "teleoperator_username": self.teleoperator_username,
+            "support_operator_username": self.support_operator_username,
+            "robot_type": self.robot_type,
+            "robot_id": self.robot_id,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "DataCollectionInfo":
+        """Create a DataCollectionInfo instance from a dictionary."""
+        return cls(**data)
+
+
+class Gr00tDatasetMetadata(LeRobotDatasetMetadata):
+    """
+    Additional metadata on top of LeRobotDatasetMetadata:
+    - modality_config: Written to `meta/modality.json`
+    - discarded_episode_indices: List of episode indices that were discarded. Written to `meta/info.json`
+    """
+
+    MODALITY_CONFIG_REL_PATH = Path("meta/modality.json")
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        with open(self.root / self.MODALITY_CONFIG_REL_PATH, "rb") as f:
+            self.modality_config = json.load(f)
+
+    @classmethod
+    def create(
+        cls,
+        modality_config: dict,
+        script_config: dict,
+        data_collection_info: DataCollectionInfo,
+        *args,
+        **kwargs,
+    ):
+        cls.validate_modality_config(modality_config)
+
+        # Create base metadata object using parent class
+        obj = super().create(*args, **kwargs)
+
+        # we also need to initialize the discarded_episode_indices
+        obj.info["script_config"] = script_config
+        obj.info["discarded_episode_indices"] = []
+        obj.info["data_collection_info"] = data_collection_info.to_dict()
+        with open(obj.root / "meta" / "info.json", "w") as f:
+            json.dump(obj.info, f, indent=4)
+
+        obj.__class__ = cls
+        with open(obj.root / cls.MODALITY_CONFIG_REL_PATH, "w") as f:
+            json.dump(modality_config, f, indent=4)
+        obj.modality_config = modality_config
+        return obj
+
+    @staticmethod
+    def validate_modality_config(modality_config: dict) -> None:
+        # verify if it contains all state, action, video, annotation keys
+        valid_keys = ["state", "action", "video", "annotation"]
+        if not all(key in modality_config for key in valid_keys):
+            raise ValueError(
+                f"Modality config must contain all of the following keys: {valid_keys}"
+            )
+
+        # verify that each key has a modality_config dict
+        for key in valid_keys:
+            if key not in modality_config:
+                raise ValueError(f"Modality config must contain a '{key}' key")
+
+
+class Gr00tDataExporter(LeRobotDataset):
+    """
+    A class for exporting data collected for a single session to LeRobot Dataset.
+
+    Intended life cycle:
+    1. Create a Gr00tDataExporter object
+    2. Add frames using add_frame()
+    3. Save the episode using save_episode()
+        - This will flush the episode buffer to disk
+        - This will also close the video writers
+        - Create a new video writer and ep buffer to start new episode
+
+    If interrupted, here's the indented behavior:
+        - Interruption before save_episode() is called: loses the current episode
+        - Interruption after save_episode() is called: keeps completed episodes
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.video_writers = self.create_video_writer()
+
+    @property
+    def repo_id(self):
+        return self.meta.repo_id
+
+    @property
+    def root(self):
+        return self.meta.root
+
+    @property
+    def local_files_only(self):
+        return self.meta.local_files_only
+
+    @property
+    def video_keys(self):
+        return self.meta.video_keys
+
+    @classmethod
+    def create(
+        cls,
+        save_root: str | Path,
+        fps: int,
+        features: dict,
+        modality_config: dict,
+        task: str,
+        script_config: ArgsConfig = ArgsConfig(),
+        data_collection_info: DataCollectionInfo = DataCollectionInfo(),
+        robot_type: str | None = None,
+        tolerance_s: float = 1e-4,
+        vcodec: str = "h264",
+        overwrite_existing: bool = False,
+        upload_bucket_path: str | None = None,
+    ) -> "Gr00tDataExporter":
+        """
+        Create a Gr00tDataExporter object.
+
+        Args:
+            save_root: The root directory to save the dataset.
+            fps: The frame rate of the dataset.
+            features: The features of the dataset.
+            modality_config: The modality config of the dataset.
+            task: The task performed during the data collection session.
+            data_collection_info: The data collection info.
+                If the dataset already exists, this argument will be ignored.
+                If data_collection_info is not provided, it will be set to an empty DataCollectionInfo object.
+            robot_type: The type of robot.
+            tolerance_s: The tolerance for the dataset.
+            image_writer_processes: The number of processes to use for the image writer.
+            image_writer_threads: The number of threads to use for the image writer.
+            vcodec: The codec to use for the video writer.
+        """
+
+        obj = cls.__new__(cls)
+        repo_id = (
+            "tmp/tmp_dataset"  # NOTE(fengyuanh): Not relevant since we are not pushing to the hub
+        )
+        if overwrite_existing and (Path(save_root)).exists():
+            print(
+                f"Found existing dataset at {save_root}",
+                "Cleaning up this directory since overwrite_existing is True.",
+            )
+            shutil.rmtree(save_root)
+
+        if (Path(save_root)).exists():
+            # Try to resume from existing dataset
+            try:
+                # Load the metadata
+                obj.meta = Gr00tDatasetMetadata(
+                    repo_id=repo_id,
+                    root=save_root,
+                )
+
+            except RepositoryNotFoundError as e:
+                raise ValueError(
+                    f"Failed to resume from corrupted dataset. Please manually check the dataset at {save_root}"
+                ) from e
+        else:
+            if not isinstance(script_config, dict):
+                script_config = script_config.to_dict()
+            obj.meta = Gr00tDatasetMetadata.create(
+                repo_id=repo_id,
+                fps=fps,
+                root=save_root,
+                # NOTE(fengyuanh): We use "robot_type" instead of this field which requires a Robot object
+                robot=None,
+                robot_type=robot_type,
+                features=features,
+                modality_config=modality_config,
+                script_config=script_config,
+                # NOTE(fengyuanh): Always use videos for exporting
+                use_videos=True,
+                data_collection_info=data_collection_info,
+            )
+        obj.tolerance_s = tolerance_s
+        obj.video_backend = (
+            "pyav"  # NOTE(fengyuanh): Only used in training, not relevant for exporting
+        )
+        obj.vcodec = vcodec
+        obj.task = task
+        obj.image_writer = None
+
+        obj.episode_buffer = obj.create_episode_buffer()
+
+        obj.episodes = None
+        obj.hf_dataset = obj.create_hf_dataset()
+        obj.image_transforms = None
+        obj.delta_timestamps = None
+        obj.delta_indices = None
+        obj.episode_data_index = None
+        obj.upload_bucket_path = upload_bucket_path
+        obj.video_writers = obj.create_video_writer()
+        return obj
+
+    def create_video_writer(self) -> dict[str, VideoWriter]:
+        video_writers = {}
+        for key in self.meta.video_keys:
+            video_writers[key] = VideoWriter(
+                self.root
+                / self.meta.get_video_file_path(self.episode_buffer["episode_index"], key),
+                self.meta.shapes[key][1],
+                self.meta.shapes[key][0],
+                self.fps,
+                self.vcodec,
+            )
+        return video_writers
+
+    # @note (k2): This function is copied from LeRobotDataset.add_frame.
+    # This is done because we want to bypass lerobot's
+    # image_writer and use our own VideoWriter class.
+    def add_frame(self, frame: dict) -> None:
+        """
+        This function only adds the frame to the episode_buffer. Videos are handled by the video_writer,
+        which uses a stream writer to write to disk.
+        """
+        frame = copy.deepcopy(frame)
+        frame["task"] = frame.get("task", self.task)
+
+        # Convert torch to numpy if needed
+        for name in frame:
+            if isinstance(frame[name], torch.Tensor):
+                frame[name] = frame[name].numpy()
+
+        validate_frame(frame, self.features)
+
+        if self.episode_buffer is None:
+            self.episode_buffer = self.create_episode_buffer()
+
+        # Automatically add frame_index and timestamp to episode buffer
+        frame_index = self.episode_buffer["size"]
+        timestamp = frame.pop("timestamp") if "timestamp" in frame else frame_index / self.fps
+        self.episode_buffer["frame_index"].append(frame_index)
+        self.episode_buffer["timestamp"].append(timestamp)
+
+        # Add frame features to episode_buffer
+        for key in frame:
+            if key == "task":
+                # Note: we associate the task in natural language to its task index during `save_episode`
+                self.episode_buffer["task"].append(frame["task"])
+                continue
+
+            if key not in self.features:
+                raise ValueError(
+                    f"An element of the frame is not in the features. '{key}' not in '{self.features.keys()}'."
+                )
+
+            if self.features[key]["dtype"] in ["image", "video"]:
+                img_path = self._get_image_file_path(
+                    episode_index=self.episode_buffer["episode_index"],
+                    image_key=key,
+                    frame_index=frame_index,
+                )
+                if frame_index == 0:
+                    img_path.parent.mkdir(parents=True, exist_ok=True)
+
+                # @note (k2): using our own VideoWriter class, bypassing the image_writer
+                self.video_writers[key].add_frame(frame[key])
+                self.episode_buffer[key].append(str(img_path))
+            else:
+                self.episode_buffer[key].append(frame[key])
+
+        self.episode_buffer["size"] += 1
+
+    def stop_video_writers(self):
+        if not hasattr(self, "video_writers"):
+            raise RuntimeError(
+                "Can't stop video writers because they haven't been initialized. Call create() first."
+            )
+        for key in self.video_writers:
+            self.video_writers[key].stop()
+
+    def skip_and_start_new_episode(
+        self,
+    ) -> None:
+        """
+        Skip the current episode and start a new one.
+        """
+        self.stop_video_writers()
+        self.episode_buffer = self.create_episode_buffer()
+        self.video_writers = self.create_video_writer()
+
+    # @note (k2): Code copied from LeRobotDataset.save_episode
+    # We override this function because we want to bypass lerobot's `compute_episode_stats` on video features
+    # since `compute_episode_stats` only works when images are written to disk.
+    def save_episode(self, episode_data: dict | None = None) -> None:
+        if not episode_data:
+            episode_buffer = self.episode_buffer
+
+        validate_episode_buffer(episode_buffer, self.meta.total_episodes, self.features)
+
+        # size and task are special cases that won't be added to hf_dataset
+        episode_length = episode_buffer.pop("size")
+        tasks = episode_buffer.pop("task")
+        episode_tasks = list(set(tasks))
+        episode_index = episode_buffer["episode_index"]
+
+        episode_buffer["index"] = np.arange(
+            self.meta.total_frames, self.meta.total_frames + episode_length
+        )
+        episode_buffer["episode_index"] = np.full((episode_length,), episode_index)
+
+        # Add new tasks to the tasks dictionary
+        for task in episode_tasks:
+            task_index = self.meta.get_task_index(task)
+            if task_index is None:
+                self.meta.add_task(task)
+
+        # Given tasks in natural language, find their corresponding task indices
+        episode_buffer["task_index"] = np.array([self.meta.get_task_index(task) for task in tasks])
+
+        for key, ft in self.features.items():
+            # index, episode_index, task_index are already processed above, and image and video
+            # are processed separately by storing image path and frame info as meta data
+            if key in ["index", "episode_index", "task_index"] or ft["dtype"] in ["image", "video"]:
+                continue
+            episode_buffer[key] = np.stack(episode_buffer[key])
+
+        self._wait_image_writer()
+        self._save_episode_table(episode_buffer, episode_index)
+
+        # @note (k2): computing only non-video features stats
+        non_video_features = {k: v for k, v in self.features.items() if v["dtype"] not in ["video"]}
+        non_vid_ep_buffer = {
+            k: v for k, v in episode_buffer.items() if k in non_video_features.keys()
+        }
+        ep_stats = compute_episode_stats(non_vid_ep_buffer, non_video_features)
+
+        if len(self.meta.video_keys) > 0:
+            video_paths = self.encode_episode_videos(episode_index)
+            for key in self.meta.video_keys:
+                episode_buffer[key] = video_paths[key]
+
+        # `meta.save_episode` be executed after encoding the videos
+        self.meta.save_episode(episode_index, episode_length, episode_tasks, ep_stats)
+
+        ep_data_index = get_episode_data_index(self.meta.episodes, [episode_index])
+        ep_data_index_np = {k: t.numpy() for k, t in ep_data_index.items()}
+        check_timestamps_sync(
+            episode_buffer["timestamp"],
+            episode_buffer["episode_index"],
+            ep_data_index_np,
+            self.fps,
+            self.tolerance_s,
+        )
+
+        video_files = list(self.root.rglob("*.mp4"))
+        assert len(video_files) == self.num_episodes * len(self.meta.video_keys)
+
+        parquet_files = list(self.root.rglob("*.parquet"))
+        assert len(parquet_files) == self.num_episodes
+
+        # delete images
+        img_dir = self.root / "images"
+        if img_dir.is_dir():
+            shutil.rmtree(self.root / "images")
+
+        if not episode_data:  # Reset the buffer and create new video writers
+            self.episode_buffer = self.create_episode_buffer()
+            self.video_writers = self.create_video_writer()
+
+        # check if all video and parquet files exist
+        for key in self.meta.video_keys:
+            video_path = os.path.join(self.root, self.meta.get_video_file_path(episode_index, key))
+            if not os.path.exists(video_path):
+                raise FileNotFoundError(
+                    f"Video path: {video_path} does not exist for episode {episode_index}"
+                )
+
+        parquet_path = os.path.join(self.root, self.meta.get_data_file_path(episode_index))
+        if not os.path.exists(parquet_path):
+            raise FileNotFoundError(
+                f"Parquet path: {parquet_path} does not exist for episode {episode_index}"
+            )
+
+    # @note (k2): Overriding LeRobotDataset.encode_episode_videos to use our own VideoWriter class
+    def encode_episode_videos(self, episode_index: int) -> dict:
+        video_paths = {}
+        for key in self.meta.video_keys:
+            video_paths[key] = self.video_writers[key].stop()
+        return video_paths
+
+    def save_episode_as_discarded(self) -> None:
+        """
+        Flag ongoing episode as discarded and save it to disk. Failed manipulations (grasp, manipulation) are
+        flagged as discarded. It will add the episode index to the discarded episode indices list in info.json.
+        """
+        self.meta.info["discarded_episode_indices"] = self.meta.info.get(
+            "discarded_episode_indices", []
+        ) + [self.episode_buffer["episode_index"]]
+        self.save_episode()
+
+
+def hf_transform_to_torch_by_features(
+    features: datasets.Sequence, items_dict: dict[torch.Tensor | None]
+):
+    """Get a transform function that convert items from Hugging Face dataset (pyarrow)
+    to torch tensors. Importantly, images are converted from PIL, which corresponds to
+    a channel last representation (h w c) of uint8 type, to a torch image representation
+    with channel first (c h w) of float32 type in range [0,1].
+    """
+    for key in items_dict:
+        first_item = items_dict[key][0]
+        if isinstance(first_item, PILImage.Image):
+            to_tensor = transforms.ToTensor()
+            items_dict[key] = [to_tensor(img) for img in items_dict[key]]
+        elif first_item is None:
+            pass
+        else:
+            if isinstance(features[key], datasets.Value):
+                dtype_str = features[key].dtype
+            elif isinstance(features[key], datasets.Sequence):
+                assert isinstance(features[key].feature, datasets.Value)
+                dtype_str = features[key].feature.dtype
+            else:
+                raise ValueError(f"Unsupported feature type for key '{key}': {features[key]}")
+            dtype_mapping = {
+                "float32": torch.float32,
+                "float64": torch.float64,
+                "int32": torch.int32,
+                "int64": torch.int64,
+            }
+            items_dict[key] = [
+                torch.tensor(x, dtype=dtype_mapping[dtype_str]) for x in items_dict[key]
+            ]
+    return items_dict
+
+
+# This is a subclass of LeRobotDataset that only fixes the data type when loading
+# By default, LeRobotDataset will automatically convert float64 to float32
+class TypedLeRobotDataset(LeRobotDataset):
+    def __init__(self, load_video=True, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if not load_video:
+            video_keys = []
+            for key in self.meta.features.keys():
+                if self.meta.features[key]["dtype"] == "video":
+                    video_keys.append(key)
+            for key in video_keys:
+                self.meta.features.pop(key)
+
+    def load_hf_dataset(self) -> datasets.Dataset:
+        """hf_dataset contains all the observations, states, actions, rewards, etc."""
+        if self.episodes is None:
+            path = str(self.root / "data")
+            hf_dataset = load_dataset("parquet", data_dir=path, split="train")
+        else:
+            files = [
+                str(self.root / self.meta.get_data_file_path(ep_idx)) for ep_idx in self.episodes
+            ]
+            hf_dataset = load_dataset("parquet", data_files=files, split="train")
+
+        # TODO(aliberts): hf_dataset.set_format("torch")
+        hf_dataset.set_transform(partial(hf_transform_to_torch_by_features, hf_dataset.features))
+        return hf_dataset

GR00T-WholeBodyControl/decoupled_wbc/data/utils.py

@@ -0,0 +1,156 @@
+from decoupled_wbc.control.robot_model.robot_model import RobotModel
+from decoupled_wbc.data.constants import RS_VIEW_CAMERA_HEIGHT, RS_VIEW_CAMERA_WIDTH
+
+
+def get_modality_config(robot_model: RobotModel, add_stereo_camera: bool = False) -> dict:
+    """
+    Get the modality config for the robot model.
+    """
+    left_hand_indices = sorted(robot_model.get_joint_group_indices("left_hand"))
+    right_hand_indices = sorted(robot_model.get_joint_group_indices("right_hand"))
+    left_arm_indices = sorted(robot_model.get_joint_group_indices("left_arm"))
+    right_arm_indices = sorted(robot_model.get_joint_group_indices("right_arm"))
+    waist_indices = sorted(robot_model.get_joint_group_indices("waist"))
+    left_leg_indices = sorted(robot_model.get_joint_group_indices("left_leg"))
+    right_leg_indices = sorted(robot_model.get_joint_group_indices("right_leg"))
+
+    modality_config = {
+        "state": {
+            "left_leg": {"start": left_leg_indices[0], "end": left_leg_indices[-1] + 1},
+            "right_leg": {"start": right_leg_indices[0], "end": right_leg_indices[-1] + 1},
+            "waist": {"start": waist_indices[0], "end": waist_indices[-1] + 1},
+            "left_arm": {"start": left_arm_indices[0], "end": left_arm_indices[-1] + 1},
+            "left_hand": {"start": left_hand_indices[0], "end": left_hand_indices[-1] + 1},
+            "right_arm": {"start": right_arm_indices[0], "end": right_arm_indices[-1] + 1},
+            "right_hand": {"start": right_hand_indices[0], "end": right_hand_indices[-1] + 1},
+            "left_wrist_pos": {"start": 0, "end": 3, "original_key": "observation.eef_state"},
+            "left_wrist_abs_quat": {
+                "start": 3,
+                "end": 7,
+                "original_key": "observation.eef_state",
+                "rotation_type": "quaternion",
+            },
+            "right_wrist_pos": {"start": 7, "end": 10, "original_key": "observation.eef_state"},
+            "right_wrist_abs_quat": {
+                "start": 10,
+                "end": 14,
+                "original_key": "observation.eef_state",
+                "rotation_type": "quaternion",
+            },
+        },
+        "action": {
+            "left_leg": {"start": left_leg_indices[0], "end": left_leg_indices[-1] + 1},
+            "right_leg": {"start": right_leg_indices[0], "end": right_leg_indices[-1] + 1},
+            "waist": {"start": waist_indices[0], "end": waist_indices[-1] + 1},
+            "left_arm": {"start": left_arm_indices[0], "end": left_arm_indices[-1] + 1},
+            "left_hand": {"start": left_hand_indices[0], "end": left_hand_indices[-1] + 1},
+            "right_arm": {"start": right_arm_indices[0], "end": right_arm_indices[-1] + 1},
+            "right_hand": {"start": right_hand_indices[0], "end": right_hand_indices[-1] + 1},
+            "left_wrist_pos": {"start": 0, "end": 3, "original_key": "action.eef"},
+            "left_wrist_abs_quat": {
+                "start": 3,
+                "end": 7,
+                "original_key": "action.eef",
+                "rotation_type": "quaternion",
+            },
+            "right_wrist_pos": {"start": 7, "end": 10, "original_key": "action.eef"},
+            "right_wrist_abs_quat": {
+                "start": 10,
+                "end": 14,
+                "original_key": "action.eef",
+                "rotation_type": "quaternion",
+            },
+            "base_height_command": {
+                "start": 0,
+                "end": 1,
+                "original_key": "teleop.base_height_command",
+            },
+            "navigate_command": {"start": 0, "end": 3, "original_key": "teleop.navigate_command"},
+        },
+        "video": {"ego_view": {"original_key": "observation.images.ego_view"}},
+        "annotation": {"human.task_description": {"original_key": "task_index"}},
+    }
+    if add_stereo_camera:
+        modality_config["video"].update(
+            {
+                "ego_view_left_mono": {"original_key": "observation.images.ego_view_left_mono"},
+                "ego_view_right_mono": {"original_key": "observation.images.ego_view_right_mono"},
+            }
+        )
+
+    return modality_config
+
+
+def get_dataset_features(robot_model: RobotModel, add_stereo_camera: bool = False) -> dict:
+    """
+    Get the dataset features for the robot model.
+    """
+    dataset_features = {
+        "observation.images.ego_view": {
+            "dtype": "video",
+            "shape": [RS_VIEW_CAMERA_HEIGHT, RS_VIEW_CAMERA_WIDTH, 3],
+            "names": ["height", "width", "channel"],
+        },
+        "observation.state": {
+            "dtype": "float64",
+            "shape": (robot_model.num_joints,),
+            "names": robot_model.joint_names,
+        },
+        "observation.eef_state": {
+            "dtype": "float64",
+            "shape": (14,),
+            "names": [
+                "left_wrist_pos",
+                "left_wrist_abs_quat",
+                "right_wrist_pos",
+                "right_wrist_abs_quat",
+            ],
+        },
+        "action": {
+            "dtype": "float64",
+            "shape": (robot_model.num_joints,),
+            "names": robot_model.joint_names,
+        },
+        "action.eef": {
+            "dtype": "float64",
+            "shape": (14,),
+            "names": [
+                "left_wrist_pos",
+                "left_wrist_abs_quat",
+                "right_wrist_pos",
+                "right_wrist_abs_quat",
+            ],
+        },
+        "observation.img_state_delta": {
+            "dtype": "float32",
+            "shape": (1,),
+            "names": "img_state_delta",
+        },
+        "teleop.navigate_command": {
+            "dtype": "float64",
+            "shape": (3,),
+            "names": ["lin_vel_x", "lin_vel_y", "ang_vel_z"],
+        },
+        "teleop.base_height_command": {
+            "dtype": "float64",
+            "shape": (1,),
+            "names": "base_height_command",
+        },
+    }
+    if add_stereo_camera:
+        dataset_features.update(
+            {
+                "observation.images.ego_view_left_mono": {
+                    "dtype": "video",
+                    "shape": [RS_VIEW_CAMERA_HEIGHT, RS_VIEW_CAMERA_WIDTH, 3],
+                    "names": ["height", "width", "channel"],
+                },
+                "observation.images.ego_view_right_mono": {
+                    "dtype": "video",
+                    "shape": [RS_VIEW_CAMERA_HEIGHT, RS_VIEW_CAMERA_WIDTH, 3],
+                    "names": ["height", "width", "channel"],
+                },
+            }
+        )
+
+    return dataset_features

GR00T-WholeBodyControl/decoupled_wbc/data/video_writer.py

@@ -0,0 +1,102 @@
+import os
+import queue
+import sys
+import threading
+import time
+
+import av
+import numpy as np
+
+
+class VideoWriter:
+    def __init__(
+        self,
+        output_path: str,
+        width: int,
+        height: int,
+        fps: float,
+        codec: str = "h264",
+        buffer_size: int = 50,
+    ):
+        self.output_path = output_path
+        self._first_frame = True  # Track first frame to suppress x264 info output
+
+        # Create output directory if it doesn't exist
+        output_dir = os.path.dirname(output_path)
+        if output_dir and not os.path.exists(output_dir):
+            os.makedirs(output_dir, exist_ok=True)
+
+        self.queue = queue.Queue(maxsize=buffer_size)
+        self.container = av.open(output_path, mode="w")
+        self.stream = self.container.add_stream(codec, rate=fps)
+        self.stream.width = width
+        self.stream.height = height
+        thread = threading.Thread(target=self._writer_worker, daemon=True)
+        thread.start()
+
+    def _assert_dimensions(self, frame: np.ndarray) -> None:
+        assert (
+            frame.shape[1] == self.stream.width and frame.shape[0] == self.stream.height
+        ), f"""Incorrect frame dimensions. Input dimensions: {frame.shape[1]}x{frame.shape[0]}. 
+            Expected dimensions: {self.stream.width}x{self.stream.height}"""
+
+    def add_frame(self, frame: np.ndarray) -> None:
+        self._assert_dimensions(frame)
+        self.queue.put(frame)
+
+    def _writer_worker(self) -> None:
+        while True:
+            frame = self.queue.get()
+            if frame is None:
+                continue
+            self._assert_dimensions(frame)
+            frame = av.VideoFrame.from_ndarray(frame, format="rgb24")
+
+            # Suppress stderr for first frame encoding (x264 prints info then)
+            if self._first_frame:
+                stderr_fd = sys.stderr.fileno()
+                old_stderr = os.dup(stderr_fd)
+                devnull = os.open(os.devnull, os.O_WRONLY)
+                os.dup2(devnull, stderr_fd)
+                try:
+                    packets = self.stream.encode(frame)
+                    for packet in packets:
+                        self.container.mux(packet)
+                finally:
+                    os.dup2(old_stderr, stderr_fd)
+                    os.close(old_stderr)
+                    os.close(devnull)
+                    self._first_frame = False
+            else:
+                packets = self.stream.encode(frame)
+                for packet in packets:
+                    self.container.mux(packet)
+
+    def _flush_stream(self) -> None:
+        packets = self.stream.encode()
+        for packet in packets:
+            self.container.mux(packet)
+
+    def stop(self) -> str:
+        """
+        Blocking call. Waits until all the frames in the queue have been written to the file
+        and the video writer has been closed.
+        """
+        if not self.queue.empty():
+            print("Waiting for video writer queue to empty...")
+            while not self.queue.empty():
+                time.sleep(0.1)
+
+        print("Video writer queue is empty, flushing stream...")
+        self._flush_stream()
+        self.container.close()
+        return self.output_path
+
+    def cancel(self) -> None:
+        """Immediately stops writing and deletes the output file"""
+        if os.path.exists(self.output_path):
+            os.remove(self.output_path)
+        self.container.close()
+
+    def __del__(self) -> None:
+        self.container.close()

GR00T-WholeBodyControl/decoupled_wbc/data/viz/rerun_viz.py

@@ -0,0 +1,213 @@
+#!/usr/bin/env python3
+"""Rerun visualization utilities for plotting data and images."""
+
+from __future__ import annotations
+
+import argparse
+import time
+from typing import Dict, List, Optional
+
+import cv2
+import numpy as np
+import rerun as rr  # pip install rerun-sdk
+import rerun.blueprint as rrb
+
+
+class RerunViz:
+    """Class for visualizing data using Rerun."""
+
+    def __init__(
+        self,
+        image_keys: List[str],
+        tensor_keys: List[str],
+        app_name: str = "rerun_visualization",
+        memory_limit: str = "1GB",
+        window_size: float = 5.0,
+        port: int = 9876,
+        in_docker: bool = False,
+    ):
+        """Initialize the RerunViz class.
+        Args:
+            app_name: Name of the Rerun application
+            memory_limit: Memory limit for Rerun
+            window_size: Size of the time window in seconds
+            image_keys: List of image keys to plot
+            tensor_keys: List of tensor keys to plot
+            in_docker: Whether running inside Docker container. If in docker,
+                forward data to outside of the container to be rendered.
+                Use `rerun --port 9876` to visualize. Expecting rerun-cli 0.22.1 outside of docker.
+                Tested with rerun-sdk 0.21.0 inside docker.
+        """
+        self.app_name = app_name
+        self.memory_limit = memory_limit
+        self.window_size = window_size
+        self.tensor_keys = tensor_keys
+        self.image_keys = image_keys
+        self.port = port
+        self.in_docker = in_docker
+        # Initialize Rerun
+        self._initialize_rerun()
+
+    def _initialize_rerun(self):
+        """Initialize Rerun and set up the blueprint."""
+        rr.init(self.app_name)
+        if not self.in_docker:
+            # support for web visualization
+            rr.spawn(memory_limit=self.memory_limit, port=self.port, connect=True)
+        else:
+            # forward data to outside of the docker container
+            rr.connect(f"127.0.0.1:{self.port}")
+        self._create_blueprint()
+
+    def _create_blueprint(self):
+        # Create a grid of plots
+        contents = []
+
+        # Add time series plots
+        for tensor_key in self.tensor_keys:
+            contents.append(
+                rrb.TimeSeriesView(
+                    origin=tensor_key,
+                    time_ranges=[
+                        rrb.VisibleTimeRange(
+                            "time",
+                            start=rrb.TimeRangeBoundary.cursor_relative(seconds=-self.window_size),
+                            end=rrb.TimeRangeBoundary.cursor_relative(),
+                        )
+                    ],
+                )
+            )
+
+        # Add image views
+        for image_key in self.image_keys:
+            contents.append(rrb.Spatial2DView(origin=image_key, name=image_key))
+
+        # Send the blueprint with collapsed panels to hide side/bottom bars
+        rr.send_blueprint(rrb.Blueprint(rrb.Grid(contents=contents), collapse_panels=True))
+
+    def set_rerun_keys(self, image_keys: List[str], tensor_keys: List[str]):
+        """Set the Rerun keys."""
+        self.image_keys = image_keys
+        self.tensor_keys = tensor_keys
+        self._create_blueprint()
+
+    def plot_images(self, images: Dict[str, np.ndarray], timestamp: Optional[float] = None):
+        """Plot image data.
+
+        Args:
+            images: Dictionary mapping image names to image data
+            timestamp: Timestamp for the data (if None, uses current time)
+        """
+        if timestamp is None:
+            timestamp = time.time()
+
+        rr.set_time_seconds("time", timestamp)
+
+        for key, image in images.items():
+            if image is None:
+                continue
+
+            if "depth" in key:
+                # Color jet
+                depth_colormap = cv2.applyColorMap(
+                    cv2.convertScaleAbs(image, alpha=0.03), cv2.COLORMAP_JET
+                )
+                rr.log(f"{key}", rr.Image(depth_colormap))
+            else:
+                # Convert to RGB
+                # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+                rr.log(f"{key}", rr.Image(image))
+
+    def plot_tensors(
+        self, data: Optional[Dict[str, np.ndarray]] = None, timestamp: Optional[float] = None
+    ):
+        """Plot tensor data.
+
+        Args:
+            data: Dictionary mapping keys to tensor values
+            timestamp: Timestamp for the data (if None, uses current time)
+        """
+        if timestamp is None:
+            timestamp = time.time()
+
+        rr.set_time_seconds("time", timestamp)
+
+        # If no data provided, use random walk generators
+        for tensor_key in self.tensor_keys:
+            for i in range(data[tensor_key].shape[0]):
+                rr.log(f"{tensor_key}/{i}", rr.Scalar(data[tensor_key][i]))
+
+    def close(self):
+        """Close the RerunViz instance."""
+        rr.rerun_shutdown()
+
+
+if __name__ == "__main__":
+    """Main function to demonstrate the RerunViz class."""
+    parser = argparse.ArgumentParser(description="Plot dashboard stress test")
+    parser.add_argument(
+        "--freq", type=float, default=20, help="Frequency of logging (applies to all series)"
+    )
+    parser.add_argument(
+        "--window-size", type=float, default=5.0, help="Size of the window in seconds"
+    )
+    parser.add_argument("--duration", type=float, default=60, help="How long to log for in seconds")
+    parser.add_argument("--use-rs", action="store_true", help="Use RealSense sensor")
+    parser.add_argument("--use-zed", action="store_true", help="Use ZED sensor")
+    parser.add_argument("--in-docker", action="store_true", help="Running inside Docker container")
+    args = parser.parse_args()
+
+    if args.use_rs:
+        image_keys = ["color_image", "depth_image"]
+        from decoupled_wbc.control.sensor.realsense import RealSenseClientSensor
+
+        sensor = RealSenseClientSensor()
+    elif args.use_zed:
+        image_keys = ["left_image", "right_image"]
+        from decoupled_wbc.control.sensor.zed import ZEDClientSensor
+
+        sensor = ZEDClientSensor()
+    else:
+        from decoupled_wbc.control.sensor.dummy import DummySensor
+
+        sensor = DummySensor()
+        image_keys = ["color_image"]
+
+    tensor_keys = ["left_arm_qpos", "left_hand_qpos", "right_arm_qpos", "right_hand_qpos"]
+
+    # Initialize the RerunViz class
+    viz = RerunViz(
+        image_keys=image_keys,
+        tensor_keys=tensor_keys,
+        window_size=args.window_size,
+        in_docker=args.in_docker,
+    )
+
+    # Run the visualization loop
+    cur_time = time.time()
+    end_time = cur_time + args.duration
+    time_per_tick = 1.0 / args.freq
+
+    while cur_time < end_time:
+        # Advance time and sleep if necessary
+        cur_time += time_per_tick
+        sleep_for = cur_time - time.time()
+        if sleep_for > 0:
+            time.sleep(sleep_for)
+
+        if sleep_for < -0.1:
+            print(f"Warning: missed logging window by {-sleep_for:.2f} seconds")
+
+        # Plot dummy tensor
+        dummy_tensor = np.random.randn(5)
+        dummy_tensor_dict = {key: dummy_tensor for key in tensor_keys}
+
+        viz.plot_tensors(dummy_tensor_dict, cur_time)
+
+        # Plot images if available
+        images = sensor.read()
+        if images is not None:
+            img_to_show = {key: images[key] for key in image_keys}
+            viz.plot_images(img_to_show, cur_time)
+
+    rr.script_teardown(args)

GR00T-WholeBodyControl/decoupled_wbc/docker/.bashrc

@@ -0,0 +1,163 @@
+# ~/.bashrc: executed by bash(1) for non-login shells.
+# see /usr/share/doc/bash/examples/startup-files (in the package bash-doc)
+# for examples
+
+# If not running interactively, don't do anything
+case $- in
+    *i*) ;;
+      *) return;;
+esac
+
+# don't put duplicate lines or lines starting with space in the history.
+# See bash(1) for more options
+HISTCONTROL=ignoreboth
+
+# append to the history file, don't overwrite it
+shopt -s histappend
+
+# for setting history length see HISTSIZE and HISTFILESIZE in bash(1)
+HISTSIZE=1000
+HISTFILESIZE=2000
+
+# check the window size after each command and, if necessary,
+# update the values of LINES and COLUMNS.
+shopt -s checkwinsize
+
+# If set, the pattern "**" used in a pathname expansion context will
+# match all files and zero or more directories and subdirectories.
+#shopt -s globstar
+
+# make less more friendly for non-text input files, see lesspipe(1)
+[ -x /usr/bin/lesspipe ] && eval "$(SHELL=/bin/sh lesspipe)"
+
+# set variable identifying the chroot you work in (used in the prompt below)
+if [ -z "${debian_chroot:-}" ] && [ -r /etc/debian_chroot ]; then
+    debian_chroot=$(cat /etc/debian_chroot)
+fi
+
+# set a fancy prompt (non-color, unless we know we "want" color)
+case "$TERM" in
+    xterm-color|*-256color) color_prompt=yes;;
+esac
+
+# uncomment for a colored prompt, if the terminal has the capability; turned
+# off by default to not distract the user: the focus in a terminal window
+# should be on the output of commands, not on the prompt
+force_color_prompt=yes
+
+if [ -n "$force_color_prompt" ]; then
+    if [ -x /usr/bin/tput ] && tput setaf 1 >&/dev/null; then
+	# We have color support; assume it's compliant with Ecma-48
+	# (ISO/IEC-6429). (Lack of such support is extremely rare, and such
+	# a case would tend to support setf rather than setaf.)
+	color_prompt=yes
+    else
+	color_prompt=
+    fi
+fi
+
+if [ "$color_prompt" = yes ]; then
+    PS1='${debian_chroot:+($debian_chroot)}\[\033[01;32m\]\u@\h\[\033[00m\]:\[\033[01;34m\]\w\[\033[00m\]\$ '
+else
+    PS1='${debian_chroot:+($debian_chroot)}\u@\h:\w\$ '
+fi
+unset color_prompt force_color_prompt
+
+# If this is an xterm set the title to user@host:dir
+case "$TERM" in
+xterm*|rxvt*)
+    PS1="\[\e]0;${debian_chroot:+($debian_chroot)}\u@\h: \w\a\]$PS1"
+    ;;
+*)
+    ;;
+esac
+
+# enable color support of ls and also add handy aliases
+if [ -x /usr/bin/dircolors ]; then
+    test -r ~/.dircolors && eval "$(dircolors -b ~/.dircolors)" || eval "$(dircolors -b)"
+    alias ls='ls --color=auto'
+    #alias dir='dir --color=auto'
+    #alias vdir='vdir --color=auto'
+
+    alias grep='grep --color=auto'
+    alias fgrep='fgrep --color=auto'
+    alias egrep='egrep --color=auto'
+fi
+
+# colored GCC warnings and errors
+export GCC_COLORS='error=01;31:warning=01;35:note=01;36:caret=01;32:locus=01:quote=01'
+
+# Set terminal type for color support
+export TERM=xterm-256color
+
+# some more ls aliases
+alias ll='ls -alF'
+alias la='ls -A'
+alias l='ls -CF'
+
+# Add an "alert" alias for long running commands.  Use like so:
+#   sleep 10; alert
+alias alert='notify-send --urgency=low -i "$([ $? = 0 ] && echo terminal || echo error)" "$(history|tail -n1|sed -e '\''s/^\s*[0-9]\+\s*//;s/[;&|]\s*alert$//'\'')"'
+
+# Alias definitions.
+# You may want to put all your additions into a separate file like
+# ~/.bash_aliases, instead of adding them here directly.
+# See /usr/share/doc/bash-doc/examples in the bash-doc package.
+
+if [ -f ~/.bash_aliases ]; then
+    . ~/.bash_aliases
+fi
+
+# enable programmable completion features (you don't need to enable
+# this, if it's already enabled in /etc/bash.bashrc and /etc/profile
+# sources /etc/bash.bashrc).
+if ! shopt -oq posix; then
+  if [ -f /usr/share/bash-completion/bash_completion ]; then
+    . /usr/share/bash-completion/bash_completion
+  elif [ -f /etc/bash_completion ]; then
+    . /etc/bash_completion
+  fi
+fi
+
+# useful commands
+bind '"\e[A": history-search-backward'
+bind '"\e[B": history-search-forward'
+
+# Store the last 10 directories in a history file
+CD_HISTFILE=~/.cd_history
+CD_HISTSIZE=10
+
+cd() {
+    local histfile="${CD_HISTFILE:-$HOME/.cd_history}"
+    local max="${CD_HISTSIZE:-10}"
+
+    case "$1" in
+        --)   [ -f "$histfile" ] && tac "$histfile" | nl -w2 -s' ' || echo "No directory history yet."; return ;;
+        -[0-9]*)
+            local idx=${1#-}
+            local dir=$(tac "$histfile" 2>/dev/null | sed -n "${idx}p")
+            [ -n "$dir" ] && builtin cd "$dir" || echo "Invalid selection: $1"
+            return ;;
+    esac
+
+    builtin cd "$@" || return
+
+    [[ $(tail -n1 "$histfile" 2>/dev/null) != "$PWD" ]] && echo "$PWD" >> "$histfile"
+    tail -n "$max" "$histfile" > "${histfile}.tmp" && mv "${histfile}.tmp" "$histfile"
+}
+
+# Make decoupled_wbc importable
+export PYTHONPATH="${DECOUPLED_WBC_DIR}:${PYTHONPATH}"
+
+# Manus to LD_LIBRARY_PATH
+export LD_LIBRARY_PATH=$DECOUPLED_WBC_DIR/decoupled_wbc/control/teleop/device/SDKClient_Linux/ManusSDK/lib:$LD_LIBRARY_PATH
+
+# CUDA support
+export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:/usr/lib:/lib/x86_64-linux-gnu:/lib64:/lib:$LD_LIBRARY_PATH 
+
+# decoupled_wbc aliases
+alias dg="python $DECOUPLED_WBC_DIR/decoupled_wbc/scripts/deploy_g1.py"
+alias rsl="python $DECOUPLED_WBC_DIR/decoupled_wbc/control/main/teleop/run_sim_loop.py"
+alias rgcl="python $DECOUPLED_WBC_DIR/decoupled_wbc/control/main/teleop/run_g1_control_loop.py"
+alias rtpl="python $DECOUPLED_WBC_DIR/decoupled_wbc/control/main/teleop/run_teleop_policy_loop.py"
+alias tgcl="pytest $DECOUPLED_WBC_DIR/decoupled_wbc/tests/control/main/teleop/test_g1_control_loop.py -s"

GR00T-WholeBodyControl/decoupled_wbc/docker/.tmux.conf

@@ -0,0 +1,84 @@
+# Enable mouse mode
+set -g mouse on
+
+# Start window numbering at 0 (default)
+set -g base-index 0
+setw -g pane-base-index 0
+
+# Increase scrollback buffer size
+set -g history-limit 50000
+
+# Use Alt-arrow keys without prefix key to switch panes
+bind -n M-Left select-pane -L
+bind -n M-Right select-pane -R
+bind -n M-Up select-pane -U
+bind -n M-Down select-pane -D
+
+# Use Alt-1,2,3... to switch windows
+bind -n M-1 select-window -t 1
+bind -n M-2 select-window -t 2
+bind -n M-3 select-window -t 3
+bind -n M-4 select-window -t 4
+bind -n M-5 select-window -t 5
+bind -n M-6 select-window -t 6
+bind -n M-7 select-window -t 7
+bind -n M-8 select-window -t 8
+bind -n M-9 select-window -t 9
+
+# Split panes using Alt-| and Alt--
+bind -n M-| split-window -h
+bind -n M-- split-window -v
+
+# Easy config reload
+bind -n M-r source-file ~/.tmux.conf \; display-message "Config reloaded!"
+
+# Status bar customization
+set -g status-style bg=colour240,fg=colour255
+set -g status-left "#[fg=colour255,bg=colour240] #S #[fg=colour240,bg=colour238]"
+set -g status-right "#[fg=colour255,bg=colour240] %H:%M #[fg=colour240,bg=colour238]"
+
+# Window status format
+setw -g window-status-format "#[fg=colour255,bg=colour238] #I:#W "
+setw -g window-status-current-format "#[fg=colour238,bg=colour255]#[fg=colour238,bg=colour255] #I:#W #[fg=colour255,bg=colour238]"
+
+# Pane border colors
+set -g pane-border-style fg=colour240
+set -g pane-active-border-style fg=colour255
+
+# Message text
+set -g message-style bg=colour238,fg=colour255
+
+# Clock mode
+setw -g clock-mode-colour colour255
+
+# Enable focus events
+set -g focus-events on
+
+# Increase escape time
+set -sg escape-time 0
+
+# Enable true color support
+set -ga terminal-overrides ",*256col*:Tc"
+
+# Set default terminal mode to 256 colors
+set -g default-terminal "screen-256color"
+
+# Display a message when a window is created
+set -g display-time 4000
+
+# Automatically set window title
+setw -g automatic-rename on
+set -g set-titles on
+set -g set-titles-string "#T"
+
+# Enable clipboard integration
+set -g @plugin 'tmux-plugins/tmux-yank'
+
+# List of plugins
+set -g @plugin 'tmux-plugins/tpm'
+set -g @plugin 'tmux-plugins/tmux-sensible'
+set -g @plugin 'tmux-plugins/tmux-resurrect'
+set -g @plugin 'tmux-plugins/tmux-continuum'
+
+# Initialize TMUX plugin manager (keep this line at the very bottom of tmux.conf)
+run '~/.tmux/plugins/tpm/tpm' 

GR00T-WholeBodyControl/decoupled_wbc/docker/70-manus-hid.rules

@@ -0,0 +1,5 @@
+  # HIDAPI/libusb
+  SUBSYSTEMS=="usb", ATTRS{idVendor}=="3325", MODE:="0666"
+
+  # HIDAPI/hidraw
+  KERNEL=="hidraw*", ATTRS{idVendor}=="3325", MODE:="0666"

GR00T-WholeBodyControl/decoupled_wbc/docker/Dockerfile.deploy

@@ -0,0 +1,130 @@
+FROM nvgear/ros-2:latest
+
+# Accept build argument for username
+ARG USERNAME
+ARG USERID
+ARG HOME_DIR
+ARG WORKTREE_NAME
+
+# Create user with the same name as host
+RUN if [ "$USERID" != "0" ]; then \
+        useradd -m -u ${USERID} -s /bin/bash ${USERNAME} && \
+        echo "${USERNAME} ALL=(ALL) NOPASSWD: ALL" >> /etc/sudoers && \
+        # Add user to video and render groups for GPU access
+        usermod -a -G video,render ${USERNAME} || true; \
+    fi
+
+# Copy .bashrc with color settings before switching user
+COPY --chown=${USERNAME}:${USERNAME} decoupled_wbc/docker/.bashrc ${HOME_DIR}/.bashrc
+
+# Install Manus udev rules
+COPY --chown=${USERNAME}:${USERNAME} decoupled_wbc/docker/70-manus-hid.rules /etc/udev/rules.d/70-manus-hid.rules
+
+# Copy tmux configuration
+COPY --chown=${USERNAME}:${USERNAME} decoupled_wbc/docker/.tmux.conf ${HOME_DIR}/.tmux.conf
+
+# Switch to user
+USER ${USERNAME}
+
+# Install tmux plugin manager and uv in parallel
+RUN git clone https://github.com/tmux-plugins/tpm ${HOME_DIR}/.tmux/plugins/tpm & \
+    curl -LsSf https://astral.sh/uv/install.sh | env UV_INSTALL_DIR=${HOME_DIR}/.cargo/bin sh & \
+    wait
+
+# Install tmux plugins automatically
+RUN ${HOME_DIR}/.tmux/plugins/tpm/bin/install_plugins || true
+
+# Add uv to PATH
+ENV PATH="${HOME_DIR}/.cargo/bin:$PATH"
+ENV UV_PYTHON=${HOME_DIR}/venv/bin/python
+
+# Create venv
+RUN uv venv --python 3.10 ${HOME_DIR}/venv
+
+# Install hardware-specific packages (x86 only - not available on ARM64/Orin)
+USER root
+COPY --chown=${USERNAME}:${USERNAME} decoupled_wbc/control/teleop/device/pico/XRoboToolkit_PC_Service_1.0.0_ubuntu_22.04_amd64.deb ${HOME_DIR}/XRoboToolkit_PC_Service_1.0.0_ubuntu_22.04_amd64.deb
+COPY --chown=${USERNAME}:${USERNAME} decoupled_wbc/control/teleop/device/pico/roboticsservice_1.0.0.0_arm64.deb ${HOME_DIR}/roboticsservice_1.0.0.0_arm64.deb
+
+RUN if [ "$(dpkg --print-architecture)" = "amd64" ]; then \
+        # Ultra Leap setup
+        wget -qO - https://repo.ultraleap.com/keys/apt/gpg | gpg --dearmor | tee /etc/apt/trusted.gpg.d/ultraleap.gpg && \
+        echo 'deb [arch=amd64] https://repo.ultraleap.com/apt stable main' | tee /etc/apt/sources.list.d/ultraleap.list && \
+        apt-get update && \
+        echo "yes" | DEBIAN_FRONTEND=noninteractive apt-get install -y ultraleap-hand-tracking libhidapi-dev && \
+        # Space Mouse udev rules
+        echo 'KERNEL=="hidraw*", SUBSYSTEM=="hidraw", MODE="0664", GROUP="plugdev"' > /etc/udev/rules.d/99-hidraw-permissions.rules && \
+        usermod -aG plugdev ${USERNAME}; \
+        # Pico setup
+        apt-get install -y xdg-utils && \
+        dpkg -i ${HOME_DIR}/XRoboToolkit_PC_Service_1.0.0_ubuntu_22.04_amd64.deb; \
+    else \
+        echo "Skipping x86-only hardware packages on $(dpkg --print-architecture)"; \
+    fi
+
+USER ${USERNAME}
+# Install hardware Python packages (x86 only) with caching
+RUN --mount=type=cache,target=${HOME_DIR}/.cache/uv,uid=${USERID},gid=${USERID} \
+    if [ "$(dpkg --print-architecture)" = "amd64" ]; then \
+        # Ultra Leap Python bindings
+        git clone https://github.com/ultraleap/leapc-python-bindings ${HOME_DIR}/leapc-python-bindings && \
+        cd ${HOME_DIR}/leapc-python-bindings && \
+        UV_CONCURRENT_DOWNLOADS=8 uv pip install -r requirements.txt && \
+        MAKEFLAGS="-j$(nproc)" ${HOME_DIR}/venv/bin/python -m build leapc-cffi && \
+        uv pip install leapc-cffi/dist/leapc_cffi-0.0.1.tar.gz && \
+        uv pip install -e leapc-python-api && \
+        # Space Mouse Python package
+        uv pip install pyspacemouse && \
+        # Pico Python bindings
+        git clone https://github.com/XR-Robotics/XRoboToolkit-PC-Service-Pybind.git ${HOME_DIR}/XRoboToolkit-PC-Service-Pybind && \
+        cd ${HOME_DIR}/XRoboToolkit-PC-Service-Pybind && \
+        uv pip install setuptools pybind11 && \
+        sed -i "s|pip install|uv pip install|g" setup_ubuntu.sh && \
+        sed -i "s|pip uninstall|uv pip uninstall|g" setup_ubuntu.sh && \
+        sed -i "s|python setup.py install|${HOME_DIR}/venv/bin/python setup.py install|g" setup_ubuntu.sh && \
+        bash setup_ubuntu.sh; \
+    fi
+
+# Install Python dependencies using uv with caching
+RUN --mount=type=cache,target=${HOME_DIR}/.cache/uv,uid=${USERID},gid=${USERID} \
+    UV_CONCURRENT_DOWNLOADS=8 uv pip install --upgrade pip ipython jupyter notebook debugpy
+
+
+# Copy entire project to the workspace directory where it will be mounted at runtime
+# NOTE: The build context must be the project root for this to work
+# Use dynamic worktree name to match runtime mount path
+COPY --chown=${USERNAME}:${USERNAME} . ${HOME_DIR}/Projects/${WORKTREE_NAME}
+
+# Install Python dependencies inside the venv with caching - split into separate commands
+RUN --mount=type=cache,target=${HOME_DIR}/.cache/uv,uid=${USERID},gid=${USERID} \
+    UV_CONCURRENT_DOWNLOADS=8 uv pip install \
+    -e ${HOME_DIR}/Projects/${WORKTREE_NAME}/external_dependencies/unitree_sdk2_python
+
+# Unlike pip, uv downloads LFS files by default. There's a bug in uv that causes LFS files
+# to fail to download (https://github.com/astral-sh/uv/issues/3312). So we need to set
+# UV_GIT_LFS=1 to prevent uv from downloading LFS files.
+# Install project packages (decoupled_wbc + gear_sonic) with caching
+RUN --mount=type=cache,target=${HOME_DIR}/.cache/uv,uid=${USERID},gid=${USERID} \
+    GIT_LFS_SKIP_SMUDGE=1 UV_CONCURRENT_DOWNLOADS=8 uv pip install \
+    -e "${HOME_DIR}/Projects/${WORKTREE_NAME}/decoupled_wbc[full,dev]" \
+    -e "${HOME_DIR}/Projects/${WORKTREE_NAME}/gear_sonic[sim]"
+
+# Clone and install robosuite with specific branch
+RUN --mount=type=cache,target=${HOME_DIR}/.cache/uv,uid=${USERID},gid=${USERID} \
+    git clone https://github.com/xieleo5/robosuite.git ${HOME_DIR}/robosuite && \
+    cd ${HOME_DIR}/robosuite && \
+    git checkout leo/support_g1_locomanip && \
+    UV_CONCURRENT_DOWNLOADS=8 uv pip install -e .
+
+# Install gr00trobocasa
+RUN --mount=type=cache,target=${HOME_DIR}/.cache/uv,uid=${USERID},gid=${USERID} \
+    UV_CONCURRENT_DOWNLOADS=8 uv pip install -e ${HOME_DIR}/Projects/${WORKTREE_NAME}/decoupled_wbc/dexmg/gr00trobocasa
+
+# Configure bash environment with virtual environment and ROS2 setup
+RUN echo "source ${HOME_DIR}/venv/bin/activate" >> ${HOME_DIR}/.bashrc && \
+    echo "source /opt/ros/humble/setup.bash" >> ${HOME_DIR}/.bashrc && \
+    echo "export ROS_LOCALHOST_ONLY=1" >> ${HOME_DIR}/.bashrc && \
+    echo "export PYTHONPATH=${HOME_DIR}/Projects/${WORKTREE_NAME}:\${PYTHONPATH}" >> ${HOME_DIR}/.bashrc
+
+# Default command (can be overridden at runtime)
+CMD ["/bin/bash"]