infer_engine.py

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