Hands On Custom State

手把手 1：从零写一个自定义状态

本课目标：写一个真正能运行的 SinWaveState。它会让一个关节按 sin() 摆动。

本课先不绑定遥控器。我们先让状态本身能跑起来；下一课会把这个思路扩展到真实 .npz / .onnx 模型动作，之后再接键盘、手柄和状态机事件。

0. 你最终会得到什么

新增文件：

src/bxi_example_py_elf3/bxi_example_py_elf3/states/__init__.py
src/bxi_example_py_elf3/bxi_example_py_elf3/states/sin_wave_state.py

修改文件：

src/bxi_example_py_elf3/setup.py
src/bxi_example_py_elf3/bxi_example_py_elf3/robot_states.py
src/bxi_example_py_elf3/config/elf3_state_machine.yaml

最终状态类：

SinWaveState

最终状态名：

sin_wave

1. 新建 states 子包

新建目录：

src/bxi_example_py_elf3/bxi_example_py_elf3/states/

新建空文件：

src/bxi_example_py_elf3/bxi_example_py_elf3/states/__init__.py

为什么要拆出来：

• robot_states.py 已经有很多状态。
• 新状态继续堆进去会越来越难读。
• 每个动作状态独立文件，后续发布保护也更容易演进到“删文件”。

2. 修改 setup.py

打开：

src/bxi_example_py_elf3/setup.py

找到：

packages=[package_name,
          f'{package_name}.inference',
          f'{package_name}.utils',
          ],

改成：

packages=[
    package_name,
    f'{package_name}.inference',
    f'{package_name}.utils',
    f'{package_name}.states',
],

为什么要改：

• Python 源码目录存在，不代表安装包里包含它。
• colcon build 后运行的是安装后的包。
• 不加 f'{package_name}.states'，运行时可能找不到 bxi_example_py_elf3.states。

3. 写最小 SinWaveState

新建：

src/bxi_example_py_elf3/bxi_example_py_elf3/states/sin_wave_state.py

先写最小可运行版本：

from __future__ import annotations

import math
from typing import TYPE_CHECKING, Optional

import numpy as np

from bxi_example_py_elf3.robot_state_base import MotorFrame, RobotControlState

if TYPE_CHECKING:
    from bxi_example_py_elf3.bxi_example_demo import BxiExample


class SinWaveState(RobotControlState):
    def __init__(
        self,
        name: str,
        state_id: int,
        joint: int = 22,
        amplitude: float = 0.4,
        frequency: float = 1.0,
    ):
        super().__init__(name, state_id)
        self.joint = joint
        self.amplitude = amplitude
        self.frequency = frequency
        self.elapsed = 0.0
        self.base_qpos: Optional[np.ndarray] = None

    def on_enter(self, ctx: BxiExample) -> None:
        self.reset_loop(ctx)
        self.elapsed = 0.0
        self.base_qpos = ctx.joint_nominal_pos.copy()

    def get_first_frame(self, ctx: BxiExample) -> Optional[MotorFrame]:
        qpos = ctx.joint_nominal_pos.copy()
        return self._motor_frame(qpos, ctx.joint_kp, ctx.joint_kd)

    def on_update(self, ctx: BxiExample, dt: float) -> None:
        if self.base_qpos is None:
            self.base_qpos = ctx.joint_nominal_pos.copy()

        self.elapsed += dt
        qpos = self.base_qpos.copy()
        qpos[self.joint] += self.amplitude * math.sin(
            2.0 * math.pi * self.frequency * self.elapsed
        )
        ctx.set_motor_target(qpos, ctx.joint_kp, ctx.joint_kd)

这个状态生命周期：

进入状态 on_enter
  -> 重置 elapsed
  -> 保存基础姿态

每个控制周期 on_update
  -> elapsed += dt
  -> 计算 sin
  -> 输出电机目标

4. 导入状态类

打开：

src/bxi_example_py_elf3/bxi_example_py_elf3/robot_states.py

在 import 区域加入：

from bxi_example_py_elf3.states.sin_wave_state import SinWaveState

原因：

• build_robot_states() 是从已导入的 RobotControlState 子类中收集状态类。
• 只创建文件但不导入，YAML 写 behavior: SinWaveState 也找不到类。

5. 先不接遥控器，直接把它设成初始状态

这是最适合新手的验证方式：先证明状态代码能跑，再接遥控器。

打开：

src/bxi_example_py_elf3/config/elf3_state_machine.yaml

临时把：

initial_state: zero_torque

改成：

initial_state: sin_wave

再在 states: 下面添加：

  sin_wave:
    behavior: SinWaveState
    params:
      joint: 22
      amplitude: 0.4
      frequency: 1.0
    transitions:
      on_event:
        zero_torque: zero_torque
        normal:
          to: normal
          transition: soft_switch

现在启动后会直接进 sin_wave。

验证完记得把 initial_state 改回：

initial_state: zero_torque

6. 编译并启动

编译：

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

加载环境：

source install/setup.bash

启动 example：

ros2 launch bxi_example_py_elf3 example_demo_hw.launch.py

看状态机信息：

ros2 topic echo /simulation/state_machine_info

你应该看到：

{
  "current": {
    "name": "sin_wave"
  }
}

如果状态没切进去，先不要继续加按键。先解决状态注册问题。

7. 加 YAML 参数

现在我们让读者明白：状态代码写一次，动作细节从 YAML 调。

SinWaveState.__init__() 已经有：

joint: int = 22
amplitude: float = 0.4
frequency: float = 1.0

所以 YAML 可以改：

  sin_wave:
    behavior: SinWaveState
    params:
      joint: 24
      amplitude: 0.25
      frequency: 1.5

规则：

• name 和 state_id 不写进 params，框架自动传。
• params 的 key 必须和 __init__() 参数名一致。
• 私有运行时变量不要放 YAML，例如 elapsed 应该只在状态内部维护。

8. 添加自动返回

现在给状态加生命周期：运行 3 秒自动回 normal。

不用改 Python，直接在 YAML 写：

  sin_wave:
    behavior: SinWaveState
    params:
      joint: 22
      amplitude: 0.4
      frequency: 1.0
    transitions:
      on_event:
        zero_torque: zero_torque
        normal:
          to: normal
          transition: soft_switch
      after:
        - seconds: 3.0
          to: normal
          transition: soft_switch

含义：

进入 sin_wave 后计时
  -> 3 秒后自动切 normal

这一步教会读者：能放配置的状态图关系，不要硬编码进状态类。

9. 添加代码内主动退出

有些结束条件必须写代码，例如播放帧结束、安全检查、模型输出触发。

我们给 SinWaveState 加一个可选 duration 参数：

def __init__(
    self,
    name: str,
    state_id: int,
    joint: int = 22,
    amplitude: float = 0.4,
    frequency: float = 1.0,
    duration: float = 0.0,
):
    super().__init__(name, state_id)
    self.joint = joint
    self.amplitude = amplitude
    self.frequency = frequency
    self.duration = duration
    self.elapsed = 0.0
    self.base_qpos: Optional[np.ndarray] = None

在 on_update() 末尾加：

if self.duration > 0.0 and self.elapsed >= self.duration:
    ctx.request_state(
        "normal",
        trigger="sin_wave_finished",
        transition="soft_switch",
    )

YAML：

params:
  joint: 22
  amplitude: 0.4
  frequency: 1.0
  duration: 3.0

什么时候用代码退出：

• 动作文件播放完。
• 模型输出满足条件。
• 姿态不安全。
• 传感器触发。

什么时候用 YAML after：

• 固定时间自动退出。

10. 添加暂停 action

现在给状态加一个不切状态的动作：暂停/继续摆动。

修改 __init__()：

self.playing = True

修改 on_enter()：

self.playing = True

修改 on_update()：

if self.playing:
    self.elapsed += dt

添加：

def on_action(self, ctx: BxiExample, action_name: str) -> bool:
    if action_name != "toggle_sin_pause":
        return False

    self.playing = not self.playing
    return True

YAML：

  sin_wave:
    behavior: SinWaveState
    transitions:
      on_event:
        toggle_dance_pause:
          action: toggle_sin_pause

这一步教会读者：

• 切状态用 to。
• 不切状态但改变状态内部变量，用 action。

11. 添加安全退出

危险动作都应该有状态内部安全判断。

在 on_update() 开头加：

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

完整顺序建议：

安全检查
  -> 更新时间
  -> 计算目标
  -> 输出电机目标
  -> 检查是否结束

12. 最终版 SinWaveState

最终代码可以整理成：

from __future__ import annotations

import math
from typing import TYPE_CHECKING, Optional

import numpy as np

from bxi_example_py_elf3.robot_state_base import MotorFrame, RobotControlState

if TYPE_CHECKING:
    from bxi_example_py_elf3.bxi_example_demo import BxiExample


class SinWaveState(RobotControlState):
    def __init__(
        self,
        name: str,
        state_id: int,
        joint: int = 22,
        amplitude: float = 0.4,
        frequency: float = 1.0,
        duration: float = 0.0,
    ):
        super().__init__(name, state_id)
        self.joint = joint
        self.amplitude = amplitude
        self.frequency = frequency
        self.duration = duration
        self.elapsed = 0.0
        self.playing = True
        self.base_qpos: Optional[np.ndarray] = None

    def on_enter(self, ctx: BxiExample) -> None:
        self.reset_loop(ctx)
        self.elapsed = 0.0
        self.playing = True
        self.base_qpos = ctx.joint_nominal_pos.copy()

    def get_first_frame(self, ctx: BxiExample) -> Optional[MotorFrame]:
        qpos = ctx.joint_nominal_pos.copy()
        return self._motor_frame(qpos, ctx.joint_kp, ctx.joint_kd)

    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

        if self.base_qpos is None:
            self.base_qpos = ctx.joint_nominal_pos.copy()

        if self.playing:
            self.elapsed += dt

        qpos = self.base_qpos.copy()
        qpos[self.joint] += self.amplitude * math.sin(
            2.0 * math.pi * self.frequency * self.elapsed
        )
        ctx.set_motor_target(qpos, ctx.joint_kp, ctx.joint_kd)

        if self.duration > 0.0 and self.elapsed >= self.duration:
            ctx.request_state(
                "normal",
                trigger="sin_wave_finished",
                transition="soft_switch",
            )

    def on_action(self, ctx: BxiExample, action_name: str) -> bool:
        if action_name != "toggle_sin_pause":
            return False

        self.playing = not self.playing
        return True

13. 本课检查清单

1. states/__init__.py 存在。
2. states/sin_wave_state.py 存在。
3. setup.py packages 包含 bxi_example_py_elf3.states。
4. robot_states.py 导入 SinWaveState。
5. elf3_state_machine.yaml 里 states.sin_wave.behavior 是 SinWaveState。
6. 初始状态临时设为 sin_wave 能运行。
7. /simulation/state_machine_info 能看到 current.name = sin_wave。
8. 验证完把 initial_state 改回 zero_torque。

下一课：手把手 2：从零添加模型动作状态。