Hands On Model State

手把手 2：从零添加模型动作状态

上一课的 SinWaveState 不依赖模型，适合理解状态机。真实动作通常会依赖：

data/*.npz
data/*.onnx

本课用一个具体例子：

wave_motion
WaveMotionState

它代表一个由 .npz 动作数据和 .onnx 策略模型驱动的动作状态。

0. 最终会改哪些文件

新增或放入文件：

src/bxi_example_py_elf3/data/wave_motion.npz
src/bxi_example_py_elf3/data/wave_motion.onnx
src/bxi_example_py_elf3/bxi_example_py_elf3/states/wave_motion_state.py

修改文件：

src/bxi_example_py_elf3/launch/example_demo.launch.py
src/bxi_example_py_elf3/launch/example_demo_hw.launch.py
src/bxi_example_py_elf3/bxi_example_py_elf3/bxi_example_demo.py
src/bxi_example_py_elf3/bxi_example_py_elf3/robot_states.py
src/bxi_example_py_elf3/config/elf3_state_machine.yaml
src/bxi_example_py_elf3/config/release_protection.yaml

如果你还没创建 states/ 子包，先完成 手把手 1：从零写一个自定义状态 的前两步。

1. 放入模型和动作数据

把文件放到：

src/bxi_example_py_elf3/data/wave_motion.npz
src/bxi_example_py_elf3/data/wave_motion.onnx

setup.py 里已经有：

def get_data_files():
    # 遍历 data/ 并安装到 share/bxi_example_py_elf3/data

所以放进 data/ 后，colcon build 会自动安装到 package share 目录。

2. 在 launch 中声明模型 key

打开：

src/bxi_example_py_elf3/launch/example_demo.launch.py
src/bxi_example_py_elf3/launch/example_demo_hw.launch.py

两个文件都要改。一个是仿真，一个是硬件。

在 npz_file_dict 添加：

npz_file_dict = {
    "recover": "data/recover.npz",
    "dance": "data/dance.npz",
    "back_flip": "data/back_flip.npz",
    "forward_flip": "data/forward_flip.npz",
    "wave_motion": "data/wave_motion.npz",
}

在 onnx_file_dict 添加：

onnx_file_dict = {
    "normal": "data/amp_terrain.onnx",
    "recover": "data/recover.onnx",
    "dance": "data/dance.onnx",
    "amp_run": "data/amp_run.onnx",
    "normal_run": "data/model_normal.onnx",
    "back_flip": "data/back_flip.onnx",
    "forward_flip": "data/forward_flip.onnx",
    "noarm": "data/arm8.onnx",
    "wave_motion": "data/wave_motion.onnx",
}

launch 会把相对路径转成 package share 下的绝对路径，然后通过 ROS 参数传给 BxiExample：

{"/npz_file_dict": json.dumps(npz_file_dict)}
{"/onnx_file_dict": json.dumps(onnx_file_dict)}

3. 在 BxiExample 里加载 policy

打开：

src/bxi_example_py_elf3/bxi_example_py_elf3/bxi_example_demo.py

找到模型加载区域：

self.dance = DanceMotionPolicy(self.npz_file_dict["dance"], self.onnx_file_dict["dance"])

添加：

self.wave_motion = DanceMotionPolicyGravityIsaaclab(
    self.npz_file_dict["wave_motion"],
    self.onnx_file_dict["wave_motion"],
    start_frame=0,
)

为什么这里用 DanceMotionPolicyGravityIsaaclab：

• 它已经在当前文件 import。
• 它接收 .npz 和 .onnx。
• 它有 timestep、start_frame、end_frame。
• 它有 inference_step()、target_dof_pos、kps、kds。

如果你的模型和 DanceMotionPolicy 更匹配，也可以换成：

self.wave_motion = DanceMotionPolicy(
    self.npz_file_dict["wave_motion"],
    self.onnx_file_dict["wave_motion"],
    start_frame=0,
)

状态类里只要按该 policy 的实际字段取 kp/kd 即可。

4. 写 WaveMotionState

新建：

src/bxi_example_py_elf3/bxi_example_py_elf3/states/wave_motion_state.py

内容：

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Optional

from bxi_example_py_elf3.robot_state_base import MotorFrame, RobotControlState
from bxi_example_py_elf3.state_machine import StateBehavior, TransitionProfile

if TYPE_CHECKING:
    from bxi_example_py_elf3.bxi_example_demo import BxiExample


class WaveMotionState(RobotControlState):
    def __init__(
        self,
        name: str,
        state_id: int,
        policy_attr: str = "wave_motion",
        reset_on_finish: bool = True,
    ):
        super().__init__(name, state_id)
        self.policy_attr = policy_attr
        self.reset_on_finish = reset_on_finish
        self.playing = True

    def _policy(self, ctx: BxiExample) -> Any:
        return getattr(ctx, self.policy_attr)

    def on_prepare_enter(
        self,
        ctx: BxiExample,
        from_state: StateBehavior[BxiExample],
        transition: TransitionProfile,
    ) -> None:
        super().on_prepare_enter(ctx, from_state, transition)
        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.reset_loop(ctx)
        self.playing = True
        policy = self._policy(ctx)
        policy.timestep = policy.start_frame
        if hasattr(policy, "timeinit"):
            policy.timeinit = 0.0

    def get_first_frame(self, ctx: BxiExample) -> Optional[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:
            return None
        return self._motor_frame(qpos, policy.kps, policy.kds)

    def on_update(self, ctx: BxiExample, dt: float) -> None:
        if ctx.is_orientation_unsafe(ctx.current_quat_xyzw):
            ctx.request_state("zero_torque", trigger="safety")
            return

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

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

        if self.playing:
            policy.timestep += 1

    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

这个状态类展示了真实动作状态常见结构：

• policy_attr 指向 ctx.wave_motion。
• on_prepare_enter() 预热模型。
• on_enter() 重置 timestep。
• get_first_frame() 给过渡态使用。
• on_update() 推理并输出电机目标。
• 播放结束后请求回 normal。
• on_action() 支持暂停。

5. 导入 WaveMotionState

打开：

src/bxi_example_py_elf3/bxi_example_py_elf3/robot_states.py

添加：

from bxi_example_py_elf3.states.wave_motion_state import WaveMotionState

6. 在状态机中注册状态

打开：

src/bxi_example_py_elf3/config/elf3_state_machine.yaml

添加 event：

remote_events:
  wave_motion:
    slot: btn_10
    value: 6

在 normal 中添加入口：

states:
  normal:
    transitions:
      on_event:
        wave_motion:
          to: wave_motion
          transition: first_frame_switch

添加状态：

  wave_motion:
    behavior: WaveMotionState
    params:
      policy_attr: wave_motion
      reset_on_finish: true
    transitions:
      on_event:
        normal:
          to: normal
          transition: first_frame_switch
        zero_torque: zero_torque
        toggle_dance_pause:
          action: toggle_motion_pause

7. 绑定遥控器入口

打开：

src/remote_controller/config/xbox_default.yaml

键盘 source：

sources:
  keyboard:
    signals:
      keyboard.wave_motion: {from: keyboard.key, key: "v"}

control：

controls:
  keyboard.wave_motion: {type: bool, source: keyboard.wave_motion}

output：

outputs:
  edge:
    - output: btn_10=6
      when: [keyboard.wave_motion]

现在按键盘 v 会触发 wave_motion event。

8. 编译并验证

编译：

colcon build --symlink-install --packages-select bxi_example_py_elf3 remote_controller

加载：

source install/setup.bash

启动遥控器：

ros2 launch remote_controller remote_controller_keyboard.launch.py

启动 example：

ros2 launch bxi_example_py_elf3 example_demo_hw.launch.py

观察：

ros2 topic echo /simulation/state_machine_info

按 v 后应看到：

{
  "events": ["wave_motion"],
  "current": {
    "name": "wave_motion"
  }
}

9. 加入发布保护

如果这个动作是内部动作，打开：

src/bxi_example_py_elf3/config/release_protection.yaml

添加：

protected_states:
  wave_motion:
    behavior: WaveMotionState
    events: [wave_motion]
    model_keys: [wave_motion]
    files:
      - ../data/wave_motion.npz
      - ../data/wave_motion.onnx

model_keys: [wave_motion] 会影响：

• launch 中的 npz_file_dict["wave_motion"]。
• launch 中的 onnx_file_dict["wave_motion"]。
• bxi_example_demo.py 中的 self.wave_motion = ... 初始化代码块。

10. 常见问题

启动时报 KeyError: 'wave_motion'：

• launch 的 npz_file_dict / onnx_file_dict 没有加 key。
• 仿真 launch 加了，硬件 launch 没加，或反过来。
• 没重新 colcon build。

状态找不到 WaveMotionState：

• robot_states.py 没导入。
• states 子包没加入 setup.py。

进入状态后第一帧不对：

• get_first_frame() 时 policy 还没推理出 target_dof_pos。
• 检查 on_prepare_enter() 是否调用 ctx.preheat_model(policy)。

动作播放完没有返回：

• 检查 policy 是否有 end_frame。
• 检查 policy.timestep 是否递增。
• 检查 self.playing 是否被 action 暂停。

11. 本课检查清单

1. data/wave_motion.npz 和 data/wave_motion.onnx 存在。
2. example_demo.launch.py 和 example_demo_hw.launch.py 都添加 model key。
3. BxiExample 加载 self.wave_motion。
4. WaveMotionState 文件存在。
5. robot_states.py 导入 WaveMotionState。
6. elf3_state_machine.yaml 有 remote_events.wave_motion。
7. normal 能切到 wave_motion。
8. xbob_default.yaml 或自定义遥控器配置能输出 btn_10=6。
9. /simulation/state_machine_info 能看到 wave_motion。
10. 如果是内部动作，release_protection.yaml 写了 model_keys 和 files。

下一课：手把手 3：从零绑定按键并接入状态机。