import os import math import numpy as np import mujoco import onnxruntime as ort # Default pose used by the policy (14 DoF) DEFAULT_POSE = np.array([ 0.0, # left_hip_yaw -0.0873, # left_hip_roll -0.4579, # left_hip_pitch -0.0049, # left_knee 0.4530, # left_ankle 0.3491, # neck_pitch 0.3491, # head_pitch 0.0, # head_yaw 0.0, # head_roll 0.0, # right_hip_yaw 0.0873, # right_hip_roll 0.4579, # right_hip_pitch 0.0049, # right_knee -0.4530, # right_ankle ], dtype=np.float32) BALL_OFFSET_X = 0.09 BALL_OFFSET_ABS_Y = 0.042 BALL_RADIUS = 0.035 class SpacePolicyInference: def __init__(self, model, data, policies_dir): self.model = model self.data = data self.action_scale = 1.0 self.use_projected_gravity = True self.new_cmd_obs = True self.switch_threshold = 0.05 # Load ONNX policies walking_path = os.path.join(policies_dir, "alpha_walking.onnx") standing_path = os.path.join(policies_dir, "alpha_stand.onnx") sitstand_path = os.path.join(policies_dir, "alpha_sitstand.onnx") ground_pick_path = os.path.join(policies_dir, "alpha_ground_pick.onnx") kick_left_path = os.path.join(policies_dir, "ball_kick_left.onnx") roulade_path = os.path.join(policies_dir, "roulade.onnx") self.walking_session = ort.InferenceSession(walking_path) if os.path.exists(walking_path) else None self.standing_session = ort.InferenceSession(standing_path) if os.path.exists(standing_path) else None self.sit_session = ort.InferenceSession(sitstand_path) if os.path.exists(sitstand_path) else None self.ground_pick_session = ort.InferenceSession(ground_pick_path) if os.path.exists(ground_pick_path) else None self.behavior_sessions = {} self.behavior_durations = {} if os.path.exists(kick_left_path): self.behavior_sessions["kick_left"] = ort.InferenceSession(kick_left_path) self.behavior_durations["kick_left"] = 3.0 if os.path.exists(roulade_path): self.behavior_sessions["roulade"] = ort.InferenceSession(roulade_path) self.behavior_durations["roulade"] = 2.0 self.is_sitstand = True self.sit_mode = False self.ground_pick_mode = False self.ground_pick_phase = 0.0 self.ground_pick_period = 4.0 self.behavior_mode = None self.behavior_time_left = 0.0 self.current_policy = "walking" if self.walking_session else "standing" self.ort_session = self.walking_session if self.walking_session else self.standing_session self.input_name = self.ort_session.get_inputs()[0].name self.output_name = self.ort_session.get_outputs()[0].name self.imu_ang_vel_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, "imu_ang_vel") self.trunk_base_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "trunk_base") _trunk_jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "trunk_base_freejoint") self._trunk_qpos_adr = int(model.jnt_qposadr[_trunk_jid]) _ball_jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "ball_free") self.ball_qpos_adr = int(model.jnt_qposadr[_ball_jid]) if _ball_jid >= 0 else None self.ball_qvel_adr = int(model.jnt_dofadr[_ball_jid]) if _ball_jid >= 0 else None self.n_joints = model.nu self.joint_qpos_indices = [int(model.jnt_qposadr[model.actuator_trnid[i, 0]]) for i in range(model.nu)] self.joint_qvel_indices = [int(model.jnt_dofadr[model.actuator_trnid[i, 0]]) for i in range(model.nu)] self.default_pose = DEFAULT_POSE[:self.n_joints] self.last_action = np.zeros(self.n_joints, dtype=np.float32) self.vel_cmd = np.zeros(3, dtype=np.float32) self.body_cmd = np.zeros(6, dtype=np.float32) self.head_offset = np.zeros(4, dtype=np.float32) self.command = np.zeros(13, dtype=np.float32) def _update_command(self): if self.behavior_mode is not None: self.command = np.zeros(13, dtype=np.float32) return cmd = np.zeros(13, dtype=np.float32) if self.current_policy == "walking": cmd[0:3] = self.vel_cmd elif self.current_policy == "sit" and self.is_sitstand: cmd[0] = 1.0 if self.sit_mode else 0.0 cmd[3:7] = self.head_offset cmd[7:13] = self.body_cmd self.command = cmd def set_vel_cmd(self, lin_vel_x=0.0, lin_vel_y=0.0, ang_vel_z=0.0): self.vel_cmd = np.array([lin_vel_x, lin_vel_y, ang_vel_z], dtype=np.float32) magnitude = float(np.linalg.norm(self.vel_cmd)) if magnitude <= self.switch_threshold and self.standing_session: self.current_policy = "standing" self.ort_session = self.standing_session elif self.walking_session: self.current_policy = "walking" self.ort_session = self.walking_session self._update_command() def trigger_behavior(self, name): if name in self.behavior_sessions: if name == "kick_left": self._place_ball() self.behavior_mode = name self.behavior_time_left = self.behavior_durations[name] self.current_policy = name self.ort_session = self.behavior_sessions[name] self._update_command() def trigger_sit(self): if self.sit_session: self.sit_mode = True self.current_policy = "sit" self.ort_session = self.sit_session self._update_command() def trigger_ground_pick(self): if self.ground_pick_session: self.ground_pick_mode = True self.ground_pick_phase = 0.0 self.current_policy = "ground_pick" self.ort_session = self.ground_pick_session def _place_ball(self): if self.ball_qpos_adr is None: return adr = self._trunk_qpos_adr x, y = float(self.data.qpos[adr]), float(self.data.qpos[adr + 1]) qw, qx, qy, qz = self.data.qpos[adr + 3:adr + 7] yaw = math.atan2(2.0 * (qw * qz + qx * qy), 1.0 - 2.0 * (qy * qy + qz * qz)) off_y = BALL_OFFSET_ABS_Y bx = x + math.cos(yaw) * BALL_OFFSET_X - math.sin(yaw) * off_y by = y + math.sin(yaw) * BALL_OFFSET_X + math.cos(yaw) * off_y self.data.qpos[self.ball_qpos_adr:self.ball_qpos_adr + 7] = [bx, by, BALL_RADIUS, 1, 0, 0, 0] if self.ball_qvel_adr is not None: self.data.qvel[self.ball_qvel_adr:self.ball_qvel_adr + 6] = 0.0 def update(self, dt: float): if self.ground_pick_mode: new_phase = self.ground_pick_phase + dt / self.ground_pick_period if new_phase >= 0.7: self.ground_pick_mode = False self.current_policy = "walking" if self.walking_session else "standing" self.ort_session = self.walking_session if self.walking_session else self.standing_session self._update_command() else: self.ground_pick_phase = new_phase self.command[0] = np.cos(2 * np.pi * self.ground_pick_phase) self.command[1] = np.sin(2 * np.pi * self.ground_pick_phase) self.command[2] = 0.0 elif self.behavior_mode is not None: self.behavior_time_left -= dt if self.behavior_time_left <= 0.0: self.behavior_mode = None self.current_policy = "walking" if self.walking_session else "standing" self.ort_session = self.walking_session if self.walking_session else self.standing_session self._update_command() def quat_rotate_inverse(self, quat, vec): w = quat[0] xyz = quat[1:4] t = np.cross(xyz, vec) * 2 return vec - w * t + np.cross(xyz, t) def get_observations(self): obs = [] sensor_adr = self.model.sensor_adr[self.imu_ang_vel_id] ang_vel = self.data.sensordata[sensor_adr:sensor_adr + 3].copy().astype(np.float32) obs.append(ang_vel) quat = self.data.xquat[self.trunk_base_id].copy().astype(np.float32) proj_grav = self.quat_rotate_inverse(quat, np.array([0.0, 0.0, -1.0], dtype=np.float32)) obs.append(proj_grav) current_pos = self.data.qpos[self.joint_qpos_indices].copy().astype(np.float32) obs.append(current_pos - self.default_pose) joint_vel = self.data.qvel[self.joint_qvel_indices].copy().astype(np.float32) obs.append(joint_vel) obs.append(self.last_action) obs.append(self.command) return np.concatenate(obs).astype(np.float32) def step(self): obs = self.get_observations().reshape(1, -1) action = self.ort_session.run([self.output_name], {self.input_name: obs})[0].squeeze(0).astype(np.float32) self.last_action = action.copy() target_positions = self.default_pose + action * self.action_scale self.data.ctrl[:] = target_positions