diff --git a/deploy/deploy_mujoco/config.yaml b/deploy/deploy_mujoco/config.yaml new file mode 100644 index 0000000..6c82fc3 --- /dev/null +++ b/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 +] \ No newline at end of file diff --git a/deploy/deploy_mujoco/deploy.py b/deploy/deploy_mujoco/deploy.py new file mode 100644 index 0000000..0cf6c8a --- /dev/null +++ b/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) diff --git a/deploy/pretrain_model/carry_box.onnx b/deploy/pretrain_model/carry_box.onnx new file mode 100644 index 0000000..29c5b5c Binary files /dev/null and b/deploy/pretrain_model/carry_box.onnx differ diff --git a/deploy/pretrain_model/carry_box.pt b/deploy/pretrain_model/carry_box.pt new file mode 100644 index 0000000..95ca39d Binary files /dev/null and b/deploy/pretrain_model/carry_box.pt differ diff --git a/legged_gym/resources/robots/g1/g1_29dof.xml b/legged_gym/resources/robots/g1/g1_29dof.xml new file mode 100644 index 0000000..78c689a --- /dev/null +++ b/legged_gym/resources/robots/g1/g1_29dof.xml @@ -0,0 +1,314 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/legged_gym/resources/robots/g1/scene.xml b/legged_gym/resources/robots/g1/scene.xml new file mode 100644 index 0000000..545d496 --- /dev/null +++ b/legged_gym/resources/robots/g1/scene.xml @@ -0,0 +1,32 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file