Skip to content

Inference Architecture

konodoki edited this page Aug 2, 2026 · 8 revisions

推理框架与具名关节架构

推理框架的目标是让策略类只保留“如何组织观测、如何解码动作”这两类代码,同时统一 后端选择、关节映射、历史帧、预热和资源释放。配置使用 Python 类,不使用 YAML 描述模型或关节契约。

目录边界

bxi_example_py_elf3/
  framework/                   # 可移植的运控框架
    inference/
      api.py                   # InferenceFrame / PolicyOutput
      contract.py              # PolicyJointContract / JointInputBinding
      policy.py                # JointPolicy 和 Builder/Decoder 组合方式
      history.py               # 无逐帧申请的历史环形缓冲
      model.py                 # 可扩展 ModelArtifact / ModelSpec
      runtime.py               # 后端注册、选择和回退
      backends/                # ONNX Runtime / OpenVINO / RKNN
    joints/                    # Layout、JointParameterSet、默认值和命令解析器
    platform/                  # 机器人接入边界和关节 I/O 编码器
  policies/                    # ELF3 内置的具体策略
    joints.py                  # 模型关节契约与策略内部顺序
    normal.py
    amp.py
    beyondmimic.py
  mods/com.bxi.normal_depth/
    depth.py                   # 深度行走 Mod 私有策略

framework/ 不导入任何 ELF3 策略或关节表;policies/ 和各 Mod 私有策略可以依赖通用框架 及项目的 policies/joints.py。完整机器人布局由平台首个合法的具名状态快照建立,不由策略关节表 决定。迁移到另一台机器人时,可以整体复用 framework/,再提供目标机器人的平台适配器、 策略、缺失关节默认目标和 Mod。

JointParameterSet 是可移植的容器,因此定义在 framework/joints/parameters.pyELF3_ISAAC_PARAMETERS 是 ELF3 的具体数据,因此定义在 policies/joints.py。通用类不能 反向放进 ELF3 策略目录,否则其他机器人会被迫依赖 ELF3 包内实现。

八个关节布局边界

消息输入布局 Incoming Message Layout
  -> 完整机器人布局 Robot Layout
  -> 策略观测布局 Policy Observation Layout
  -> 模型观测布局 Model Observation Layout
  -> backend tensor
  -> 模型动作布局 Model Action Layout
  -> 策略动作布局 Policy Action Layout
  -> 状态自然输出布局 State Output Layout
  -> 可选的硬件固定布局 Hardware Fixed Layout

布局由不可变的 JointLayout 表示,关节名称是跨组件的语义主键。数字索引只在首次绑定 或消息顺序改变时编译,不能作为组件间的隐式协议。

  • Incoming Message Layout 可以任意排序,由平台归一化到稳定 Robot Layout。
  • Robot Layout 是平台当前实际具有的完整关节集合,也是最终 MotorFrame 的布局。
  • Policy Observation Layout 是策略从机器人状态读取的关节及顺序。
  • Model Observation Layout 是写入模型输入张量的真实关节顺序。
  • Model Action Layout 是模型输出张量中动作元素的真实关节顺序。
  • Policy Action Layout 是策略解码、补充或合成后对外输出的关节顺序,可以与前三者不同。
  • State Output Layout 是 Policy 与话题、IK、轨迹等来源合成后的自然输出,可多于或少于 Robot。
  • Hardware Layout 只用于不携带关节名称、只能接收固定数组顺序的接口。

多数单模型策略中 Policy Observation 就是 Model Observation,Model Action 也就是 Policy Action,因此只需要一个 PolicyJointContract。但它们不是概念上的同一个边界:例如 withoutarm.onnx 的模型输入和输出都是 15 个腰腿关节,AMP Policy 可以为剩余关节补默认 目标后对外输出完整 29 关节。此时 15 关节的 model observation/action layout 也必须用 JointLayout 明确声明,不能藏在 [:15] 里。

每周期的实际路径是:

完整 RobotObservation(N)
  -> PolicyJointContract 按名称选择 policy observation(P)
  -> 按具名 model observation 组织输入张量(M)
  -> backend 输出具名 model action(K)
  -> 策略解码/合成自己的 action JointTargetView(A)
  -> 可选 JointCommandComposer 合成 State Output MotorFrame(S)
  -> JointCommandResolver:S + 平台显式默认目标 -> 完整 Robot Frame(N)
  -> Transition:两端都先解析成 N,再插值
  -> 具名消息直接发布,或映射到 Hardware Layout

PMKASN 可以不同。policy observation 名称必须存在于 Robot Layout;状态输出少于 Robot 时必须为缺失硬件关节提供平台默认目标,状态输出多于 Robot 时按名称裁剪并 warning 一次。 模型输出、Policy Action、State Output 与 Hardware Layout 的完整区别见 关节布局与映射

策略类中的具名声明是唯一权威。下面这个模型的 metadata 明确采用 Isaac 顺序,因此不能 为了迎合 ELF3 消息顺序把 contract 写成 ELF3_POLICY_JOINTS

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


class WalkPolicy:
    joint_contract = PolicyJointContract(
        observation=ELF3_ISAAC_PARAMETERS.layout,
        action=ELF3_ISAAC_PARAMETERS.layout,
    )

模型 metadata 只能校验该声明,不能在运行时反过来修改契约。声明和 metadata 不一致应在 加载时失败,不能猜测顺序。没有 joint metadata 的资产必须由策略代码声明一个有名字的 JointLayout;同一策略内部的 NPZ、动作 clip 和 ONNX 如果顺序不同,也要分别声明布局并在 初始化时用 CompiledJointMap 编译映射。

禁止再维护 ISAAC_TO_MUJOCOMUJOCO_TO_ISAAC 之类的手写数字表。这类表没有方向类型, 也无法证明某个数组究竟属于哪一侧。数字索引只能由两个 JointLayout 编译得到,并且只能 保存在组件内部作为热路径缓存。

调用策略

平台每周期更新一个长期复用的 InferenceFrame

InferenceFrame(
    joints=joint_state_view,       # JointStateView(layout, position, velocity)
    quat_wxyz=quat_wxyz,
    angular_velocity=omega,
    command=command,
    base_linear_velocity=None,
    world_position=None,
    depth=None,
    depth_frame_id=None,
)

所有内置策略使用同一调用形式:

policy.reset(frame)
output = policy.step(frame, dt=0.02, advance=True)

output.joints.layout
output.joints.position
output.joints.kp
output.joints.kd
output.estimated_velocity
output.completed

调用者不再把裸 q, dq, quat, omega 按约定顺序传入。JointInputBinding 根据策略类声明 自动取得关节;来源布局与策略布局完全相同时直接使用原 view,不发生复制。布局不同或 来源包含额外关节时,映射写入策略长期持有的 float32 缓冲。

advance=False 用于预热和 Transition 双侧采样。策略不得在这种调用中推进时间、动作帧 或历史索引;HistoryBuffer.preview_append_into() 可以预览“如果追加当前帧”的连续输入, 但不修改环形缓冲状态。

一个可落地的 AMP 策略

下面是当前公共 API 的精简示例。关节契约、后端、稳定输入、历史和稳定输出均完整保留; 实际项目可继续加入相位、VAE 速度或不同 observation 字段。

import numpy as np

from bxi_example_py_elf3.framework.inference import (
    HistoryBuffer,
    InferenceFrame,
    InferenceRuntime,
    JointPolicy,
    ModelSpec,
    PolicyJointContract,
    PolicyOutput,
    default_runtime,
)
from bxi_example_py_elf3.policies.joints import ELF3_ISAAC_PARAMETERS


class AmpPolicy(JointPolicy):
    joint_contract = PolicyJointContract(
        observation=ELF3_ISAAC_PARAMETERS.layout,
        action=ELF3_ISAAC_PARAMETERS.layout,
    )

    def __init__(
        self,
        model_path: str,
        *,
        runtime: InferenceRuntime | None = None,
        backend: str = "auto",
    ) -> None:
        super().__init__()  # 创建具名输入绑定和稳定 PolicyOutput
        self._runtime = runtime or default_runtime()
        self._parameters = ELF3_ISAAC_PARAMETERS
        model = ModelSpec.portable_onnx(
            model_path,
            input_names=("obs",),
            output_names=("actions",),
        )
        self._backend = self._runtime.open_backend(model, backend=backend)

        count = self._parameters.layout.dof_num

        self._single = np.empty(96, dtype=np.float32)
        self._obs = np.empty((1, 960), dtype=np.float32)
        self._inputs = {"obs": self._obs}  # 字典和数组身份始终稳定
        self._history = HistoryBuffer(10, 96, dtype=np.float32)
        self._action = np.zeros(count, dtype=np.float32)
        self._scaled = np.empty(count, dtype=np.float32)
        self._joint_delta = np.empty(count, dtype=np.float32)
        self._gravity = np.empty(3, dtype=np.float32)

        # 直接在 JointPolicy 的输出缓冲中生成位置,不增加一次复制。
        self._target = self._target_buffer.position
        np.copyto(self._target, self._parameters.default_position)
        self.publish_output(
            self._target,
            self._parameters.kp,
            self._parameters.kd,
        )

    def reset(self, frame: InferenceFrame) -> None:
        joints = self.bind_joints(frame)  # 自动按类契约选取和排列
        self._action.fill(0.0)
        self._build_single(frame, joints.position, joints.velocity)
        self._history.fill(self._single)
        self._history.write_into(self._obs[0])
        np.copyto(self._target, self._parameters.default_position)

    def step(
        self,
        frame: InferenceFrame,
        dt: float,
        *,
        advance: bool = True,
    ) -> PolicyOutput:
        del dt
        joints = self.bind_joints(frame)
        self._build_single(frame, joints.position, joints.velocity)
        if advance:
            self._history.append(self._single)
            self._history.write_into(self._obs[0])
        else:
            self._history.preview_append_into(self._single, self._obs[0])

        raw = self._backend.run(self._inputs)["actions"]
        np.copyto(self._action, np.asarray(raw).reshape(-1))
        np.multiply(
            self._action,
            self._parameters.action_scale,
            out=self._scaled,
        )
        np.add(
            self._parameters.default_position,
            self._scaled,
            out=self._target,
        )
        return self.publish_output(self._target)

    def _build_single(self, frame, q, dq) -> None:
        if frame.command is None:
            raise ValueError("AMP policy requires frame.command")
        obs = self._single
        obs[0:3] = frame.angular_velocity
        self._project_gravity(frame.quat_wxyz, self._gravity)
        obs[3:6] = self._gravity
        obs[6:9] = frame.command
        np.subtract(
            q,
            self._parameters.default_position,
            out=self._joint_delta,
        )
        obs[9:38] = self._joint_delta
        obs[38:67] = dq
        obs[67:96] = self._action

    @staticmethod
    def _project_gravity(quat, out) -> None:
        w, x, y, z = quat
        out[0] = 2.0 * (w * y - x * z)
        out[1] = -2.0 * (w * x + y * z)
        out[2] = 2.0 * (x * x + y * y) - 1.0

    def close(self) -> None:
        self._backend.close()

这里没有四组脱离关节名的构造参数。ELF3_ISAAC_PARAMETERS.layout 同时作为模型契约, default_positionkpkdaction_scale 永远跟着这个布局。内置表按行维护:

from bxi_example_py_elf3.framework.joints import JointParameterSet


ELF3_ISAAC_PARAMETERS = JointParameterSet.from_rows(
    ELF3_ISAAC_JOINTS,
    (
        # name, default_position, kp, kd, action_scale
        ("l_shoulder_y_joint", 0.2, 54.224, 3.452, 0.231),
        ("r_shoulder_y_joint", 0.2, 54.224, 3.452, 0.231),
        # 其余关节继续逐行声明;名称和布局不一致会在导入时失败
    ),
)

真实完整的 29 行表位于 policies/joints.py。这里的省略只用于解释格式,不能原样复制为 可运行配置。若参数来自 ONNX metadata,也必须先校验 metadata 的 joint_names 与 contract 完全一致,再用 JointParameterSet.from_arrays(layout, ...) 把四列绑定到同一个布局。 禁止在推理类中维护四组无名称的平行数组。

这是“手写策略”的方式。希望进一步拆分单元测试时,也可以使用 InputBuilder + OutputDecoder + Policy 组合模板;两种方式共享完全相同的 InferenceFramePolicyOutput、关节契约和后端运行时。

29 策略与 31 关节机器人

本节只说明推理侧契约;完整的 29→31 补齐、31→29 裁剪、多来源合成和硬件输出规则见 关节布局与映射

31 关节状态进入 29 关节策略是安全的:策略 observation 是 source 的子集,绑定时按名称 抽取 29 个关节。反方向不能默认成立:29 个动作不足以形成 31 关节完整硬件命令。

policies/joints.py 中的 ELF3_POLICY_JOINTS 只固定旧模型的 29 维契约;完整 Robot Layout 直接来自首个合法的 ActuatorStates.name。因此机器人增加两个关节不会改变旧模型 的输入宽度。新 31 关节模型只需声明自己的 31 关节 observation/action layout。

旧模型没有输出的新关节必须在平台类配置中给出固定安全目标:

from bxi_example_py_elf3.framework.joints import (
    JointCommandDefaults,
    JointDefault,
)

ROBOT_COMMAND_DEFAULTS = JointCommandDefaults(
    {
        "left_gripper_joint": JointDefault(position=0.0, kp=20.0, kd=0.5),
        "right_gripper_joint": JointDefault(position=0.0, kp=20.0, kd=0.5),
    }
)

把它传给 RobotControlRuntime(command_defaults=ROBOT_COMMAND_DEFAULTS)。状态和策略产生的 MotorFrame 保留自己的自然布局;JointCommandResolver 在该布局第一次出现时完成以下 工作:

  1. 对模型输出中机器人不存在的关节发出一次 warning,并在编译映射中裁剪。
  2. 为 Robot Layout 中缺失的关节查找显式默认目标。
  3. 编译 source→robot 数字索引并缓存。
  4. 控制周期只向长期复用的完整输出缓冲写入,不查名称、不申请数组。

若模型缺少机器人关节且缺少任一默认项,首次绑定立即报错;不能静默补零或沿用上帧。 若模型输出多于机器人,则只发布名称交集并明确 warning。裁剪索引和复用缓冲只编译一次, 因此后续周期不会重复名称查找、日志输出或申请数组。

Transition 也先把源状态和目标状态各自解析到完整 Robot Layout,再执行插值。因此从 31 关节状态切到旧 29 关节状态时,新增关节会从源状态平滑过渡到显式默认目标,而不是 在切换开始时突跳。

输出到 ROS 或固定顺序硬件

携带名称的消息无需强制变成硬件固定顺序:

publish(names=frame.layout.names, position=frame.qpos)

框架最终返回的 MotorFrame 已经是完整 Robot Layout。只有绕过状态机、直接处理 PolicyOutput 且下游协议确实支持局部命令时,才使用:

encoder = NamedJointCommandEncoder(supports_partial=True)

不携带名称的硬件数组必须声明固定布局并编码:

encoder = FixedOrderJointCommandEncoder(hardware_layout)
hardware_target = encoder.encode(policy_output.joints)

encoder 在 source layout 首次出现时编译索引,之后复用输出数组。

后端模型与回退

ModelSpec 保存按优先级排列的开放 ModelArtifact,没有固定的“onnx/openvino/rknn 字段”。 增加后端只需实现新的 artifact 和 factory,不需要扩大策略接口。 策略类传入 ONNX 路径时统一使用 ModelSpec.portable_onnx(...),默认顺序是 rknn -> openvino -> onnxruntime:先找同名 .rknn,需要时在显式环境变量授权下 由 ONNX 转换,再回退到 OpenVINO 和 ONNX Runtime。

from bxi_example_py_elf3.framework.inference import (
    ModelSpec, OnnxArtifact, OpenVinoArtifact, RknnArtifact,
)

model = ModelSpec(
    artifacts=(
        RknnArtifact(
            "walk.rknn",
            source_onnx="walk.onnx",
            target="rk3588",
            input_shapes=(("obs", (1, 960)),),
            output_shapes=(("actions", (1, 29)),),
        ),
        OpenVinoArtifact("walk.onnx", device="CPU"),
        OnnxArtifact("walk.onnx"),
    ),
    input_names=("obs",),
    output_names=("actions",),
)

运行时按 artifacts 顺序选择第一个可用后端。缺少库时会给出安装命令并继续回退;显式 指定不可用后端则抛错。平台上真正的优先顺序应以 benchmark 为准,不能只根据“已安装” 推断速度。

OpenVINO 的 Intel GPU 插件可能枚举到 NVIDIA GPU。框架会拒绝这种设备,避免 CL_INVALID_COMMAND_QUEUE 导致进程崩溃;NVIDIA 应使用 ONNX Runtime CUDA/TensorRT。

RKNN 按需转换

ONNX 到 RKNN 是构建步骤,不是每帧推理步骤。默认绝不转换;只有显式设置环境变量才在 加载时检查缓存并按需转换:

BXI_RKNN_CONVERT_ON_LOAD=rk3588 python3 your_program.py

BXI_RKNN_CONVERT_ON_LOAD='{
  "target": "rk3588",
  "do_quantization": true,
  "dataset": "/data/calibration.txt",
  "force_rebuild": false,
  "outputs": ["actions"],
  "config": {"optimization_level": 3}
}' python3 your_program.py

转换缓存通过源模型和设置指纹判断是否有效。x86 是否能转换取决于 Rockchip 提供的 rknn-toolkit2 wheel;真正执行 .rknn 通常使用目标板上的 rknn-toolkit-lite2outputs 用于裁剪 RKNN 中实际需要的输出,也会进入缓存指纹。benchmark 默认只转换 actions;这既减少无用图节点,也避开 RKNN Toolkit 2.3.2 在合并不同 shape 的并行 Gather 输出时可能触发的优化器错误。需要多个运行时输出的策略必须显式列出全部名称。

RKNN 多输入 / 多输出注意事项

框架层面支持 N 个输入和 N 个输出,但 RKNN Lite 的运行时接口只返回按顺序排列的 list,不携带输出名。因此策略类必须把真实需要的输入、输出名称声明清楚:

model = ModelSpec.portable_onnx(
    "policy.onnx",
    input_names=("obs", "depth", "command"),
    output_names=("actions", "value", "latent"),
)

portable_onnx(...) 会从 ONNX 读取模型描述,生成同名 .rknn artifact,并把 output_names 作为 RKNN 转换时的 outputs。运行时再按相同顺序把 RKNN 返回的 list 映射回字典。只要转换时的 outputs 顺序和策略声明一致,N 输入 / N 输出不会因为框架 结构本身出错。

部署时仍需注意这些边界:

  1. 只声明策略实际使用的输出。只需要 actions 就不要保留调试输出,RKNN 会裁掉无用图, 转换更稳,推理也更轻。
  2. 多输出模型如果触发 RKNN Toolkit 优化器 bug,框架会回退 OpenVINO / ONNX Runtime; 这属于具体 Toolkit 和 ONNX 图结构兼容性问题,不是控制周期逻辑错误。
  3. RKNN 转换更适合静态 shape。动态 batch 或动态空间尺寸需要先导出成固定 shape,或通过 artifact / 环境变量明确转换参数。
  4. RKNN artifact 本身不保存 ONNX metadata。框架会从 source_onnx 读取关节名、增益、 shape 等信息;部署 .rknn 时建议保留旁边的 .onnx,至少保留同源 metadata 文件的能力。
  5. 如果某个输出只用于训练或离线分析,不要放进 output_names。运行时输出越少,后端选择 和跨平台 benchmark 越可控。

热路径性能规则

  1. JointLayout 和映射只在启动、首次绑定或消息名称顺序变化时创建。
  2. 相同 observation layout 使用零复制 fast path;不同布局使用 np.take(..., out=...)
  3. 状态输出布局第一次出现时编译命令映射;后续只写预分配的完整 Robot Frame。
  4. InferenceFramePolicyOutput、输入字典、输入数组、历史存储和目标数组长期复用。
  5. 不在控制周期创建 Session、InferRequest、深拷贝或按名称执行 list.index()
  6. ONNX Runtime 使用 I/O Binding;OpenVINO 使用持久 InferRequest 和共享输入 Tensor。
  7. 监控默认关闭,关闭时不读取性能时钟。

验证与 benchmark

# 所有发现到的模型 × 本机可用后端,子进程隔离 native driver 崩溃
python3 tools/benchmark/backend_benchmark.py --quick

# 策略级 Git 基线、分位延迟和 Python allocation
python3 tools/benchmark/inference_benchmark.py --baseline-ref HEAD

# 31→29 观测映射、29→31 默认补齐、31 关节重排和完整布局 fast path
python3 tools/benchmark/joint_mapping_benchmark.py

跨平台对比时保持模型、warmup、iterations、电源模式和 CPU 负载一致。报告目录 tools/benchmark/results/ 被 Git 忽略,benchmark 程序本身纳入版本管理。

Clone this wiki locally