Hands On Model State

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

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

data/*.npz
data/*.onnx

本课用一个具体例子：

wave_motion
WaveMotionState

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

这一课仍然保持低门槛：状态类继续直接写在 robot_states.py 里。不新建 states/ 子包，不新建 wave_motion_state.py，不改 setup.py。

0. 最终会改哪些文件

新增或放入模型文件：

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

修改文件：

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/remote_controller/config/xbox_default.yaml

如果这个动作属于内部高危动作，还会修改：

src/bxi_example_py_elf3/config/release_protection.yaml

1. 放入模型和动作数据

把文件放到：

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

data/ 目录里的文件会被当前 package 安装逻辑自动安装到 share/bxi_example_py_elf3/data，所以这里不需要改 setup.py。

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. 在 robot_states.py 里添加 WaveMotionState

打开：

src/bxi_example_py_elf3/bxi_example_py_elf3/robot_states.py

推荐把 WaveMotionState 放在 DanceState 附近，因为它们都是模型动作状态。

直接添加：

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 get_motor_frame(self, ctx: BxiExample, dt: float) -> Optional[MotorFrame]:
        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:
            policy.timestep += 1

        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

        frame = self.get_motor_frame(ctx, dt)
        if frame is not None:
            ctx.set_motor_target(*frame)

    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

这段代码使用的 Any、Optional、MotorFrame、RobotControlState、StateBehavior、TransitionProfile 在当前 robot_states.py 顶部已经有，不需要额外新建文件或导入子包。

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

• policy_attr 指向 ctx.wave_motion。
• on_prepare_enter() 预热模型。
• on_enter() 重置 timestep。
• get_first_frame() 给过渡态使用。
• get_motor_frame() 只负责计算电机帧。
• on_update() 做安全检查、播放结束检查和输出。
• on_action() 支持暂停。

如果模型需要速度输入

有些模型的推理接口需要 cmd_vel。这种状态不要直接读 ctx.raw_cmd_vel，而是在状态里调用 self.get_cmd_vel(ctx)：

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
    ctx.preheat_model(
        policy,
        with_cmd_vel=True,
        cmd_vel=self.get_cmd_vel(ctx),
    )

def get_motor_frame(self, ctx: BxiExample, dt: float) -> Optional[MotorFrame]:
    policy = self._policy(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,
    )
    return self._motor_frame(qpos, policy.kps, policy.kds)

并在状态 YAML 里声明：

  wave_motion:
    behavior: WaveMotionState
    speed_profile: normal

没有 speed_profile 时，get_cmd_vel(ctx) 会返回零速度。

如果这个模型还需要状态私有的速度滤波，不要重写 get_cmd_vel()，而是重写：

def process_cmd_vel(
    self,
    ctx: BxiExample,
    cmd_vel: np.ndarray,
) -> Optional[np.ndarray]:
    cmd_vel[1] = 0.0
    return cmd_vel

get_cmd_vel() 会自动把返回值写入 ctx.current_cmd_vel。

如果模型状态需要额外订阅

例如模型状态要订阅一个外部播放速度、目标点或动作选择话题，实现 on_bind(ctx)：

def on_bind(self, ctx: BxiExample) -> None:
    self.selector_sub = ctx.create_subscription(
        String,
        "wave_motion_selector",
        self._selector_callback,
        10,
    )

def _selector_callback(self, msg: String) -> None:
    self.selected_clip = msg.data

on_bind(ctx) 只执行一次，适合创建 ROS 订阅、client 或 timer。__init__() 里只放 policy_attr、playing 这类普通私有变量。

如果使用上面的 String 示例，记得在 robot_states.py 顶部加：

from std_msgs.msg import String

5. 在状态机中注册状态

打开：

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

注意：

• behavior: WaveMotionState 必须和 Python 类名一致。
• params.policy_attr: wave_motion 表示状态会访问 ctx.wave_motion。
• ctx.wave_motion 来自上一节在 BxiExample 里添加的 self.wave_motion = ...。

6. 绑定遥控器入口

打开：

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 会触发：

keyboard.wave_motion
  -> btn_10=6
  -> remote_events.wave_motion
  -> normal 切到 wave_motion

7. 编译并验证

编译：

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.launch.py

观察：

ros2 topic echo /simulation/state_machine_info

按 v 后应看到：

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

如果你在真实硬件上验证，把 launch 换成：

ros2 launch bxi_example_py_elf3 example_demo_hw.launch.py

硬件验证前必须先在仿真里确认动作安全。

8. 加入发布保护

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

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 = ... 初始化代码块。

如果多个受保护状态共用一个基类或辅助类，可以把 behavior 写成数组：

behavior:
  - SomeSharedBaseState
  - WaveMotionState

但入门阶段一般只写当前状态类即可。

9. 常见问题

启动时报 KeyError: 'wave_motion'：

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

状态找不到 WaveMotionState：

• WaveMotionState 没写进 robot_states.py。
• YAML 里的 behavior 拼错了。
• 修改后没重新 colcon build 和 source install/setup.bash。

进入状态后第一帧不对：

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

动作播放完没有返回：

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

10. 本课检查清单

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 直接写在 robot_states.py 里。
5. 没有新建 states/ 文件夹。
6. 没有修改 setup.py。
7. elf3_state_machine.yaml 有 remote_events.wave_motion。
8. normal 能切到 wave_motion。
9. xbox_default.yaml 或自定义遥控器配置能输出 btn_10=6。
10. /simulation/state_machine_info 能看到 wave_motion。
11. 如果是内部动作，release_protection.yaml 写了 model_keys 和 files。

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