[WIP] Add Mujoco sim2sim deployment pipeline for carrybox task

deploy/deploy_mujoco/config.yaml

@@ -0,0 +1,69 @@
+policy_path: "{LEGGED_GYM_ROOT_DIR}/../deploy/pretrain_model/carry_box.onnx"
+xml_path: "{LEGGED_GYM_ROOT_DIR}/resources/robots/g1/scene.xml"
+
+# Simulation
+simulation_duration: 600.0
+simulation_dt: 0.005
+control_decimation: 4
+
+# Dimensions
+num_actions: 29
+num_obs: 123
+num_history: 6
+
+# Scales
+ang_vel_scale: 0.25
+dof_pos_scale: 1.0
+dof_vel_scale: 0.05
+lin_vel_scale: 2.0
+action_scale: 0.25
+
+# ---------------------------------------------------------
+# Joint Order & Params
+# Order: 
+# 1. Left Leg (Pitch, Roll, Yaw, Knee, AnkleP, AnkleR)
+# 2. Right Leg (Pitch, Roll, Yaw, Knee, AnkleP, AnkleR)
+# 3. Waist (Yaw, Roll, Pitch)
+# 4. Left Arm (ShoP, ShoR, ShoY, Elbow, WristR, WristP, WristY)
+# 5. Right Arm (ShoP, ShoR, ShoY, Elbow, WristR, WristP, WristY)
+# ---------------------------------------------------------
+
+kps: [
+    # Left Leg (6)
+    150, 150, 150, 300, 40, 40,
+    # Right Leg (6)
+    150, 150, 150, 300, 40, 40,
+    # Waist (3)
+    300, 300, 300,
+    # Left Arm (7)
+    200, 200, 200, 100, 20, 20, 20,
+    # Right Arm (7)
+    200, 200, 200, 100, 20, 20, 20
+]
+
+kds: [
+    # Left Leg
+    2, 2, 2, 4, 1, 1,
+    # Right Leg
+    2, 2, 2, 4, 1, 1,
+    # Waist
+    4, 4, 4,
+    # Left Arm
+    3, 3, 3, 1, 0.5, 0.5, 0.5,
+    # Right Arm
+    3, 3, 3, 1, 0.5, 0.5, 0.5
+]
+
+# Default Angles based on G1Cfg.init_state
+default_angles: [
+    # Left Leg: HipP, HipR, HipY, Knee, AnkP, AnkR
+    -0.1, 0.0, 0.0, 0.3, -0.2, 0.0,
+    # Right Leg: HipP, HipR, HipY, Knee, AnkP, AnkR
+    -0.1, 0.0, 0.0, 0.3, -0.2, 0.0,
+    # Waist: Yaw, Roll, Pitch
+    0.0, 0.0, 0.0,
+    # Left Arm: ShoP, ShoR, ShoY, Elbow, WristR, WristP, WristY
+    0.0, 0.1, 0.0, 1.2, 0.0, 0.0, 0.0,
+    # Right Arm: ShoP, ShoR, ShoY, Elbow, WristR, WristP, WristY
+    0.0, -0.1, 0.0, 1.2, 0.0, 0.0, 0.0
+]

deploy/deploy_mujoco/deploy.py

@@ -0,0 +1,219 @@
+import os
+import time
+import mujoco.viewer
+import mujoco
+import numpy as np
+from legged_gym import LEGGED_GYM_ROOT_DIR
+import torch
+import yaml
+from collections import deque
+import onnxruntime as ort
+from scipy.spatial.transform import Rotation as R
+
+np.set_printoptions(precision=4, suppress=True, floatmode='fixed')
+
+def quat_rotate_inverse(q, v):
+    """
+    Rotate vector v by inverse of quaternion q.
+    Transforms vector v from World Frame to Local Frame defined by q.
+    q: [w, x, y, z] (Mujoco)
+    v: [x, y, z]
+    """
+    # Mujoco [w, x, y, z] -> Scipy [x, y, z, w]
+    r = R.from_quat([q[1], q[2], q[3], q[0]])
+    return r.inv().apply(v)
+
+def get_gravity_orientation(quaternion):
+    """
+    Computes gravity vector (0,0,-1) in the local frame.
+    """
+    # Just rotate standard gravity vector by inverse orientation
+    gravity_world = np.array([0.0, 0.0, -1.0])
+    return quat_rotate_inverse(quaternion, gravity_world)
+
+def quat_to_rot6d_local(q_ref, q_target):
+    """
+    Compute relative rotation of q_target w.r.t q_ref in 6D.
+    Matches IsaacGym's quat_to_tan_norm (X-axis and Z-axis).
+    """
+    r_ref = R.from_quat([q_ref[1], q_ref[2], q_ref[3], q_ref[0]])
+    r_tar = R.from_quat([q_target[1], q_target[2], q_target[3], q_target[0]])
+
+    r_local = r_ref.inv() * r_tar
+    matrix = r_local.as_matrix()
+    col_x = matrix[:, 0]
+    col_z = matrix[:, 2]
+
+    return np.concatenate([col_x, col_z])
+
+def pd_control(target_q, q, kp, target_dq, dq, kd):
+    return (target_q - q) * kp + (target_dq - dq) * kd
+
+def load_onnx_policy(path):
+    model = ort.InferenceSession(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="cuda:0")
+    return run_inference
+
+if __name__ == "__main__":
+    config_file = "config.yaml"
+
+    config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), config_file)
+    with open(config_path, "r") as f:
+        config = yaml.load(f, Loader=yaml.FullLoader)
+
+    policy_path = config["policy_path"].replace("{LEGGED_GYM_ROOT_DIR}", LEGGED_GYM_ROOT_DIR)
+    xml_path = config["xml_path"].replace("{LEGGED_GYM_ROOT_DIR}", LEGGED_GYM_ROOT_DIR)
+
+    simulation_duration = config["simulation_duration"]
+    simulation_dt = config["simulation_dt"]
+    control_decimation = config["control_decimation"]
+
+    kps = np.array(config["kps"], dtype=np.float32)
+    kds = np.array(config["kds"], dtype=np.float32)
+    default_angles = np.array(config["default_angles"], dtype=np.float32)
+
+    ang_vel_scale = config["ang_vel_scale"]
+    dof_pos_scale = config["dof_pos_scale"]
+    dof_vel_scale = config["dof_vel_scale"]
+    action_scale = config["action_scale"]
+
+    num_actions = config["num_actions"]
+    num_single_obs = config["num_obs"]
+    num_history = config.get("num_history", 6)
+
+    # --- Init Buffers ---
+    action = np.zeros(num_actions, dtype=np.float32)
+    last_action = np.zeros(num_actions, dtype=np.float32)
+    target_dof_pos = default_angles.copy()
+
+    obs_history = deque(maxlen=num_history)
+    for _ in range(num_history):
+        obs_history.append(np.zeros(num_single_obs, dtype=np.float32))
+
+    m = mujoco.MjModel.from_xml_path(xml_path)
+    d = mujoco.MjData(m)
+    m.opt.timestep = simulation_dt
+    renderer = mujoco.Renderer(m, height=480, width=640)
+
+    # policy = torch.jit.load(policy_path)
+    policy = load_onnx_policy(policy_path)
+    # policy.eval()
+
+    torso_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, "torso_link")
+
+    gyro_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_SENSOR, "imu-torso-angular-velocity")
+
+    ee_names = ["left_palm_link", "right_palm_link", "left_ankle_pitch_link", "right_ankle_pitch_link", "mid360_link"]
+    ee_ids = []
+    for name in ee_names:
+        bid = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, name)
+        if bid == -1:
+            print(f"Error: Body {name} not found! FK obs will be wrong.")
+        ee_ids.append(bid)
+
+    box_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, "box")
+    box_jnt_adr = m.body_jntadr[box_id] if box_id != -1 else -1
+
+    goal_pos_world = np.array([2.0, 2.0, 0.6]) 
+
+    print(f"Starting simulation. Torso ID: {torso_id}, Gyro ID: {gyro_id}")
+
+    counter = 0
+    with mujoco.viewer.launch_passive(m, d) as viewer: 
+        start = time.time()
+        while viewer.is_running() and time.time() - start < simulation_duration:
+            step_start = time.time()
+
+            tau = pd_control(target_dof_pos, d.qpos[7:7+num_actions], kps, np.zeros_like(kds), d.qvel[6:6+num_actions], kds)
+            d.ctrl[:] = tau
+
+            mujoco.mj_step(m, d)
+
+            counter += 1
+            if counter % control_decimation == 0:
+                torso_quat = d.xquat[torso_id] # [w, x, y, z] Global orientation
+                # print(f"{torso_quat=}")
+                torso_pos = d.xpos[torso_id]
+                root_pos = d.qpos[:3] # Root (Pelvis) position
+
+                # Proprioception 108
+
+                # Base Angular Velocity 3
+                # 使用传感器数据有一帧时延 issue #2419
+                # if gyro_id != -1:
+                #     ang_vel = d.sensordata[m.sensor_adr[gyro_id]:m.sensor_adr[gyro_id]+3] * ang_vel_scale
+                # else:
+                #     ang_vel = np.zeros(3) 
+                ang_vel_world = d.cvel[torso_id, :3]
+                ang_vel = quat_rotate_inverse(torso_quat, ang_vel_world) * ang_vel_scale
+                # print(f"{ang_vel=}")
+
+                # Projected Gravity 3
+                gravity_vec = get_gravity_orientation(torso_quat)
+                # print(f"{gravity_vec=}")
+
+                # Joint Pos/Vel 29+29
+                qj = (d.qpos[7:7+num_actions] - default_angles) * dof_pos_scale
+                dqj = d.qvel[6:6+num_actions] * dof_vel_scale
+
+                # End Effector Positions 15
+                ee_list = []
+
+                for bid in ee_ids:
+                    if bid != -1:
+                        ee_world = d.xpos[bid]
+                        ee_local = quat_rotate_inverse(torso_quat, ee_world - root_pos)
+                        ee_list.append(ee_local)
+                    else:
+                        ee_list.append(np.zeros(3))
+                ee_flat = np.concatenate(ee_list)
+                # print(f"{ee_flat=}")
+
+                # Task Obs 15
+                if box_id != -1 and box_jnt_adr != -1:
+                    qpos_adr = m.jnt_qposadr[box_jnt_adr]
+                    box_pos_world = d.qpos[qpos_adr:qpos_adr+3]
+                    box_quat_world = d.qpos[qpos_adr+3:qpos_adr+7]
+                    box_local = quat_rotate_inverse(torso_quat, box_pos_world - root_pos)
+                    # print(f"{box_local=}")
+
+                    box_rot6d = quat_to_rot6d_local(torso_quat, box_quat_world)
+                    box_size = np.array([0.3, 0.3, 0.25])
+                else:
+                    box_local = np.zeros(3)
+                    box_rot6d = np.zeros(6)
+                    box_size = np.zeros(3)
+
+                # Goal 3
+                goal_local = quat_rotate_inverse(torso_quat, goal_pos_world - root_pos)
+
+                obs_t = np.concatenate([
+                    ang_vel,        # 3
+                    gravity_vec,    # 3
+                    qj,             # 29
+                    dqj,            # 29
+                    ee_flat,        # 15
+                    last_action,    # 29
+                    box_local,      # 3
+                    box_rot6d,      # 6
+                    box_size,       # 3
+                    goal_local      # 3
+                ]).astype(np.float32)
+
+                obs_history.append(obs_t)
+                obs_tensor = torch.from_numpy(np.concatenate(list(obs_history))).unsqueeze(0)
+
+                with torch.no_grad():
+                    raw_action = policy(obs_tensor).detach().cpu().numpy().squeeze()
+                # print(f"{raw_action=}")
+                last_action = np.clip(raw_action, -100., 100.)
+                target_dof_pos = raw_action * action_scale + default_angles
+
+            viewer.sync()
+
+            time_until_next_step = m.opt.timestep - (time.time() - step_start)
+            if time_until_next_step > 0:
+                time.sleep(time_until_next_step)