Hands On Model State

手把手 2：把模型动作封装成 Mod

本课在独立 Mod 中放置 .npz 和 .onnx，通过可配置加载策略的资源交给状态使用。

1. 目录

mods/com.example.motion/
  mod.yaml
  plugin.py
  state.py
  assets/
    motion.npz
    motion.onnx

2. 先声明模型资产的真实关节顺序

不要从机器人消息顺序推断模型顺序。先查看 ONNX 的 joint_names metadata、训练导出脚本 和配套 NPZ 的定义，确定模型 observation 与 action 的真实具名布局：

from bxi_example_py_elf3.framework.inference import PolicyJointContract
from bxi_example_py_elf3.policies.joints import ELF3_ISAAC_JOINTS


class ExampleMotionPolicy:
    joint_contract = PolicyJointContract(
        observation=ELF3_ISAAC_JOINTS,
        action=ELF3_ISAAC_JOINTS,
    )

当前 DanceMotionPolicyGravityIsaaclab 已做了这个声明，并在加载时检查 ONNX metadata； back-flip、forward-flip 等 Isaac 模型可以直接使用它。recover 一类采用另一顺序的模型应使用 对应的 Mjlab Policy。声明错了就应在加载阶段失败，不能靠 ISAAC_TO_MUJOCO、 MUJOCO_TO_ISAAC 或看不出含义的数字切片补救。

如果模型只输入/输出局部关节，也要单独声明真实子布局。例如 15 关节 no-arm 模型使用 ELF3_LOWER_BODY_JOINTS。Policy 可以把这 15 个模型动作与默认姿态、话题、IK 或程序轨迹 合成为更大的 Policy Action，但 Model Action Layout 与 Policy Action Layout 必须分开说明。

NPZ 通常不带关节名。此时训练/导出方必须把其顺序作为代码中的 JointLayout 固化；若它 与 ONNX 不同，只能在初始化阶段用 CompiledJointMap 按名称编译一次映射。

关节顺序不仅约束 observation 和 action，也约束 default position、kp、kd 和 action scale。 不要在策略类中分别放四组 29 项裸数组。固定参数按关节逐行声明：

from bxi_example_py_elf3.framework.joints import JointParameterSet


PARAMETERS = JointParameterSet.from_rows(
    MODEL_JOINTS,
    (
        # name, default_position, kp, kd, action_scale
        ("joint_a", 0.0, 50.0, 2.0, 0.25),
        ("joint_b", 0.1, 40.0, 1.5, 0.20),
    ),
)

然后让 joint_contract 引用 PARAMETERS.layout，推理代码只读取 self._parameters.default_position/kp/kd/action_scale。若四列来自 ONNX metadata，先严格 校验 metadata 的 joint_names，再通过 JointParameterSet.from_arrays(MODEL_JOINTS, ...) 绑定。模型没有 joint_names 时，策略代码必须提供真实 MODEL_JOINTS；不能从数组长度、 机器人话题顺序或类名猜测。

3. 状态

复用框架的 MotionReplayState：

from bxi_example_py_elf3.policies import (
    DanceMotionPolicyGravityIsaaclab,
)
from bxi_example_py_elf3.framework.mod_api import ResourceHandle
from bxi_example_py_elf3.framework.mod_api import MotionReplayState


class ExampleMotionState(
    MotionReplayState[DanceMotionPolicyGravityIsaaclab]
):
    def __init__(self, name, state_id, policy):
        super().__init__(
            name,
            state_id,
            policy,
            finish_state="com.bxi.basic_actions/normal",
            finish_trigger="example_motion_finished",
            end_frame_trim=20,
            end_transition={
                "profile": "dual_running_blend",
                "duration": 0.6,
                "sample_from": True,
            },
        )

基类负责重置策略、预热、进入帧、运行采样、播放结束请求 finish_state 和暂停 action。结束目标必须显式声明；框架不会猜测哪个状态是当前机器人的“normal”。

MotionReplayState 从 policy.output.joints.layout 取得策略的具名输出布局，不假设固定 29 关节。策略输出少于状态最终需要的关节时，用 多来源关节命令合成接入话题、IK 或程序控制；状态输出与机器人 29/31/N 布局不同时，由关节布局与映射中的规则处理。

4. 注册资源及加载策略

plugin.py：

from bxi_example_py_elf3.policies import (
    DanceMotionPolicyGravityIsaaclab,
)
from bxi_example_py_elf3.framework.mod_api import (
    ModDefinition,
    ModLoadContext,
    ResourceKey,
    ResourceLoadContext,
)

from .state import ExampleMotionState

POLICY = ResourceKey[DanceMotionPolicyGravityIsaaclab](
    "com.example.motion/policy"
)


def _load(context: ResourceLoadContext):
    return DanceMotionPolicyGravityIsaaclab(
        str(context.asset("assets/motion.npz")),
        str(context.asset("assets/motion.onnx")),
        start_frame=0,
    )


def create_mod(context: ModLoadContext) -> ModDefinition:
    context.register_resource(POLICY, _load, loading="lazy")
    policy = context.resource(POLICY)
    return ModDefinition(
        state_factories={
            "motion": lambda state: ExampleMotionState(
                state.name, state.state_id, policy
            )
        }
    )

这里选择 lazy，所以加载 Mod 不会创建推理器；状态首次访问 handle 时才加载 文件。如果模型初始化时间可能超过控制周期，改为 loading="eager"，框架会在 控制循环启动前完成加载，失败则直接终止启动。

5. 清单

schema: 1
id: com.example.motion
name: 模型动作
version: 1.0.0
api: ">=2,<3"
enable: true
entrypoint: plugin:create_mod
visibility: public
requires:
  - {id: com.bxi.basic_actions, version: ">=1,<2"}
conflicts: []
python_exports: []
runtime_requirements:
  python: []
  ros: []
  system: []

events:
  activate: {slot: btn_10, value: 8}
  toggle_pause: {slot: btn_9, value: 1}

states:
  motion:
    manifest:
      label: 模型动作
      priority: 100
      group: Customer
      icon: animation
      confirm: true
      confirm_message: 请确保周围安全

routes:
  - {from: com.bxi.basic_actions/normal, event: activate, to: motion, transition: soft_switch}
  - {from: motion, event: com.bxi.basic_actions/normal, to: com.bxi.basic_actions/normal, transition: dual_running_blend}
  - {from: motion, event: com.bxi.basic_actions/zero_torque, to: com.bxi.basic_actions/zero_torque}

actions:
  - {from: motion, event: toggle_pause, action: toggle_pause, manifest: {label: 暂停/继续}}

Transition 是两个状态之间的临时控制过程：状态图已经决定要从哪里切到哪里，Transition 再决定切换期间怎样生成 qpos/kp/kd。模型第一帧和当前电机帧可能差异很大，因此通常需要保持、增益渐变或双状态混合，避免直接换帧造成突跳。上例进入前使用短暂保持，返回 normal 时混合动作模型与行走模型的运行帧。

6. 调试顺序

1. 先离线加载 Mod，确认资源仍是未加载状态。
2. 仿真进入动作，确认 on_prepare() 预热成功。
3. 检查第一帧、结束帧和 end_frame_trim。
4. 用乱序具名状态和带额外关节的状态验证模型输入仍按声明布局排列。
5. 检查安全退出与 normal 返回。
6. 最后再进入真机低增益验证。

下一课：绑定按键。