Hands On Model State

手把手 2：把模型动作封装成状态

本课把一个有 timestep 的动作 policy 封装成强类型状态，并支持预热、进入帧、运行帧采样、暂停和自动返回。

1. 先给 policy 明确类型

不要让 _policy() 返回 Any。如果状态可能使用多种 policy，用联合类型：

if TYPE_CHECKING:
    from bxi_example_py_elf3.inference.beyondmimic import (
        DanceMotionPolicyGravityIsaaclab,
        DanceMotionPolicyGravityIsaaclabV3,
        DanceMotionPolicyMjlab,
    )

    MotionPolicy: TypeAlias = (
        DanceMotionPolicyMjlab
        | DanceMotionPolicyGravityIsaaclab
        | DanceMotionPolicyGravityIsaaclabV3
    )

2. 状态实现

class WaveMotionState(
    RobotControlState,
    EntryFrameProvider,
    RunningFrameProvider,
):
    def __init__(self, name: str, state_id: int, policy_attr: str):
        super().__init__(name, state_id)
        self.policy_attr = policy_attr
        self.playing = True

    def _policy(self, ctx: BxiExample) -> "MotionPolicy":
        return cast("MotionPolicy", getattr(ctx, self.policy_attr))

    def on_prepare(
        self,
        ctx: BxiExample,
        from_state: StateBehavior[BxiExample],
    ) -> None:
        policy = self._policy(ctx)
        policy.timestep = policy.start_frame
        if hasattr(policy, "timeinit"):
            policy.timeinit = 0.0
        ctx.preheat_model(policy)

    def on_enter(self, ctx: BxiExample) -> None:
        self.playing = True
        policy = self._policy(ctx)
        policy.timestep = policy.start_frame
        if hasattr(policy, "timeinit"):
            policy.timeinit = 0.0

    def get_entry_frame(self, ctx: BxiExample) -> MotorFrame:
        policy = self._policy(ctx)
        qpos = getattr(policy, "target_dof_pos", None)
        if qpos is None:
            qpos = getattr(policy, "default_dof_pos", None)
        if qpos is None:
            raise ValueError(f"state '{self.name}' policy has no entry position")
        return self._motor_frame(qpos, policy.kps, policy.kds)

    def sample_running_frame(
        self,
        ctx: BxiExample,
        dt: float,
        *,
        advance: bool,
    ) -> MotorFrame | None:
        policy = self._policy(ctx)
        if policy.timestep > policy.end_frame:
            return None

        qpos = policy.inference_step(
            ctx.current_q,
            ctx.current_dq,
            ctx.current_quat_wxyz,
            ctx.current_omega,
        )

        if self.playing and advance:
            policy.timestep += 50.0 * dt
        return self._motor_frame(qpos, policy.kps, policy.kds)

    def on_update(self, ctx: BxiExample, dt: float) -> None:
        policy = self._policy(ctx)

        if ctx.is_orientation_unsafe(ctx.current_quat_xyzw):
            ctx.request_state("zero_torque", trigger="safety")
            return

        if policy.timestep > policy.end_frame:
            ctx.request_state(
                "normal",
                trigger=f"{self.name}_finished",
                transition="soft_switch",
            )
            return

        self._apply_frame(
            ctx,
            self.sample_running_frame(ctx, dt, advance=True),
        )

    def on_action(self, ctx: BxiExample, action_name: str) -> bool:
        if action_name != "toggle_motion_pause":
            return False
        self.playing = not self.playing
        return True

3. 为什么准备和进入都要设置 timestep

• on_prepare() 在 Session 创建前预热模型，让进入帧可用。
• 活动过渡可能被中断，因此准备不等于正式进入。
• on_enter() 在过渡成功提交后建立正式状态运行起点。
• 如果准备分配了临时资源，在 on_prepare_cancel() 中释放。

4. 模型需要速度输入

def on_prepare(self, ctx: BxiExample, from_state) -> None:
    policy = self._policy(ctx)
    policy.timestep = policy.start_frame
    ctx.preheat_model(
        policy,
        with_cmd_vel=True,
        cmd_vel=self.get_cmd_vel(ctx),
    )

运行采样中同样调用状态自己的速度接口：

cmd_vel = self.get_cmd_vel(ctx)
qpos, _ = policy.inference_step(
    ctx.current_q,
    ctx.current_dq,
    ctx.current_quat_wxyz,
    ctx.current_omega,
    cmd_vel,
)

不要直接读取原始遥控器字段，否则会绕过 speed_profile。

5. YAML

states:
  wave_motion:
    behavior: WaveMotionState
    params:
      policy_attr: wave_motion
    transitions:
      on_event:
        normal_event:
          to: normal
          transition: soft_switch
        zero_torque_event: zero_torque
        toggle_dance_pause_event:
          action: toggle_motion_pause

入口边：

wave_motion_event:
  to: wave_motion
  transition: first_frame_switch

6. 使用运行混合

由于状态实现了两个能力，可以使用：

wave_motion_event:
  to: wave_motion
  transition: dual_running_blend

是否推进 policy 由 profile 的 advance_from / advance_to 决定。状态必须保证 advance=False 时不增加 timestep。

7. 调试

进入能力错误：

• 检查状态是否继承并实现 EntryFrameProvider。
• 检查预热后 target_dof_pos 或 default_dof_pos 是否存在。

动作重复推进：

• 检查 advance=False 路径。
• 确认只有 advance=True 时增加 timestep。

动作结束不返回：

• 检查 end_frame 与 timestep。
• 检查 pause action 是否让 playing=False。

过渡被打断后模型资源异常：

• 实现 on_prepare_cancel()。
• 不要在 on_prepare() 中假设目标一定会正式进入。

8. 检查清单

1. policy 有明确联合类型。
2. on_prepare() 预热但不输出电机。
3. 进入帧无法生成时明确抛错。
4. 运行采样遵守 advance。
5. on_update() 保留安全判断。
6. 动作完成主动请求下一个状态。