sim/base_sim.py

import argparse
import logging
import pathlib
from pathlib import Path
import threading
from threading import Thread
from typing import Dict

import mujoco
import mujoco.viewer

import numpy as np
try:
    import rclpy
    HAS_RCLPY = True
except ImportError:
    HAS_RCLPY = False
    print("ROS 2 integration unavailable; camera images use the ZMQ publisher.")
from unitree_sdk2py.core.channel import ChannelFactoryInitialize
import yaml
import os
from .image_publish_utils import ImagePublishProcess
from .keyboard_forward import forward_key
from .metric_utils import check_contact
from .sim_utils import get_subtree_body_names
from .unitree_sdk2py_bridge import ElasticBand, UnitreeSdk2Bridge
from .mjcf.compose import HAND_COLLISION_BIT
from .model_config import select_end_effector

# AmazingHand servo targets (deg) of the open hand, as lerobot's AMAZING_HAND open_q: (servo 1, servo 2) per finger.
AMAZING_HAND_OPEN_DEG = (-35.0, 35.0)
WORLD_EDGE_MARGIN = 0.05  # m from a world object's edge to the table's edge, whatever WORLD_RANDOMIZE
AMAZING_HAND_CLOSED_DEG = (60.0, -60.0)
# Kinematic grasp: closure thresholds and the zone (min, max corners, right hand) in the {side}_tcp frame between
# the palm face and the closed fingertips; the left hand mirrors y.
GRASP_ATTACH_CLOSURE = 0.6
GRASP_RELEASE_CLOSURE = 0.4
GRASP_ZONE = (np.array([0.03, 0.0, -0.07]), np.array([0.11, 0.07, 0.05]))
logger = logging.getLogger(__name__)
GR00T_WBC_ROOT = Path(__file__).resolve().parent.parent  # Points to mujoco_sim_g1/


def amazing_hand_closure(servo_q) -> float:
    """Closure 0..1 of an AmazingHand's fingers 1-3 from their six servo angles (rad), as lerobot's q_to_closure."""
    q = np.degrees(np.asarray(servo_q, float))
    opened = np.tile(AMAZING_HAND_OPEN_DEG, 3)
    closed = np.tile(AMAZING_HAND_CLOSED_DEG, 3)
    return float(np.clip(np.mean((q - opened) / (closed - opened)), 0.0, 1.0))


class DefaultEnv:
    """Base environment class that handles simulation environment setup and step"""

    def __init__(
        self,
        config: Dict[str, any],
        env_name: str = "default",
        camera_configs: Dict[str, any] = None,
        onscreen: bool = False,
        offscreen: bool = False,
    ):
        # Avoid mutable default argument gotcha
        if camera_configs is None:
            camera_configs = {}
        
        # global_view is only set up for this specifc scene for now.
        if config["ROBOT_SCENE"] == "gr00t_wbc/control/robot_model/model_data/g1/scene_29dof.xml":
            camera_configs["global_view"] = {
                "height": 400,
                "width": 400,
            }
        self.config = config
        self.env_name = env_name
        self.num_body_dof = self.config["NUM_JOINTS"]
        self.num_hand_dof = self.config["NUM_HAND_JOINTS"]
        self.sim_dt = self.config["SIMULATE_DT"]
        self.obs = None
        self.torque_limit = np.array(self.config["motor_effort_limit_list"])
        self.camera_configs = camera_configs
        
        # Debug: print camera config
        if len(camera_configs) > 0:
            print(f"✓ DefaultEnv initialized with {len(camera_configs)} camera(s): {list(camera_configs.keys())}")


        # Unitree bridge will be initialized by the simulator
        self.unitree_bridge = None

        # Store display mode
        self.onscreen = onscreen

        # Initialize scene (defined in subclasses)
        self.init_scene()

        # Setup offscreen rendering if needed (lazy init - renderers created on first use)
        self.offscreen = offscreen
        self.renderers = {}  # Will be lazily initialized
        self._renderers_initialized = False
        self.image_dt = self.config.get("IMAGE_DT", 0.033333)
        
        # Image publishing subprocess (initialized separately)
        self.image_publish_process = None

    def init_scene(self):
        """Initialize the default robot scene"""
        assets_root = Path(__file__).parent.parent
        self.mj_model = mujoco.MjModel.from_xml_path(
            str(assets_root / self.config["ROBOT_SCENE"])
        )
        self.mj_data = mujoco.MjData(self.mj_model)
        # Set valid floating base quaternion (MjData initializes qpos to zeros)
        self.mj_data.qpos[3:7] = [1.0, 0.0, 0.0, 0.0]
        self.mj_model.opt.timestep = self.sim_dt
        self.torso_index = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, "torso_link")
        self.root_body = "pelvis"
        # Enable the elastic band
        if self.config["ENABLE_ELASTIC_BAND"]:
            self.elastic_band = ElasticBand()
            if "g1" in self.config["ROBOT_TYPE"]:
                if self.config["enable_waist"]:
                    self.band_attached_link = self.mj_model.body("pelvis").id
                else:
                    self.band_attached_link = self.mj_model.body("torso_link").id
            elif "h1" in self.config["ROBOT_TYPE"]:
                self.band_attached_link = self.mj_model.body("torso_link").id
            else:
                self.band_attached_link = self.mj_model.body("base_link").id

            if self.onscreen:
                self.viewer = mujoco.viewer.launch_passive(
                    self.mj_model,
                    self.mj_data,
                    key_callback=self._viewer_key_callback,
                    show_left_ui=False,
                    show_right_ui=False,
                )
            else:
                mujoco.mj_forward(self.mj_model, self.mj_data)
                self.viewer = None
        else:
            if self.onscreen:
                self.viewer = mujoco.viewer.launch_passive(
                    self.mj_model, self.mj_data, show_left_ui=False, show_right_ui=False
                )
            else:
                mujoco.mj_forward(self.mj_model, self.mj_data)
                self.viewer = None

        if self.viewer:
            # viewer camera
            self.viewer.cam.azimuth = 120  # Horizontal rotation in degrees
            self.viewer.cam.elevation = -30  # Vertical tilt in degrees
            self.viewer.cam.distance = 2.0  # Distance from camera to target
            self.viewer.cam.lookat = np.array([0, 0, 0.5])  # Point the camera is looking at

        # Body DDS indices exclude the end effectors. Map each scalar joint
        # explicitly: actuator order need not match the model's joint order.
        # AmazingHand fingers are not DDS-managed like dex1/dex3 (a later change wires them
        # to their own control path), so they are excluded from this DDS hand scan entirely.
        is_dds_hand = self.config.get("END_EFFECTOR") != "amazing_hand"
        actuated_joint_ids = set(self.mj_model.actuator_trnid[
            self.mj_model.actuator_trntype == mujoco.mjtTrn.mjTRN_JOINT, 0
        ].tolist())
        self.body_joint_index = []
        self.left_hand_index = []
        self.right_hand_index = []
        for i in range(self.mj_model.njnt):
            name = self.mj_model.joint(i).name
            # A world object's free joint is never the robot's, whatever its name (e.g. `chips_can` has "hip").
            if self.mj_model.jnt_type[i] == mujoco.mjtJoint.mjJNT_FREE:
                continue
            if any(
                [
                    part_name in name
                    for part_name in ["hip", "knee", "ankle", "waist", "shoulder", "elbow", "wrist"]
                ]
            ):
                self.body_joint_index.append(i)
            elif not is_dds_hand or i not in actuated_joint_ids:
                continue
            elif "left_hand" in name or name.startswith("left_dex1_finger_joint_"):
                self.left_hand_index.append(i)
            elif "right_hand" in name or name.startswith("right_dex1_finger_joint_"):
                self.right_hand_index.append(i)

        assert len(self.body_joint_index) == self.config["NUM_JOINTS"], \
            f"Expected {self.config['NUM_JOINTS']} body joints, got {len(self.body_joint_index)}"
        expected_hands = self.config.get("NUM_HAND_JOINTS", 0)
        if len(self.left_hand_index) != expected_hands or len(self.right_hand_index) != expected_hands:
            raise ValueError(f"Expected {expected_hands} joints per end effector, got left={len(self.left_hand_index)}, right={len(self.right_hand_index)}")

        for prefix, attribute in (("body", "body_joint_index"), ("left_hand", "left_hand_index"), ("right_hand", "right_hand_index")):
            joint_ids = np.asarray(getattr(self, attribute), dtype=int)
            setattr(self, attribute, joint_ids)
            setattr(self, prefix + "_qpos_index", self.mj_model.jnt_qposadr[joint_ids])
            setattr(self, prefix + "_dof_index", self.mj_model.jnt_dofadr[joint_ids])
            actuator_ids = []
            for joint_id in joint_ids:
                matches = np.flatnonzero(
                    (self.mj_model.actuator_trntype == mujoco.mjtTrn.mjTRN_JOINT)
                    & (self.mj_model.actuator_trnid[:, 0] == joint_id)
                )
                if len(matches) != 1:
                    raise ValueError(f"Expected one actuator for {self.mj_model.joint(joint_id).name}, got {len(matches)}")
                actuator_ids.append(matches[0])
            setattr(self, prefix + "_actuator_index", np.asarray(actuator_ids, dtype=int))
        self.dds_actuator_index = np.concatenate(
            (self.body_actuator_index, self.left_hand_actuator_index, self.right_hand_actuator_index)
        ).astype(int)
        self.torques = np.zeros(self.mj_model.nu)
        if self.config.get("FREE_BASE", False):
            self.torque_limit = np.concatenate((np.zeros(6), self.torque_limit))
        if self.torque_limit.shape[0] < self.torques.shape[0]:
            # motor_effort_limit_list only covers the DDS-managed actuators (body + dex1/dex3
            # hands). Non-DDS actuators (e.g. the D455 pan/tilt head, AmazingHand fingers) are
            # never written from self.torques (see dds_actuator_index below), so their clip
            # bound is a don't-care; pad with +inf rather than reordering the DDS-index layout.
            pad = self.torques.shape[0] - self.torque_limit.shape[0]
            self.torque_limit = np.concatenate((self.torque_limit, np.full(pad, np.inf)))
        if self.torque_limit.shape != self.torques.shape:
            raise ValueError("motor_effort_limit_list must match the scene's actuator count")
        base_body = self.mj_model.body("pelvis").id
        self.world_object_joints = np.asarray(
            [
                i
                for i in range(self.mj_model.njnt)
                if self.mj_model.jnt_type[i] == mujoco.mjtJoint.mjJNT_FREE
                and self.mj_model.jnt_bodyid[i] != base_body
            ],
            dtype=int,
        )
        object_addresses = self.mj_model.jnt_qposadr[self.world_object_joints]
        self.world_object_qpos0 = self.mj_model.qpos0[object_addresses[:, None] + np.arange(7)]
        self.world_rng = np.random.default_rng()
        self._init_grasp()
        self.hand_open_state = self._settle_open_hands() if self.config.get("WORLD") else None
        # Jittered objects keep their edge this far inside the world's support (a box geom named by world.yaml).
        self.world_support = None
        support = self.config.get("WORLD_SUPPORT")
        if support:
            geom = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, support)
            if geom < 0 or self.mj_model.geom_type[geom] != mujoco.mjtGeom.mjGEOM_BOX:
                raise ValueError(f"World support {support!r} is not a box geom of {self.config['ROBOT_SCENE']}")
            # In the world frame: the support's body can be placed and turned anywhere (e.g. the table_room tables).
            data = mujoco.MjData(self.mj_model)
            mujoco.mj_kinematics(self.mj_model, data)
            rot = data.geom_xmat[geom].reshape(3, 3)[:2, :2]
            half = np.abs(rot) @ self.mj_model.geom_size[geom, :2]
            self.world_support = (data.geom_xpos[geom, :2].copy(), half)
        for name in self.camera_configs:
            if mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_CAMERA, name) < 0:
                raise ValueError(f"Camera {name!r} does not exist in {self.config['ROBOT_SCENE']}")
        self.reset()

    def init_renderers(self):
        # Initialize camera renderers
        self.renderers = {}
        for camera_name, camera_config in self.camera_configs.items():
            renderer = mujoco.Renderer(
                self.mj_model, height=camera_config["height"], width=camera_config["width"]
            )
            self.renderers[camera_name] = renderer
    
    def start_image_publish_subprocess(self, start_method: str = "spawn", camera_port: int = 5555):
        """Start image publishing subprocess using ZMQ"""
        # Use spawn method for better GIL isolation, or configured method
        if len(self.camera_configs) == 0:
            print(
                "Warning: No camera configs provided, image publishing subprocess will not be started"
            )
            return
        start_method = self.config.get("MP_START_METHOD", "spawn")
        self.image_publish_process = ImagePublishProcess(
            camera_configs=self.camera_configs,
            image_dt=self.image_dt,
            zmq_port=camera_port,
            start_method=start_method,
            verbose=self.config.get("verbose", False),
        )
        self.image_publish_process.start_process()
        print(f"✓ Started image publishing subprocess on ZMQ port {camera_port}")

    def compute_body_torques(self) -> np.ndarray:
        """Compute body torques based on the current robot state"""
        body_torques = np.zeros(self.num_body_dof)
        if self.unitree_bridge is not None and self.unitree_bridge.low_cmd:
            # DDS command slots may be sparse (see UnitreeSdk2Bridge.joint_slots); index i is
            # the MuJoCo actuator order, joint_slots[i] is the matching DDS motor_cmd slot.
            for i, slot in enumerate(self.unitree_bridge.joint_slots):
                if self.unitree_bridge.use_sensor:
                    body_torques[i] = (
                        self.unitree_bridge.low_cmd.motor_cmd[slot].tau
                        + self.unitree_bridge.low_cmd.motor_cmd[slot].kp
                        * (self.unitree_bridge.low_cmd.motor_cmd[slot].q - self.mj_data.sensordata[i])
                        + self.unitree_bridge.low_cmd.motor_cmd[slot].kd
                        * (
                            self.unitree_bridge.low_cmd.motor_cmd[slot].dq
                            - self.mj_data.sensordata[i + self.unitree_bridge.num_body_motor]
                        )
                    )
                else:
                    body_torques[i] = (
                        self.unitree_bridge.low_cmd.motor_cmd[slot].tau
                        + self.unitree_bridge.low_cmd.motor_cmd[slot].kp
                        * (
                            self.unitree_bridge.low_cmd.motor_cmd[slot].q
                            - self.mj_data.qpos[self.body_qpos_index[i]]
                        )
                        + self.unitree_bridge.low_cmd.motor_cmd[slot].kd
                        * (
                            self.unitree_bridge.low_cmd.motor_cmd[slot].dq
                            - self.mj_data.qvel[self.body_dof_index[i]]
                        )
                    )
        return body_torques

    def compute_hand_torques(self) -> np.ndarray:
        """Compute hand torques based on the current robot state"""
        left_hand_torques = np.zeros(self.num_hand_dof)
        right_hand_torques = np.zeros(self.num_hand_dof)
        if self.unitree_bridge is not None and self.unitree_bridge.low_cmd:
            for i in range(self.unitree_bridge.num_hand_motor):
                left_hand_torques[i] = (
                    self.unitree_bridge.left_hand_cmd.motor_cmd[i].tau
                    + self.unitree_bridge.left_hand_cmd.motor_cmd[i].kp
                    * (
                        self.unitree_bridge.left_hand_cmd.motor_cmd[i].q
                        - self.mj_data.qpos[self.left_hand_qpos_index[i]]
                    )
                    + self.unitree_bridge.left_hand_cmd.motor_cmd[i].kd
                    * (
                        self.unitree_bridge.left_hand_cmd.motor_cmd[i].dq
                        - self.mj_data.qvel[self.left_hand_dof_index[i]]
                    )
                )
                right_hand_torques[i] = (
                    self.unitree_bridge.right_hand_cmd.motor_cmd[i].tau
                    + self.unitree_bridge.right_hand_cmd.motor_cmd[i].kp
                    * (
                        self.unitree_bridge.right_hand_cmd.motor_cmd[i].q
                        - self.mj_data.qpos[self.right_hand_qpos_index[i]]
                    )
                    + self.unitree_bridge.right_hand_cmd.motor_cmd[i].kd
                    * (
                        self.unitree_bridge.right_hand_cmd.motor_cmd[i].dq
                        - self.mj_data.qvel[self.right_hand_dof_index[i]]
                    )
                )
        return np.concatenate((left_hand_torques, right_hand_torques))

    def compute_body_qpos(self) -> np.ndarray:
        """Compute body joint positions based on the current command"""
        body_qpos = np.zeros(self.num_body_dof)
        if self.unitree_bridge is not None and self.unitree_bridge.low_cmd:
            for i, slot in enumerate(self.unitree_bridge.joint_slots):
                body_qpos[i] = self.unitree_bridge.low_cmd.motor_cmd[slot].q
        return body_qpos

    def compute_hand_qpos(self) -> np.ndarray:
        """Compute hand joint positions based on the current command"""
        hand_qpos = np.zeros(self.num_hand_dof * 2)
        if self.unitree_bridge is not None and self.unitree_bridge.low_cmd:
            for i in range(self.unitree_bridge.num_hand_motor):
                hand_qpos[i] = self.unitree_bridge.left_hand_cmd.motor_cmd[i].q
                hand_qpos[i + self.num_hand_dof] = self.unitree_bridge.right_hand_cmd.motor_cmd[i].q
        return hand_qpos

    def prepare_obs(self) -> Dict[str, any]:
        """Prepare observation dictionary from the current robot state"""
        obs = {}
        obs["floating_base_pose"] = self.mj_data.qpos[:7]
        obs["floating_base_vel"] = self.mj_data.qvel[:6]
        obs["floating_base_acc"] = self.mj_data.qacc[:6]
        obs["secondary_imu_quat"] = self.mj_data.xquat[self.torso_index]
        obs["secondary_imu_vel"] = self.mj_data.cvel[self.torso_index]
        obs["body_q"] = self.mj_data.qpos[self.body_qpos_index]
        obs["body_dq"] = self.mj_data.qvel[self.body_dof_index]
        obs["body_ddq"] = self.mj_data.qacc[self.body_dof_index]
        obs["body_tau_est"] = self.mj_data.actuator_force[self.body_actuator_index]
        if self.num_hand_dof > 0:
            obs["left_hand_q"] = self.mj_data.qpos[self.left_hand_qpos_index]
            obs["left_hand_dq"] = self.mj_data.qvel[self.left_hand_dof_index]
            obs["left_hand_ddq"] = self.mj_data.qacc[self.left_hand_dof_index]
            obs["left_hand_tau_est"] = self.mj_data.actuator_force[self.left_hand_actuator_index]
            obs["right_hand_q"] = self.mj_data.qpos[self.right_hand_qpos_index]
            obs["right_hand_dq"] = self.mj_data.qvel[self.right_hand_dof_index]
            obs["right_hand_ddq"] = self.mj_data.qacc[self.right_hand_dof_index]
            obs["right_hand_tau_est"] = self.mj_data.actuator_force[self.right_hand_actuator_index]
        obs["time"] = self.mj_data.time
        return obs

    def sim_step(self):
        self.obs = self.prepare_obs()
        self.unitree_bridge.PublishLowState(self.obs)
        if self.unitree_bridge.joystick:
            self.unitree_bridge.PublishWirelessController()
        if self.config["ENABLE_ELASTIC_BAND"]:
            if self.elastic_band.enable:
                # Get Cartesian pose and velocity of the band_attached_link
                pose = np.concatenate(
                    [
                        self.mj_data.xpos[self.band_attached_link],  # link position in world
                        self.mj_data.xquat[
                            self.band_attached_link
                        ],  # link quaternion in world [w,x,y,z]
                        np.zeros(6),  # placeholder for velocity
                    ]
                )

                # Get velocity in world frame
                mujoco.mj_objectVelocity(
                    self.mj_model,
                    self.mj_data,
                    mujoco.mjtObj.mjOBJ_BODY,
                    self.band_attached_link,
                    pose[7:13],
                    0,  # 0 for world frame
                )

                # Reorder velocity from [ang, lin] to [lin, ang]
                pose[7:10], pose[10:13] = pose[10:13], pose[7:10].copy()
                self.mj_data.xfrc_applied[self.band_attached_link] = self.elastic_band.Advance(pose)
            else:
                # explicitly resetting the force when the band is not enabled
                self.mj_data.xfrc_applied[self.band_attached_link] = np.zeros(6)
        body_torques = self.compute_body_torques()
        hand_torques = self.compute_hand_torques()
        self.torques[self.body_actuator_index] = body_torques
        if self.num_hand_dof > 0:
            self.torques[self.left_hand_actuator_index] = hand_torques[: self.num_hand_dof]
            self.torques[self.right_hand_actuator_index] = hand_torques[self.num_hand_dof :]

        self.torques = np.clip(self.torques, -self.torque_limit, self.torque_limit)

        # Only the DDS-managed slots (body + dex1/dex3 hands) are written here; any other
        # actuator (e.g. the ZMQ-driven head/AmazingHand pan-tilt/fingers) keeps whatever
        # ctrl another writer (SimHeadHandDevice) has set, instead of being reset to zero.
        self.mj_data.ctrl[self.dds_actuator_index] = self.torques[self.dds_actuator_index]
        mujoco.mj_step(self.mj_model, self.mj_data)
        self._update_grasp()
        # self.check_self_collision()

    def kinematics_step(self):
        """
        Run kinematics only: compute the qpos of the robot and directly set the qpos.
        For debugging purposes.
        """
        if self.unitree_bridge is not None:
            self.unitree_bridge.PublishLowState(self.prepare_obs())
            if self.unitree_bridge.joystick:
                self.unitree_bridge.PublishWirelessController()

        if self.config["ENABLE_ELASTIC_BAND"]:
            if self.elastic_band.enable:
                # Get Cartesian pose and velocity of the band_attached_link
                pose = np.concatenate(
                    [
                        self.mj_data.xpos[self.band_attached_link],  # link position in world
                        self.mj_data.xquat[
                            self.band_attached_link
                        ],  # link quaternion in world [w,x,y,z]
                        np.zeros(6),  # placeholder for velocity
                    ]
                )

                # Get velocity in world frame
                mujoco.mj_objectVelocity(
                    self.mj_model,
                    self.mj_data,
                    mujoco.mjtObj.mjOBJ_BODY,
                    self.band_attached_link,
                    pose[7:13],
                    0,  # 0 for world frame
                )

                # Reorder velocity from [ang, lin] to [lin, ang]
                pose[7:10], pose[10:13] = pose[10:13], pose[7:10].copy()

                self.mj_data.xfrc_applied[self.band_attached_link] = self.elastic_band.Advance(pose)
            else:
                # explicitly resetting the force when the band is not enabled
                self.mj_data.xfrc_applied[self.band_attached_link] = np.zeros(6)

        body_qpos = self.compute_body_qpos()  # (num_body_dof,)
        hand_qpos = self.compute_hand_qpos()  # (num_hand_dof * 2,)

        self.mj_data.qpos[self.body_qpos_index] = body_qpos
        self.mj_data.qpos[self.left_hand_qpos_index] = hand_qpos[: self.num_hand_dof]
        self.mj_data.qpos[self.right_hand_qpos_index] = hand_qpos[self.num_hand_dof :]

        mujoco.mj_kinematics(self.mj_model, self.mj_data)
        mujoco.mj_comPos(self.mj_model, self.mj_data)

    def apply_perturbation(self, key):
        """Apply perturbation to the robot"""
        # Add velocity perturbations in body frame
        perturbation_x_body = 0.0  # forward/backward in body frame
        perturbation_y_body = 0.0  # left/right in body frame
        if key == "up":
            perturbation_x_body = 1.0  # forward
        elif key == "down":
            perturbation_x_body = -1.0  # backward
        elif key == "left":
            perturbation_y_body = 1.0  # left
        elif key == "right":
            perturbation_y_body = -1.0  # right

        # Transform body frame velocity to world frame using MuJoCo's rotation
        vel_body = np.array([perturbation_x_body, perturbation_y_body, 0.0])
        vel_world = np.zeros(3)
        base_quat = self.mj_data.qpos[3:7]  # [w, x, y, z] quaternion

        # Use MuJoCo's robust quaternion rotation (handles invalid quaternions automatically)
        mujoco.mju_rotVecQuat(vel_world, vel_body, base_quat)

        # Apply to base linear velocity in world frame
        self.mj_data.qvel[0] += vel_world[0]  # world X velocity
        self.mj_data.qvel[1] += vel_world[1]  # world Y velocity

        # Update dynamics after velocity change
        mujoco.mj_forward(self.mj_model, self.mj_data)

    def _viewer_key_callback(self, key):
        """Viewer keys drive the elastic band (7/8/9) and are forwarded to the keyboard teleop."""
        self.elastic_band.MujuocoKeyCallback(key)
        forward_key(key)

    def update_viewer(self):
        if self.viewer is not None:
            self.viewer.sync()

    def update_viewer_camera(self):
        if self.viewer is not None:
            if self.viewer.cam.type == mujoco.mjtCamera.mjCAMERA_TRACKING:
                self.viewer.cam.type = mujoco.mjtCamera.mjCAMERA_FREE
            else:
                self.viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING

    def set_unitree_bridge(self, unitree_bridge):
        """Set the unitree bridge from the simulator"""
        self.unitree_bridge = unitree_bridge

    def get_privileged_obs(self):
        """Get privileged observation. Should be implemented by subclasses."""
        return {}

    def update_render_caches(self):
        """Update render cache and shared memory for subprocess."""
        # Lazy init renderers on first call (creates OpenGL context in calling thread)
        if not self._renderers_initialized and self.offscreen:
            self.init_renderers()
            self._renderers_initialized = True
            print(f"✓ Renderers initialized lazily in thread {__import__('threading').current_thread().name}")
        
        render_caches = {}
        for camera_name, camera_config in self.camera_configs.items():
            renderer = self.renderers.get(camera_name)
            if renderer is None:
                continue
            if "params" in camera_config:
                renderer.update_scene(self.mj_data, camera=camera_config["params"])
            else:
                renderer.update_scene(self.mj_data, camera=camera_name)
            render_caches[camera_name + "_image"] = renderer.render()
        
        # Update shared memory if image publishing process is available
        if self.image_publish_process is not None:
            self.image_publish_process.update_shared_memory(render_caches)

        return render_caches

    def handle_keyboard_button(self, key):
        if self.elastic_band is not None:
            self.elastic_band.handle_keyboard_button(key)

        if key == "backspace":
            self.reset()
        if key == "v":
            self.update_viewer_camera()
        if key in ["up", "down", "left", "right"]:
            self.apply_perturbation(key)

    def check_fall(self):
        """Check if the robot has fallen"""
        self.fall = False
        if self.mj_data.qpos[2] < 0.2:
            self.fall = True
            print(f"Warning: Robot has fallen, height: {self.mj_data.qpos[2]:.3f} m")

        if self.fall:
            self.reset()

    def check_self_collision(self):
        """Check for self-collision of the robot"""
        robot_bodies = get_subtree_body_names(self.mj_model, self.mj_model.body(self.root_body).id)
        self_collision, contact_bodies = check_contact(
            self.mj_model, self.mj_data, robot_bodies, robot_bodies, return_all_contact_bodies=True
        )
        if self_collision:
            print(f"Warning: Self-collision detected: {contact_bodies}")
        return self_collision

    def reset(self):
        mujoco.mj_resetData(self.mj_model, self.mj_data)
        # Set valid floating base quaternion (identity: w=1, x=y=z=0)
        # mj_resetData sets qpos to zeros, which gives invalid [0,0,0,0] quaternion
        self.mj_data.qpos[3:7] = [1.0, 0.0, 0.0, 0.0]
        if self.config.get("END_EFFECTOR") == "dex1":
            for side in ("left", "right"):
                joints = getattr(self, side + "_hand_index")
                addresses = getattr(self, side + "_hand_qpos_index")
                self.mj_data.qpos[addresses] = self.mj_model.jnt_range[joints, 1]
        self._clear_grasp()
        self.reset_world_objects()
        if self.hand_open_state is not None:
            qpos_index, qpos, ctrl_index, ctrl = self.hand_open_state
            self.mj_data.qpos[qpos_index] = qpos
            self.mj_data.ctrl[ctrl_index] = ctrl
        # Propagate qpos to derived quantities (xquat, xpos, etc.)
        mujoco.mj_forward(self.mj_model, self.mj_data)

    def _init_grasp(self):
        """Per-hand state of the world's kinematic grasp: the `{side}_grasp_{body}` welds, finger servos and tcp site."""
        model = self.mj_model
        self.grasp_hands = {}
        self.grasp_state = {}
        self.grasp_geoms = np.zeros(0, int)
        object_geoms = {}
        for side, first in (("right", 1), ("left", 11)):
            welds = {
                int(model.eq_obj2id[i]): i
                for i in range(model.neq)
                if model.eq_type[i] == mujoco.mjtEq.mjEQ_WELD and model.eq(i).name.startswith(f"{side}_grasp_")
            }
            servos = [f"{side}_hand_motor{first + i}_joint" for i in range(6)]
            if not welds or mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, servos[0]) < 0:
                continue
            low, high = (GRASP_ZONE[0].copy(), GRASP_ZONE[1].copy())
            if side == "left":
                low[1], high[1] = -GRASP_ZONE[1][1], -GRASP_ZONE[0][1]
            self.grasp_hands[side] = {
                "welds": welds,
                "held": None,
                "qpos": np.array([model.jnt_qposadr[model.joint(name).id] for name in servos]),
                "site": model.site(f"{side}_tcp").id,
                "zone": (low, high),
            }
            self.grasp_state[side] = False
            for body in welds:
                object_geoms[body] = np.flatnonzero(model.geom_bodyid == body)
        self.grasp_object_geoms = object_geoms
        if object_geoms:
            self.grasp_geoms = np.concatenate(list(object_geoms.values()))
            self.grasp_conaffinity = model.geom_conaffinity[self.grasp_geoms].copy()

    def _clear_grasp(self):
        """Drop every kinematic grasp (reset): welds inactive, hand-object contacts back."""
        for side, hand in self.grasp_hands.items():
            self.grasp_state[side] = False
            hand["held"] = None
            for weld in hand["welds"].values():
                self.mj_data.eq_active[weld] = self.mj_model.eq_active0[weld]
        if self.grasp_hands:
            self.mj_model.geom_conaffinity[self.grasp_geoms] = self.grasp_conaffinity

    def _update_grasp(self):
        """Attach the object to a hand that has closed on it, release it when the hand opens (WORLD_GRASP=attach)."""
        model, data = self.mj_model, self.mj_data
        for side, hand in self.grasp_hands.items():
            closure = amazing_hand_closure(data.qpos[hand["qpos"]])
            if self.grasp_state[side]:
                if closure <= GRASP_RELEASE_CLOSURE:
                    self._release_grasp(side)
                continue
            if closure < GRASP_ATTACH_CLOSURE or any(self.grasp_state.values()):
                continue
            mujoco.mj_kinematics(model, data)
            tcp = data.site_xpos[hand["site"]]
            rotation = data.site_xmat[hand["site"]].reshape(3, 3)
            inside = []
            for body, weld in hand["welds"].items():
                offset = data.xpos[body] - tcp
                local = rotation.T @ offset
                if np.all(local >= hand["zone"][0]) and np.all(local <= hand["zone"][1]):
                    inside.append((float(np.linalg.norm(offset)), weld))
            if inside:
                self._attach_grasp(side, min(inside)[1], closure)

    def _attach_grasp(self, side, weld, closure):
        model, data, hand = self.mj_model, self.mj_data, self.grasp_hands[side]
        body1, body2 = model.eq_obj1id[weld], model.eq_obj2id[weld]
        relpos = np.zeros(3)
        relquat = np.zeros(4)
        inverse = np.zeros(4)
        mujoco.mju_negQuat(inverse, data.xquat[body1])
        mujoco.mju_rotVecQuat(relpos, data.xpos[body2] - data.xpos[body1], inverse)
        mujoco.mju_mulQuat(relquat, inverse, data.xquat[body2])
        model.eq_data[weld, :3] = 0.0
        model.eq_data[weld, 3:6] = relpos
        model.eq_data[weld, 6:10] = relquat
        data.eq_active[weld] = 1
        self._set_object_contacts(body2, hands=False)
        hand["held"] = weld
        self.grasp_state[side] = True
        logger.info("%s hand attached %s (closure %.2f)", side, model.body(body2).name, closure)

    def _release_grasp(self, side):
        hand = self.grasp_hands[side]
        weld = hand["held"]
        self.mj_data.eq_active[weld] = 0
        self._set_object_contacts(int(self.mj_model.eq_obj2id[weld]), hands=True)
        hand["held"] = None
        self.grasp_state[side] = False
        logger.info("%s hand released %s", side, self.mj_model.body(self.mj_model.eq_obj2id[weld]).name)

    def _set_object_contacts(self, body, hands):
        """Let the end effectors' collision bit touch a world object's geoms again, or not."""
        geoms = self.grasp_object_geoms[body]
        original = self.grasp_conaffinity[np.isin(self.grasp_geoms, geoms)]
        self.mj_model.geom_conaffinity[geoms] = original if hands else original & ~HAND_COLLISION_BIT

    def _settle_open_hands(self):
        """The AmazingHand joint state of the open hand, so a world starts with open rather than curled hands.

        Returns None without AmazingHands. The loop-closing passive joints need the hand to be simulated to
        its open pose, with everything else of the robot held.
        """
        model = self.mj_model
        names = [f"{side}_hand_motor{first + i}_joint" for side, first in (("right", 1), ("left", 11)) for i in range(8)]
        if mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, names[0]) < 0:
            return None
        ctrl_index = np.array([model.actuator(name).id for name in names])
        ctrl = np.radians(np.tile(AMAZING_HAND_OPEN_DEG, 8))
        hand_q = np.zeros(model.nq, bool)
        hand_v = np.zeros(model.nv, bool)
        for j in range(model.njnt):
            if "_hand_" in model.joint(j).name:
                size = (4, 3) if model.jnt_type[j] == mujoco.mjtJoint.mjJNT_BALL else (1, 1)
                hand_q[model.jnt_qposadr[j] : model.jnt_qposadr[j] + size[0]] = True
                hand_v[model.jnt_dofadr[j] : model.jnt_dofadr[j] + size[1]] = True
        data = mujoco.MjData(model)
        data.qpos[:] = model.qpos0
        held = data.qpos.copy()
        data.ctrl[ctrl_index] = ctrl
        for _ in range(int(1.5 / model.opt.timestep)):
            mujoco.mj_step(model, data)
            data.qpos[~hand_q] = held[~hand_q]
            data.qvel[~hand_v] = 0.0
        qpos_index = np.flatnonzero(hand_q)
        return qpos_index, data.qpos[qpos_index].copy(), ctrl_index, ctrl

    def reset_world_objects(self):
        """Put every free object of the world back at its spawn pose, jittered by WORLD_RANDOMIZE and WORLD_RANDOMIZE_YAW.

        mj_resetData already restores qpos0, but the pose is set explicitly so it never depends on that.
        """
        jitter = float(self.config.get("WORLD_RANDOMIZE", 0.0) or 0.0)
        max_yaw = np.radians(float(self.config.get("WORLD_RANDOMIZE_YAW", 15.0)))
        for joint, spawn in zip(self.world_object_joints, self.world_object_qpos0):
            address = self.mj_model.jnt_qposadr[joint]
            dof = self.mj_model.jnt_dofadr[joint]
            pose = spawn.copy()
            if jitter > 0:
                pose[:2] += self.world_rng.uniform(-jitter, jitter, 2)
                if self.world_support is not None:
                    geom = np.flatnonzero(self.mj_model.geom_bodyid == self.mj_model.jnt_bodyid[joint])[0]
                    reach = self.mj_model.geom_size[geom, 0] + WORLD_EDGE_MARGIN
                    centre, half = self.world_support
                    pose[:2] = np.clip(pose[:2], centre - half + reach, centre + half - reach)
                yaw = self.world_rng.uniform(-max_yaw, max_yaw)
                mujoco.mju_mulQuat(pose[3:], np.array([np.cos(yaw / 2), 0.0, 0.0, np.sin(yaw / 2)]), spawn[3:])
            self.mj_data.qpos[address : address + 7] = pose
            self.mj_data.qvel[dof : dof + 6] = 0.0


class BaseSimulator:
    """Base simulator class that handles initialization and running of simulations"""

    def __init__(self, config: Dict[str, any], env_name: str = "default", **kwargs):
        config = select_end_effector(config)
        self.config = config
        self.env_name = env_name

        # Initialize ROS 2 node (optional, only if rclpy is available)
        if HAS_RCLPY:
            if not rclpy.ok():
                rclpy.init()
                self.node = rclpy.create_node("sim_mujoco")
                self.thread = threading.Thread(target=rclpy.spin, args=(self.node,), daemon=True)
                self.thread.start()
            else:
                self.thread = None
                executor = rclpy.get_global_executor()
                self.node = executor.get_nodes()[0]  # will only take the first node
        else:
            self.node = None
            self.thread = None

        # Set update frequencies
        self.sim_dt = self.config["SIMULATE_DT"]
        self.image_dt = self.config.get("IMAGE_DT", 0.033333)
        self.viewer_dt = self.config.get("VIEWER_DT", 0.02)

        # Create the environment
        self.sim_env = DefaultEnv(config, env_name, **kwargs)

        # Initialize the DDS communication layer - should be safe to call multiple times

        try:
            if self.config.get("INTERFACE", None):
                ChannelFactoryInitialize(self.config["DOMAIN_ID"], self.config["INTERFACE"])
            else:
                ChannelFactoryInitialize(self.config["DOMAIN_ID"])
        except Exception as e:
            # If it fails because it's already initialized, that's okay
            print(f"Note: Channel factory initialization attempt: {e}")

        # Initialize the unitree bridge and pass it to the environment
        self.init_unitree_bridge()
        self.sim_env.set_unitree_bridge(self.unitree_bridge)

        # Initialize additional components
        self.init_subscriber()
        self.init_publisher()

        self.sim_thread = None

    def start_as_thread(self):
        # Create simulation thread
        self.sim_thread = Thread(target=self.start)
        self.sim_thread.start()
    
    def start_image_publish_subprocess(self, start_method: str = "spawn", camera_port: int = 5555):
        """Start the image publish subprocess"""
        self.sim_env.start_image_publish_subprocess(start_method, camera_port)

    def init_subscriber(self):
        """Initialize subscribers. Can be overridden by subclasses."""
        pass

    def init_publisher(self):
        """Initialize publishers. Can be overridden by subclasses."""
        pass

    def init_unitree_bridge(self):
        """Initialize the unitree SDK bridge and auto-detect joystick."""
        self.unitree_bridge = UnitreeSdk2Bridge(self.config)
        self.unitree_bridge.SetupJoystick(
            device_id=self.config.get("JOYSTICK_DEVICE", 0),
            js_type=self.config.get("JOYSTICK_TYPE", "xbox"),
        )

    def start(self):
        """Main simulation loop"""
        import time
        sim_cnt = 0
        last_time = time.time()

        print(f"Starting simulation loop. Viewer: {self.sim_env.viewer is not None}")
        
        try:
            while (
                self.sim_env.viewer and self.sim_env.viewer.is_running()
            ) or self.sim_env.viewer is None:
                # Run simulation step
                self.sim_env.sim_step()

                # Update viewer at viewer rate
                if sim_cnt % int(self.viewer_dt / self.sim_dt) == 0:
                    self.sim_env.update_viewer()

                # Update render caches at image rate
                if sim_cnt % int(self.image_dt / self.sim_dt) == 0:
                    self.sim_env.update_render_caches()

                # Sleep to maintain correct rate (simple timing without ROS)
                elapsed = time.time() - last_time
                sleep_time = max(0, self.sim_dt - elapsed)
                if sleep_time > 0:
                    time.sleep(sleep_time)
                last_time = time.time()

                sim_cnt += 1
            
            print(f"Loop exited. Viewer running: {self.sim_env.viewer.is_running() if self.sim_env.viewer else 'No viewer'}")
        except KeyboardInterrupt:
            # User pressed Ctrl+C - exit cleanly
            print("Keyboard interrupt received")
            pass
        except Exception as e:
            print(f"Exception in simulation loop: {e}")
            import traceback
            traceback.print_exc()
            self.close()

    def __del__(self):
        """Clean up resources when simulator is deleted"""
        self.close()

    def reset(self):
        """Reset the simulation. Can be overridden by subclasses."""
        self.unitree_bridge.reset()
        self.sim_env.reset()

    def close(self):
        """Close the simulation. Can be overridden by subclasses."""
        try:
            # Stop image publishing subprocess
            if hasattr(self.sim_env, "image_publish_process") and self.sim_env.image_publish_process is not None:
                self.sim_env.image_publish_process.stop()
                self.sim_env.image_publish_process = None
            
            # Close viewer
            if hasattr(self.sim_env, "viewer") and self.sim_env.viewer is not None:
                self.sim_env.viewer.close()

            for renderer in self.sim_env.renderers.values():
                renderer.close()
            self.sim_env.renderers.clear()
            self.sim_env._renderers_initialized = False

            # Shutdown ROS (if available)
            if HAS_RCLPY and rclpy.ok():
                rclpy.shutdown()
        except Exception as e:
            print(f"Warning during close: {e}")

    def get_privileged_obs(self):
        obs = self.sim_env.get_privileged_obs()
        # TODO: add ros2 topic to get privileged obs
        return obs

    def handle_keyboard_button(self, key):
        # Only handles keyboard buttons for default env.
        if self.env_name == "default":
            self.sim_env.handle_keyboard_button(key)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Robot")
    parser.add_argument(
        "--config",
        type=str,
        default="./gr00t_wbc/control/main/teleop/configs/g1_29dof_gear_wbc.yaml",
        help="config file",
    )
    args = parser.parse_args()

    with open(args.config, "r") as file:
        config = yaml.load(file, Loader=yaml.FullLoader)

    if config.get("INTERFACE", None):
        ChannelFactoryInitialize(config["DOMAIN_ID"], config["INTERFACE"])
    else:
        ChannelFactoryInitialize(config["DOMAIN_ID"])

    simulation = BaseSimulator(config)
    simulation.start_as_thread()