First Commit

Deploy_Tienkung/policy/beyond_mimic/fsm_beyond_mimic.py

@@ -0,0 +1,547 @@
+from FSM.fsm_base import FSMState, FSMStateName
+import numpy as np
+import yaml
+import os
+from types import SimpleNamespace
+from typing import Optional
+try:
+    import onnx
+except ImportError:  # pragma: no cover - runtime fallback when onnx isn't installed
+    onnx = None
+import onnxruntime
+try:
+    import torch
+except ImportError:  # pragma: no cover - torch is optional for warm start prep
+    torch = None
+from common.robot_data import RobotData
+from common.joystick import ControlFlag
+import time
+
+
+DEFAULT_BODY_NAMES = [
+    "pelvis",
+    "hip_pitch_l_link",
+    "hip_roll_l_link",
+    "hip_yaw_l_link",
+    "knee_pitch_l_link",
+    "ankle_pitch_l_link",
+    "ankle_roll_l_link",
+    "hip_pitch_r_link",
+    "hip_roll_r_link",
+    "hip_yaw_r_link",
+    "knee_pitch_r_link",
+    "ankle_pitch_r_link",
+    "ankle_roll_r_link",
+    "waist_yaw_link",
+    "waist_roll_link",
+    "waist_pitch_link",
+    "shoulder_pitch_l_link",
+    "shoulder_roll_l_link",
+    "shoulder_yaw_l_link",
+    "elbow_pitch_l_link",
+    "shoulder_pitch_r_link",
+    "shoulder_roll_r_link",
+    "shoulder_yaw_r_link",
+    "elbow_pitch_r_link",
+]
+
+class FSMStateBeyondMimic(FSMState):
+    def __init__(self, robot_data: RobotData, config_path: Optional[str] = None, variant_name: str = "default"):
+        super().__init__(robot_data)
+        self.motion_phase = 0
+        self.counter_step = 0
+        self.ref_motion_phase = 0
+        self.variant_name = variant_name
+        
+        
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        if config_path is None:
+            config_path = os.path.join(current_dir, "config", "BeyondMimic.yaml")
+        self.config_path = os.path.abspath(config_path)
+        with open(self.config_path, "r") as f:
+            config = yaml.load(f, Loader=yaml.FullLoader)
+            # 兼容多策略：模型路径仍然默认指向 policy/beyond_mimic/model 下
+            self.onnx_path = config["onnx_path"]
+            if not os.path.isabs(self.onnx_path):
+                self.onnx_path = os.path.join(current_dir, "model", self.onnx_path)
+            # self.motion_length = config["motion_length"]
+            # self.history_length = config["history_length"]
+            self.physical_dt = config["physical_dt"]
+            self.decimation_ = config["decimation"]
+            self.num_actions = config["num_actions"]
+            self.motor_nums = config["motor_nums"]
+            self.warm_start_time = config["warm_start_time"]
+            self.kps = config["kps"]
+            self.kds = config["kds"]
+            self.hold_final_reference = config.get("hold_final_reference", False)
+            self.motion_length: Optional[int] = config.get("motion_length")
+            self.body_names = config.get("body_names", DEFAULT_BODY_NAMES)
+            self.locked_joint_map = config["locked_joint_map"]
+            self.anchor_body_name = config.get("anchor_body", "pelvis")
+            if self.anchor_body_name not in self.body_names:
+                raise ValueError(f"Anchor body {self.anchor_body_name} missing from body list.")
+            self.anchor_body_index = self.body_names.index(self.anchor_body_name)
+            self.num_bodies = len(self.body_names)
+            self._warm_start_from_lab = np.zeros(self.num_actions, dtype=np.float32)
+            self._warm_start_to_lab = np.zeros(self.num_actions, dtype=np.float32)
+            self._warm_start_prev_target = np.zeros(self.num_actions, dtype=np.float32)
+
+            self.last_run_time = time.perf_counter()
+            
+            self.qj_obs = np.zeros(self.num_actions, dtype=np.float32)  # 初始化为最大可能大小
+            self.dqj_obs = np.zeros(self.num_actions, dtype=np.float32)
+            self.num_obs = None  # set after loading onnx
+            self.obs = None
+            self.action = np.zeros(self.num_actions, dtype=np.float32)
+
+            self.ref_joint_pos = np.zeros(self.num_actions, dtype=np.float32)
+            self.ref_joint_vel = np.zeros(self.num_actions, dtype=np.float32)
+            self.ref_body_pos_w = np.zeros((1, self.num_bodies, 3), dtype=np.float32)
+            self.ref_body_quat_w = np.zeros((1, self.num_bodies, 4), dtype=np.float32)
+            self.ref_body_lin_vel_w = np.zeros((1, self.num_bodies, 3), dtype=np.float32)
+            self.ref_body_ang_vel_w = np.zeros((1, self.num_bodies, 3), dtype=np.float32)
+            # load policy
+            self.onnx_model = None
+            metadata_props = []
+            if onnx is not None and hasattr(onnx, "load"):
+                try:
+                    self.onnx_model = onnx.load(self.onnx_path)
+                    metadata_props = getattr(self.onnx_model, "metadata_props", [])
+                except Exception as exc:
+                    print(f"[FSMStateBeyondMimic] Failed to load ONNX model via onnx.load: {exc}")
+            else:
+                print("[FSMStateBeyondMimic] Python onnx package unavailable, falling back to onnxruntime metadata.")
+
+            self.ort_session = onnxruntime.InferenceSession(self.onnx_path)
+            if not metadata_props:
+                model_meta = self.ort_session.get_modelmeta()
+                custom_map = getattr(model_meta, "custom_metadata_map", {})
+                metadata_props = [SimpleNamespace(key=k, value=v) for k, v in custom_map.items()]
+
+            input = self.ort_session.get_inputs()
+            self.input_name = []
+            for i, inpt in enumerate(input):
+                self.input_name.append(inpt.name)
+            obs_input = self.ort_session.get_inputs()[0]
+            last_dim = obs_input.shape[-1]
+            if isinstance(last_dim, int):
+                self.num_obs = last_dim
+            else:
+                self.num_obs = config.get("num_obs", 154)
+            self.obs = np.zeros(self.num_obs, dtype=np.float32)
+
+            # 从ONNX模型中读取参数
+            self.joint_seq = None
+            self.joint_pos_array_seq = None
+            self.action_scale = None
+            # self.stiffness_array_seq = None
+            # self.damping_array_seq = None
+            
+            for prop in metadata_props:
+                if prop.key == "joint_names":
+                    self.joint_seq = prop.value.split(",")
+                if prop.key == "default_joint_pos":   
+                    self.joint_pos_array_seq = np.array([float(x) for x in prop.value.split(",")])
+                # if prop.key == "joint_stiffness":
+                #     self.stiffness_array_seq = np.array([float(x) for x in prop.value.split(",")])
+                # if prop.key == "joint_damping":
+                #     self.damping_array_seq = np.array([float(x) for x in prop.value.split(",")])
+                if prop.key == "action_scale":
+                    self.action_scale = np.array([float(x) for x in prop.value.split(",")])
+                if prop.key in ("motion_length", "time_step_total"):
+                    try:
+                        self.motion_length = int(float(prop.value))
+                    except (TypeError, ValueError):
+                        print(f"[FSMStateBeyondMimic] Invalid motion_length metadata value: {prop.value}")
+
+            if self.motion_length is not None:
+                try:
+                    self.motion_length = int(self.motion_length)
+                except (TypeError, ValueError):
+                    print(f"[FSMStateBeyondMimic] Invalid motion_length config value: {self.motion_length}")
+                    self.motion_length = None
+            
+            # # 根据YAML配置设置关节映射
+            # self.mj2lab = np.array(config["mj2lab"], dtype=np.int32)
+            
+            # 设置从序列到实验室顺序的映射
+            self.joint_xml = [
+                "hip_pitch_l_joint", "hip_roll_l_joint", "hip_yaw_l_joint",
+                "knee_pitch_l_joint", "ankle_pitch_l_joint", "ankle_roll_l_joint",
+                "hip_pitch_r_joint", "hip_roll_r_joint", "hip_yaw_r_joint",
+                "knee_pitch_r_joint", "ankle_pitch_r_joint", "ankle_roll_r_joint",
+                "waist_yaw_joint", "waist_roll_joint", "waist_pitch_joint",
+                "shoulder_pitch_l_joint",
+                "elbow_pitch_l_joint",
+                "shoulder_pitch_r_joint",
+                "elbow_pitch_r_joint", 
+            ]
+            # 从MjXUML顺序映射到实验室顺序
+            self.mj2lab = np.array([self.joint_xml.index(joint) for joint in self.joint_seq])
+
+            # 从实验室顺序映射到MjXUML顺序
+            self.joint_pos_array = np.array([self.joint_pos_array_seq[self.joint_seq.index(joint)] for joint in self.joint_xml])
+            self.default_angles_lab = self.joint_pos_array_seq
+            self.action_scale_lab = self.action_scale
+
+            print("BeyondMimic-like policy initializing ...")
+            self._warmup_inference_counter = 0
+            self.warm_start_steps = 0
+            # Cache for the last motion frame so we can keep sending it after motion ends.
+            self._final_ref_cached = False
+            self._final_ref_joint_pos = np.zeros_like(self.ref_joint_pos)
+            self._final_ref_joint_vel = np.zeros_like(self.ref_joint_vel)
+            self._final_ref_body_pos_w = np.zeros_like(self.ref_body_pos_w)
+            self._final_ref_body_quat_w = np.zeros_like(self.ref_body_quat_w)
+            self._final_ref_body_lin_vel_w = np.zeros_like(self.ref_body_lin_vel_w)
+            self._final_ref_body_ang_vel_w = np.zeros_like(self.ref_body_ang_vel_w)
+    
+    def on_enter(self):
+        self.ref_motion_phase = 0.
+        self.motion_time = 0
+        self.counter_step = 0
+        self._warmup_inference_counter = 0
+        print(f"[FSMStateBeyondMimic] enter variant={self.variant_name}, config={self.config_path}")
+        if self.warm_start_time > 0:
+            step = self.decimation_ * self.physical_dt
+            self.warm_start_steps = max(1, int(self.warm_start_time / step))
+        else:
+            self.warm_start_steps = 0
+
+        observation = {}
+        observation[self.input_name[0]] = np.zeros((1, self.num_obs), dtype=np.float32)
+        observation[self.input_name[1]] = np.zeros((1, 1), dtype=np.float32)
+        outputs_result = self.ort_session.run(None, observation)
+        (
+            self.action,
+            self.ref_joint_pos,
+            self.ref_joint_vel,
+            self.ref_body_pos_w,
+            self.ref_body_quat_w,
+            self.ref_body_lin_vel_w,
+            self.ref_body_ang_vel_w,
+        ) = outputs_result
+
+        self.qj_obs = np.zeros(self.num_actions, dtype=np.float32)
+        self.dqj_obs = np.zeros(self.num_actions, dtype=np.float32)
+        self.obs = np.zeros(self.num_obs)
+        self._final_ref_cached = False
+        self._warm_start_from_lab = self._get_current_joint_pos_lab()
+        self._warm_start_to_lab = self._get_onnx_first_frame_lab()
+        self._warm_start_prev_target = np.array(self._warm_start_from_lab, copy=True)
+
+        # self.action = np.zeros(self.num_actions)
+
+        pass
+        
+    def quat_mul(self, q1, q2):
+        w1, x1, y1, z1 = q1[0], q1[1], q1[2], q1[3]
+        w2, x2, y2, z2 = q2[0], q2[1], q2[2], q2[3]
+        # perform multiplication
+        ww = (z1 + x1) * (x2 + y2)
+        yy = (w1 - y1) * (w2 + z2)
+        zz = (w1 + y1) * (w2 - z2)
+        xx = ww + yy + zz
+        qq = 0.5 * (xx + (z1 - x1) * (x2 - y2))
+        w = qq - ww + (z1 - y1) * (y2 - z2)
+        x = qq - xx + (x1 + w1) * (x2 + w2)
+        y = qq - yy + (w1 - x1) * (y2 + z2)
+        z = qq - zz + (z1 + y1) * (w2 - x2)
+        return np.array([w, x, y, z])
+        
+    def matrix_from_quat(self, q):
+        w, x, y, z = q
+        return np.array([
+            [1 - 2 * (y**2 + z**2), 2 * (x * y - z * w), 2 * (x * z + y * w)],
+            [2 * (x * y + z * w), 1 - 2 * (x**2 + z**2), 2 * (y * z - x * w)],
+            [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x**2 + y**2)]
+        ])
+
+    def yaw_quat(self, q):
+        w, x, y, z = q
+        yaw = np.arctan2(2 * (w * z + x * y), 1 - 2 * (y**2 + z**2))
+        return np.array([np.cos(yaw / 2), 0, 0, np.sin(yaw / 2)])
+    
+    def euler_single_axis_to_quat(self, angle, axis, degrees=False):
+        """
+        将单个欧拉角转换为四元数
+        
+        参数:
+            angle: 旋转角度
+            axis: 旋转轴，可以是 'x', 'y', 'z' 或者单位向量 [x, y, z]
+            degrees: 如果为True，输入角度为度数；如果为False，输入角度为弧度
+        
+        返回:
+            四元数 (w, x, y, z)
+        """
+        # 转换角度为弧度
+        if degrees:
+            angle = np.radians(angle)
+        
+        # 计算半角
+        half_angle = angle * 0.5
+        cos_half = np.cos(half_angle)
+        sin_half = np.sin(half_angle)
+        
+        # 根据旋转轴确定四元数分量
+        if isinstance(axis, str):
+            if axis.lower() == 'x':
+                return np.array([cos_half, sin_half, 0.0, 0.0])
+            elif axis.lower() == 'y':
+                return np.array([cos_half, 0.0, sin_half, 0.0])
+            elif axis.lower() == 'z':
+                return np.array([cos_half, 0.0, 0.0, sin_half])
+            else:
+                raise ValueError("axis must be 'x', 'y', 'z' or a 3D unit vector")
+        else:
+            # 假设axis是一个3D向量 [x, y, z]
+            axis = np.array(axis, dtype=np.float32)
+            # 归一化轴向量
+            axis_norm = np.linalg.norm(axis)
+            if axis_norm == 0:
+                raise ValueError("axis vector cannot be zero")
+            axis = axis / axis_norm
+            
+            # 计算四元数分量
+            w = cos_half
+            x = sin_half * axis[0]
+            y = sin_half * axis[1]
+            z = sin_half * axis[2]
+            
+            return np.array([w, x, y, z])
+
+    def inner_run(self):
+        robot_quat = self.robot_data_.get_robot_quat()
+        qj = self.robot_data_.get_joint_pos()
+        # 将29dof自由度的数据映射回锁住腕部6关节，之后的逻辑和和之前没区别
+        qj = qj[self.locked_joint_map]
+
+        qj = qj[self.mj2lab]
+        qj = (qj - self.default_angles_lab)
+
+        # IMU mounted on pelvis, so directly use measured orientation.
+        ref_anchor_pos_w, ref_anchor_ori_w = self._get_ref_anchor_pose()
+
+        # 在第一帧提取当前机器人yaw方向，与参考动作yaw方向做差（与beyond mimic一致）
+        if(self.counter_step < 2):
+            init_to_anchor = self.matrix_from_quat(self.yaw_quat(ref_anchor_ori_w))
+            world_to_anchor = self.matrix_from_quat(self.yaw_quat(robot_quat))
+            self.init_to_world = world_to_anchor @ init_to_anchor.T
+            print("self.init_to_world: ", self.init_to_world)
+            self.counter_step += 1
+            return
+
+        robot_rot_mat = self.matrix_from_quat(robot_quat)
+        motion_anchor_ori_b = robot_rot_mat.T @ self.init_to_world @ self.matrix_from_quat(ref_anchor_ori_w)
+
+        ang_vel = self.robot_data_.get_angular_velocity()
+
+        dqj = self.robot_data_.get_joint_vel()
+        #映射到23dof
+        dqj = dqj[self.locked_joint_map]
+
+        use_warmstart = (
+            self.warm_start_steps > 0
+            and self._warmup_inference_counter < self.warm_start_steps
+        )
+        blended_target = None
+        if use_warmstart:
+            blend = (self._warmup_inference_counter + 1) / self.warm_start_steps
+            blended_target = (1.0 - blend) * self._warm_start_from_lab + blend * self._warm_start_to_lab
+            blended_vel = (blended_target - self._warm_start_prev_target) / (
+                self.decimation_ * self.physical_dt
+            )
+            self._warm_start_prev_target = blended_target
+
+            command_joint_pos = blended_target.reshape(1, -1)
+            command_joint_vel = blended_vel.reshape(1, -1)
+            safe_scale = np.where(self.action_scale_lab == 0, 1.0, self.action_scale_lab)
+            action_for_history = (blended_target - self.default_angles_lab) / safe_scale
+        else:
+            command_joint_pos = self.ref_joint_pos
+            command_joint_vel = self.ref_joint_vel
+            action_for_history = self.action
+
+        command_root = self.matrix_from_quat(ref_anchor_ori_w)
+        command_vec = np.concatenate(
+            (
+                command_joint_pos.squeeze(0),
+                command_joint_vel.squeeze(0),
+                # ref_anchor_pos_w,
+                # command_root[:, :2].reshape(-1),
+            ),
+            dtype=np.float32,
+        )
+
+        mimic_obs_buf = np.concatenate(
+            (
+                command_vec,
+                motion_anchor_ori_b[:, :2].reshape(-1),
+                ang_vel,
+                qj,
+                dqj[self.mj2lab],
+                np.asarray(action_for_history, dtype=np.float32).reshape(-1),
+            ),
+            axis=-1,
+            dtype=np.float32,
+        )
+        if mimic_obs_buf.shape[0] != self.num_obs:
+            raise RuntimeError(f"Observation length mismatch. Expected {self.num_obs}, got {mimic_obs_buf.shape[0]}.")
+        
+        if torch is not None:
+            mimic_obs_tensor = torch.from_numpy(mimic_obs_buf).unsqueeze(0).cpu().numpy()
+        else:
+            mimic_obs_tensor = np.expand_dims(mimic_obs_buf, axis=0)
+        observation = {}
+
+        # obs0 是网络观测，obs1 是当前时间步，用于输出参考动作信息
+        observation[self.input_name[0]] = mimic_obs_tensor
+        time_index = max(self.counter_step - self.warm_start_steps, 0)
+        
+        if (
+            self.hold_final_reference
+            and self.motion_length is not None
+            and self.motion_length > 0
+        ):
+            if self.motion_length is not None and self.motion_length > 0:
+                time_index = min(time_index, self.motion_length - 1)
+
+        observation[self.input_name[1]] = np.array([[time_index]], dtype=np.float32)
+        outputs_result = self.ort_session.run(None, observation)
+        (
+            self.action,
+            self.ref_joint_pos,
+            self.ref_joint_vel,
+            self.ref_body_pos_w,
+            self.ref_body_quat_w,
+            self.ref_body_lin_vel_w,
+            self.ref_body_ang_vel_w,
+        ) = outputs_result
+
+        if (
+            self.hold_final_reference
+            and self.motion_length is not None
+            and self.motion_length > 0
+        ):
+            if time_index == self.motion_length - 1 and not self._final_ref_cached:
+                self._cache_final_ref()
+            elif self.counter_step >= self.motion_length and self._final_ref_cached:
+                self._apply_final_ref()
+        target_dof_pos_mj = np.zeros(29)
+        target_dof_pos_mj_19dof = np.zeros(19)
+        if use_warmstart and blended_target is not None:
+            target_dof_pos_lab = blended_target
+            # Keep action history aligned with the inserted warm trajectory.
+            self.action = np.asarray(action_for_history, dtype=np.float32).reshape(1, -1)
+        else:
+            target_dof_pos_lab = self.action * self.action_scale_lab + self.default_angles_lab
+            if target_dof_pos_lab.ndim > 1:
+                target_dof_pos_lab = np.squeeze(target_dof_pos_lab, axis=0)
+
+        if self.warm_start_steps > 0:
+            self._warmup_inference_counter += 1
+            if self._warmup_inference_counter <= self.warm_start_steps:
+                blend = self._warmup_inference_counter / self.warm_start_steps
+                if not use_warmstart:
+                    target_dof_pos_lab = (1.0 - blend) * self._warm_start_from_lab + blend * self._warm_start_to_lab
+
+        target_dof_pos_mj_19dof[self.mj2lab] = target_dof_pos_lab
+        target_dof_pos_mj[self.locked_joint_map] = target_dof_pos_mj_19dof
+
+        # Set joint commands exactly like C++
+        for i in range(self.motor_nums):
+            # C++: robot_data_->q_d_(35 - motor_num_ + i)
+            joint_idx = 35 - self.motor_nums + i
+            self.robot_data_.q_d_[joint_idx] = target_dof_pos_mj[i]
+            self.robot_data_.q_dot_d_[joint_idx] = 0.0
+            self.robot_data_.tau_d_[joint_idx] = 0.0
+
+        # update motion phase
+        self.counter_step += 1
+
+    def _cache_final_ref(self):
+        if not self.hold_final_reference:
+            return
+        self._final_ref_cached = True
+        self._final_ref_joint_pos = np.array(self.ref_joint_pos, copy=True)
+        self._final_ref_joint_vel = np.array(self.ref_joint_vel, copy=True)
+        self._final_ref_body_pos_w = np.array(self.ref_body_pos_w, copy=True)
+        self._final_ref_body_quat_w = np.array(self.ref_body_quat_w, copy=True)
+        self._final_ref_body_lin_vel_w = np.array(self.ref_body_lin_vel_w, copy=True)
+        self._final_ref_body_ang_vel_w = np.array(self.ref_body_ang_vel_w, copy=True)
+
+    def _apply_final_ref(self):
+        if not self.hold_final_reference or not self._final_ref_cached:
+            return
+        self.ref_joint_pos = np.array(self._final_ref_joint_pos, copy=True)
+        self.ref_joint_vel = np.array(self._final_ref_joint_vel, copy=True)
+        self.ref_body_pos_w = np.array(self._final_ref_body_pos_w, copy=True)
+        self.ref_body_quat_w = np.array(self._final_ref_body_quat_w, copy=True)
+        self.ref_body_lin_vel_w = np.array(self._final_ref_body_lin_vel_w, copy=True)
+        self.ref_body_ang_vel_w = np.array(self._final_ref_body_ang_vel_w, copy=True)
+
+    def run(self, flag: ControlFlag):
+        if int(self.robot_data_.time_now_ / self.physical_dt) % self.decimation_ == 0:
+            current_time = time.perf_counter()
+            print(f"Inference hz: {1/(current_time - self.last_run_time)}")
+            self.last_run_time = current_time
+            self.inner_run()
+        self.set_kp_kd()
+    def set_kp_kd(self):
+        # Set kp/kd gains
+        self.robot_data_.joint_kp_p_[:self.motor_nums] = self.kps
+        self.robot_data_.joint_kd_p_[:self.motor_nums] = self.kds
+    def on_exit(self):
+        self.action = np.zeros(self.num_actions, dtype=np.float32)
+        # self.action_buf = np.zeros(23 * self.history_length, dtype=np.float32)
+        self.ref_motion_phase = 0.
+        # self.ref_motion_phase_buf = np.zeros(1 * self.history_length, dtype=np.float32)
+        self.motion_time = 0
+        self.counter_step = 0
+        self._final_ref_cached = False
+        
+        print("exited")
+
+    
+    def check_transition(self, flag: ControlFlag) -> FSMStateName:
+        """检查状态转换"""
+        if flag.fsm_state_command == "gotoSTOP":
+            return FSMStateName.STOP
+        elif flag.fsm_state_command == "gotoZERO":
+            return FSMStateName.ZERO
+        elif flag.fsm_state_command == "gotoBEYONDMIMIC":
+            return FSMStateName.BEYONDMIMIC
+        elif flag.fsm_state_command == "gotoBEYONDZERO":
+            return FSMStateName.BEYONDZERO
+        elif flag.fsm_state_command == "gotoWALKAMP":
+            return FSMStateName.WALKAMP
+        else:
+            return None  # 无状态转换
+
+    def _get_ref_anchor_pose(self):
+        ref_pos = self.ref_body_pos_w[:, self.anchor_body_index].squeeze(0)
+        ref_quat = self.ref_body_quat_w[:, self.anchor_body_index].squeeze(0)
+        return ref_pos.astype(np.float32), ref_quat.astype(np.float32)
+
+    def _get_current_joint_pos_lab(self) -> np.ndarray:
+        try:
+            current_q = self.robot_data_.get_joint_pos()
+            current_q = current_q[self.locked_joint_map]
+            current_q_lab = current_q[self.mj2lab]
+            return current_q_lab.astype(np.float32)
+        except Exception as exc:
+            print(f"[FSMStateBeyondMimic] Failed to read current joint pose: {exc}")
+            return np.array(self.default_angles_lab, copy=True)
+
+    def _get_onnx_first_frame_lab(self) -> np.ndarray:
+        try:
+            action = self.action
+            if action is None:
+                raise RuntimeError("ONNX action output is None")
+            if action.ndim > 1:
+                action = np.squeeze(action, axis=0)
+            first_frame = action * self.action_scale_lab + self.default_angles_lab
+            return first_frame.astype(np.float32)
+        except Exception as exc:
+            print(f"[FSMStateBeyondMimic] Failed to read ONNX first frame: {exc}")
+            return np.array(self.default_angles_lab, copy=True)