gr00t-WholeBodyControl/decoupled_wbc/control/envs/g1/g1_env.py

from copy import deepcopy
from typing import Dict

import gymnasium as gym
import numpy as np
from scipy.spatial.transform import Rotation as R

from decoupled_wbc.control.base.humanoid_env import Hands, HumanoidEnv
from decoupled_wbc.control.envs.g1.g1_body import G1Body
from decoupled_wbc.control.envs.g1.g1_hand import G1ThreeFingerHand
from decoupled_wbc.control.envs.g1.sim.simulator_factory import SimulatorFactory, init_channel
from decoupled_wbc.control.envs.g1.utils.joint_safety import JointSafetyMonitor
from decoupled_wbc.control.robot_model.instantiation.g1 import instantiate_g1_robot_model
from decoupled_wbc.control.robot_model.robot_model import RobotModel
from decoupled_wbc.control.utils.ros_utils import ROSManager


class G1Env(HumanoidEnv):
    def __init__(
        self,
        env_name: str = "default",
        robot_model: RobotModel = None,
        wbc_version: str = "v2",
        config: Dict[str, any] = None,
        **kwargs,
    ):
        super().__init__()
        self.robot_model = deepcopy(robot_model)  # need to cache FK results
        self.config = config

        # Initialize safety monitor (visualization disabled)
        self.safety_monitor = JointSafetyMonitor(
            robot_model, enable_viz=False, env_type=self.config.get("ENV_TYPE", "real")
        )
        self.last_obs = None
        self.last_safety_ok = True  # Track last safety status from queue_action

        init_channel(config=self.config)

        # Initialize body and hands
        self._body = G1Body(config=self.config)

        self.with_hands = config.get("with_hands", False)

        # Gravity compensation settings
        self.enable_gravity_compensation = config.get("enable_gravity_compensation", False)
        self.gravity_compensation_joints = config.get("gravity_compensation_joints", ["arms"])

        if self.enable_gravity_compensation:
            print(
                f"Gravity compensation enabled for joint groups: {self.gravity_compensation_joints}"
            )
        if self.with_hands:
            self._hands = Hands()
            self._hands.left = G1ThreeFingerHand(is_left=True)
            self._hands.right = G1ThreeFingerHand(is_left=False)

        # Initialize simulator if in simulation mode
        self.use_sim = self.config.get("ENV_TYPE") == "sim"

        if self.use_sim:
            # Create simulator using factory

            kwargs.update(
                {
                    "onscreen": self.config.get("ENABLE_ONSCREEN", True),
                    "offscreen": self.config.get("ENABLE_OFFSCREEN", False),
                }
            )
            self.sim = SimulatorFactory.create_simulator(
                config=self.config,
                env_name=env_name,
                wbc_version=wbc_version,
                body_ik_solver_settings_type=kwargs.get("body_ik_solver_settings_type", "default"),
                **kwargs,
            )
        else:
            self.sim = None

            # using the real robot
            self.calibrate_hands()

        # Initialize ROS 2 node
        self.ros_manager = ROSManager(node_name="g1_env")
        self.ros_node = self.ros_manager.node

        self.delay_list = []
        self.visualize_delay = False
        self.print_delay_interval = 100
        self.cnt = 0

    def start_simulator(self):
        # imag epublish disabled since the sim is running in a sub-thread
        SimulatorFactory.start_simulator(self.sim, as_thread=True, enable_image_publish=False)

    def step_simulator(self):
        sim_num_steps = int(self.config["REWARD_DT"] / self.config["SIMULATE_DT"])
        for _ in range(sim_num_steps):
            self.sim.sim_env.sim_step()
        self.sim.sim_env.update_viewer()

    def body(self) -> G1Body:
        return self._body

    def hands(self) -> Hands:
        if not self.with_hands:
            raise RuntimeError(
                "Hands not initialized. Use --with_hands True to enable hand functionality."
            )
        return self._hands

    def observe(self) -> Dict[str, any]:
        # Get observations from body and hands
        body_obs = self.body().observe()

        body_q = body_obs["body_q"]
        body_dq = body_obs["body_dq"]
        body_ddq = body_obs["body_ddq"]
        body_tau_est = body_obs["body_tau_est"]

        if self.with_hands:
            left_hand_obs = self.hands().left.observe()
            right_hand_obs = self.hands().right.observe()
            left_hand_q = left_hand_obs["hand_q"]
            right_hand_q = right_hand_obs["hand_q"]
            left_hand_dq = left_hand_obs["hand_dq"]
            right_hand_dq = right_hand_obs["hand_dq"]
            left_hand_ddq = left_hand_obs["hand_ddq"]
            right_hand_ddq = right_hand_obs["hand_ddq"]
            left_hand_tau_est = left_hand_obs["hand_tau_est"]
            right_hand_tau_est = right_hand_obs["hand_tau_est"]

            # Body and hand joint measurements come in actuator order, so we need to convert them to joint order
            whole_q = self.robot_model.get_configuration_from_actuated_joints(
                body_actuated_joint_values=body_q,
                left_hand_actuated_joint_values=left_hand_q,
                right_hand_actuated_joint_values=right_hand_q,
            )
            whole_dq = self.robot_model.get_configuration_from_actuated_joints(
                body_actuated_joint_values=body_dq,
                left_hand_actuated_joint_values=left_hand_dq,
                right_hand_actuated_joint_values=right_hand_dq,
            )
            whole_ddq = self.robot_model.get_configuration_from_actuated_joints(
                body_actuated_joint_values=body_ddq,
                left_hand_actuated_joint_values=left_hand_ddq,
                right_hand_actuated_joint_values=right_hand_ddq,
            )
            whole_tau_est = self.robot_model.get_configuration_from_actuated_joints(
                body_actuated_joint_values=body_tau_est,
                left_hand_actuated_joint_values=left_hand_tau_est,
                right_hand_actuated_joint_values=right_hand_tau_est,
            )
        else:
            # Body and hand joint measurements come in actuator order, so we need to convert them to joint order
            whole_q = self.robot_model.get_configuration_from_actuated_joints(
                body_actuated_joint_values=body_q,
            )
            whole_dq = self.robot_model.get_configuration_from_actuated_joints(
                body_actuated_joint_values=body_dq,
            )
            whole_ddq = self.robot_model.get_configuration_from_actuated_joints(
                body_actuated_joint_values=body_ddq,
            )
            whole_tau_est = self.robot_model.get_configuration_from_actuated_joints(
                body_actuated_joint_values=body_tau_est,
            )

        eef_obs = self.get_eef_obs(whole_q)

        obs = {
            "q": whole_q,
            "dq": whole_dq,
            "ddq": whole_ddq,
            "tau_est": whole_tau_est,
            "floating_base_pose": body_obs["floating_base_pose"],
            "floating_base_vel": body_obs["floating_base_vel"],
            "floating_base_acc": body_obs["floating_base_acc"],
            "wrist_pose": np.concatenate([eef_obs["left_wrist_pose"], eef_obs["right_wrist_pose"]]),
            "torso_quat": body_obs["torso_quat"],
            "torso_ang_vel": body_obs["torso_ang_vel"],
        }

        if self.use_sim and self.sim:
            obs.update(self.sim.get_privileged_obs())

        # Store last observation for safety checking
        self.last_obs = obs

        return obs

    @property
    def observation_space(self) -> gym.Space:
        # @todo: check if the low and high bounds are correct for body_obs.
        q_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.robot_model.num_dofs,))
        dq_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.robot_model.num_dofs,))
        ddq_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.robot_model.num_dofs,))
        tau_est_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.robot_model.num_dofs,))
        floating_base_pose_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(7,))
        floating_base_vel_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(6,))
        floating_base_acc_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(6,))
        wrist_pose_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(7 + 7,))
        return gym.spaces.Dict(
            {
                "floating_base_pose": floating_base_pose_space,
                "floating_base_vel": floating_base_vel_space,
                "floating_base_acc": floating_base_acc_space,
                "q": q_space,
                "dq": dq_space,
                "ddq": ddq_space,
                "tau_est": tau_est_space,
                "wrist_pose": wrist_pose_space,
            }
        )

    def queue_action(self, action: Dict[str, any]):
        # Safety check
        if self.last_obs is not None:
            safety_result = self.safety_monitor.handle_violations(self.last_obs, action)
            action = safety_result["action"]

        # Map action from joint order to actuator order
        body_actuator_q = self.robot_model.get_body_actuated_joints(action["q"])

        self.body().queue_action(
            {
                "body_q": body_actuator_q,
                "body_dq": np.zeros_like(body_actuator_q),
                "body_tau": np.zeros_like(body_actuator_q),
            }
        )

        if self.with_hands:
            left_hand_actuator_q = self.robot_model.get_hand_actuated_joints(
                action["q"], side="left"
            )
            right_hand_actuator_q = self.robot_model.get_hand_actuated_joints(
                action["q"], side="right"
            )
            self.hands().left.queue_action({"hand_q": left_hand_actuator_q})
            self.hands().right.queue_action({"hand_q": right_hand_actuator_q})

    def action_space(self) -> gym.Space:
        return gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.robot_model.num_dofs,))

    def calibrate_hands(self):
        """Calibrate the hand joint qpos if real robot"""
        if self.with_hands:
            print("calibrating left hand")
            self.hands().left.calibrate_hand()
            print("calibrating right hand")
            self.hands().right.calibrate_hand()
        else:
            print("Skipping hand calibration - hands disabled")

    def set_ik_indicator(self, teleop_cmd):
        """Set the IK indicators for the simulator"""
        if self.config["SIMULATOR"] == "robocasa":
            if "left_wrist" in teleop_cmd and "right_wrist" in teleop_cmd:
                left_wrist_input_pose = teleop_cmd["left_wrist"]
                right_wrist_input_pose = teleop_cmd["right_wrist"]
                ik_wrapper = self.sim.env.env.unwrapped.env
                ik_wrapper.set_target_poses_outside_env(
                    [left_wrist_input_pose, right_wrist_input_pose]
                )
        else:
            raise NotImplementedError("IK indicators are only implemented for robocasa simulator")

    def set_sync_mode(self, sync_mode: bool, steps_per_action: int = 4):
        """When set to True, the simulator will wait for the action to be sent to it"""
        if self.config["SIMULATOR"] == "robocasa":
            self.sim.set_sync_mode(sync_mode, steps_per_action)

    def reset(self):
        if self.sim:
            self.sim.reset()

    def close(self):
        if self.sim:
            self.sim.close()

    def robot_model(self) -> RobotModel:
        return self.robot_model

    def get_reward(self):
        if self.sim:
            return self.sim.get_reward()

    def reset_obj_pos(self):
        if hasattr(self.sim, "base_env") and hasattr(self.sim.base_env, "reset_obj_pos"):
            self.sim.base_env.reset_obj_pos()

    def get_eef_obs(self, q: np.ndarray) -> Dict[str, np.ndarray]:
        self.robot_model.cache_forward_kinematics(q)
        eef_obs = {}
        for side in ["left", "right"]:
            wrist_placement = self.robot_model.frame_placement(
                self.robot_model.supplemental_info.hand_frame_names[side]
            )
            wrist_pos, wrist_quat = wrist_placement.translation[:3], R.from_matrix(
                wrist_placement.rotation
            ).as_quat(scalar_first=True)
            eef_obs[f"{side}_wrist_pose"] = np.concatenate([wrist_pos, wrist_quat])

        return eef_obs

    def get_joint_safety_status(self) -> bool:
        """Get current joint safety status from the last queue_action safety check.

        Returns:
            bool: True if joints are safe (no shutdown required), False if unsafe
        """
        return self.last_safety_ok

    def handle_keyboard_button(self, key):
        # Only handles keyboard buttons for the mujoco simulator for now.
        if self.use_sim and self.config.get("SIMULATOR", "mujoco") == "mujoco":
            self.sim.handle_keyboard_button(key)


if __name__ == "__main__":
    env = G1Env(robot_model=instantiate_g1_robot_model(), wbc_version="gear_wbc")
    while True:
        print(env.observe())