Skip to content

Inference Architecture

konodoki edited this page Jul 28, 2026 · 8 revisions

推理框架架构

推理框架位于 bxi_example_py_elf3/inference/,目标是在不增加策略开发门槛的前提下, 统一管理模型后端、输入输出、历史帧、性能监控和资源释放。模型策略只保留观测构造和 动作解码;ONNX Runtime、OpenVINO、RKNN 以及未来后端的加载和回退都由公共运行时负责。

配置只使用 Python 类,不使用 YAML。ModelSpec 只描述逻辑输入输出和按优先级排列的 模型产物,不包含固定的后端字段,因此增加新后端不需要修改 ModelSpec

总体数据流

RobotObservation / 传感器数据
  -> InferenceFrame
  -> InputBuilder.build_into()       # 原地更新稳定输入缓冲
  -> Mapping[str, numpy.ndarray]
  -> Policy
       -> InferenceRuntime.open_backend(ModelSpec)
            -> 按 artifacts 顺序检查 BackendFactory
            -> 加载第一个可用 InferenceBackend
            -> 缺库、模型不兼容或初始化失败时继续回退
       -> InferenceBackend.run()
  -> OutputDecoder.decode_into()     # 原地解码到复用输出
  -> PolicyOutput
  -> State / MotorFrame

后端选择与策略数学解耦。策略看到的始终是具名 NumPy 输入和输出,不需要导入 onnxruntimeopenvinorknnlite

目录与责任

文件 责任
api.py 后端无关的 InferenceFramePolicyOutputFloatArray
model.py ModelArtifact、各后端产物类和 ModelSpec
runtime.py 后端注册、按顺序选择、初始化失败回退和统一告警
policy.py InputBuilder -> backend -> OutputDecoder 组合模板
history.py 固定容量、无逐帧分配的历史环形缓冲
monitor.py 可选的 input/backend/output/total 分段耗时统计
backends/base.py 后端与工厂的最小接口
backends/onnxruntime.py ONNX Runtime 会话和 I/O Binding
backends/openvino.py OpenVINO CompiledModel、持久 InferRequest 和共享输入 Tensor
backends/rknn.py Rockchip RKNN Lite 加载和推理
backends/rknn_builder.py 显式授权的 ONNX 到 RKNN 缓存转换

已有的 normal.pyamp.pybeyondmimic.py 和深度策略是具体策略实现。它们可以 继续提供原有的领域接口,但模型加载应统一交给 InferenceRuntime

最简单的模型配置

只有 ONNX Runtime 时:

from bxi_example_py_elf3.inference import ModelSpec

model = ModelSpec.onnx(
    "model.onnx",
    input_names=("obs",),
    output_names=("actions",),
)

同一个 ONNX 文件优先尝试 OpenVINO,再回退到 ONNX Runtime:

model = ModelSpec.portable_onnx(
    "model.onnx",
    input_names=("obs",),
    output_names=("actions",),
    openvino_device="AUTO",
)

需要精确控制优先级时直接排列产物:

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

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

InferenceRuntime 严格按 artifacts 顺序尝试。部署平台的最优顺序应由该平台上的 benchmark 决定,不应仅因为安装了某个库就假设它更快。也可以在调试时显式传入 backend="onnxruntime""openvino""rknn"

策略组合接口

新策略推荐拆成三个小对象:

  • InputBuilder 拥有固定输入数组,并在 build_into() 中原地写入。
  • OutputDecoder 拥有一个可复用的 PolicyOutput,并在 decode_into() 中原地更新。
  • Policy 只负责 build -> run -> decode、预热、监控和关闭。
policy = Policy(
    model,
    input_builder=MyInputBuilder(),
    output_decoder=MyOutputDecoder(),
    name="walk",
)

policy.prepare(frame)
output = policy.step(frame, dt=0.02)
policy.close()

prepare() 会 reset、用 advance=False 构造输入并执行预热。advance=False 的调用 不得推进策略时间、历史索引或动作帧;它用于状态准备和双状态采样。每个策略实例拥有 自己的后端实例,使用结束后必须调用 close()

完整示例:手写一个 AMP 推理类

下面实现一个典型的 AMP 行走策略:29 个关节,每帧观测为 omega(3) + gravity(3) + command(3) + q(29) + dq(29) + last_action(29) = 96 维,模型读取连续 10 帧,即输入 shape 为 (1, 960),输出为 29 维动作。

示例是可运行的真实接口,不省略后端加载、历史初始化、reset、推理和资源释放。为了突出 架构,关节重排等具体机器人差异通过调用方传入的数据顺序解决。

from collections.abc import Mapping

import numpy as np

from bxi_example_py_elf3.inference import (
    HistoryBuffer,
    InferenceFrame,
    InputBuilder,
    ModelSpec,
    OutputDecoder,
    Policy,
    PolicyOutput,
)


class AmpInputBuilder(InputBuilder):
    """AMP 只负责定义观测,不接触任何推理后端。"""

    def __init__(self, default_position: np.ndarray) -> None:
        self.default_position = np.asarray(
            default_position, dtype=np.float32
        ).reshape(29)

        # 所有热路径数组都只在构造时申请一次。
        self._obs = np.empty((1, 960), dtype=np.float32)
        self._single_obs = np.empty(96, dtype=np.float32)
        self._joint_delta = np.empty(29, dtype=np.float32)
        self._last_action = np.zeros(29, dtype=np.float32)
        self._gravity = np.empty(3, dtype=np.float32)
        self._history = HistoryBuffer(10, 96, dtype=np.float32)
        self._inputs = {"obs": self._obs}

    @property
    def inputs(self) -> Mapping[str, np.ndarray]:
        # 字典和其中的数组对象在整个策略生命周期内保持不变,方便后端绑定。
        return self._inputs

    def reset(self, frame: InferenceFrame) -> None:
        self._last_action.fill(0.0)
        self._build_single(frame)
        self._history.fill(self._single_obs)
        self._history.write_into(self._obs[0])

    def build_into(
        self,
        frame: InferenceFrame,
        dt: float,
        *,
        advance: bool,
    ) -> None:
        del dt  # 这个 AMP 模型没有显式时间输入。
        self._build_single(frame)
        if advance:
            self._history.append(self._single_obs)
            self._history.write_into(self._obs[0])

    def accept_action(self, action: np.ndarray) -> None:
        """本次推理完成后保存动作,供下一帧观测使用。"""

        np.copyto(self._last_action, action)

    def _build_single(self, frame: InferenceFrame) -> None:
        obs = self._single_obs
        obs[0:3] = frame.angular_velocity
        self._project_gravity(frame.quat_wxyz, self._gravity)
        obs[3:6] = self._gravity
        if frame.command is None:
            obs[6:9] = 0.0
        else:
            obs[6:9] = frame.command
        np.subtract(frame.q, self.default_position, out=self._joint_delta)
        obs[9:38] = self._joint_delta
        obs[38:67] = frame.dq
        obs[67:96] = self._last_action

    @staticmethod
    def _project_gravity(quat: np.ndarray, out: np.ndarray) -> None:
        """把世界系重力方向投影到机身系,结果原地写入 out。"""

        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


class AmpOutputDecoder(OutputDecoder):
    """把模型动作原地转换为关节目标位置。"""

    def __init__(
        self,
        default_position: np.ndarray,
        action_scale: np.ndarray,
    ) -> None:
        self._default = np.asarray(default_position, dtype=np.float32).reshape(29)
        self._scale = np.asarray(action_scale, dtype=np.float32).reshape(29)
        self.action = np.zeros(29, dtype=np.float32)
        self._scaled = np.empty(29, dtype=np.float32)
        self._output = PolicyOutput(
            joint_position=np.empty(29, dtype=np.float32)
        )

    @property
    def output(self) -> PolicyOutput:
        return self._output

    def reset(self) -> None:
        self.action.fill(0.0)
        np.copyto(self._output.joint_position, self._default)

    def decode_into(self, outputs: Mapping[str, np.ndarray]) -> None:
        np.copyto(self.action, np.asarray(outputs["actions"]).reshape(-1))
        np.multiply(self.action, self._scale, out=self._scaled)
        np.add(self._default, self._scaled, out=self._output.joint_position)


class AmpPolicy:
    """给机器人控制代码使用的极薄领域接口。"""

    def __init__(
        self,
        model_path: str,
        default_position: np.ndarray,
        action_scale: np.ndarray,
        *,
        backend: str = "auto",
    ) -> None:
        builder = AmpInputBuilder(default_position)
        decoder = AmpOutputDecoder(default_position, action_scale)
        model = ModelSpec.portable_onnx(
            model_path,
            input_names=("obs",),
            output_names=("actions",),
        )
        self._builder = builder
        self._decoder = decoder
        self._policy = Policy(
            model,
            builder,
            decoder,
            name="amp",
            backend=backend,
        )

    def prepare(self, frame: InferenceFrame) -> None:
        self._policy.prepare(frame)

    def reset(self, frame: InferenceFrame) -> None:
        self._policy.reset(frame)

    def step(
        self,
        frame: InferenceFrame,
        dt: float = 0.02,
        *,
        advance: bool = True,
    ) -> np.ndarray:
        output = self._policy.step(frame, dt, advance=advance)
        if advance:
            self._builder.accept_action(self._decoder.action)
        return output.joint_position

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

调用方只需创建一次策略和 InferenceFrame,控制周期中原地更新 frame 内已有数组:

amp = AmpPolicy("walk.onnx", default_position, action_scale)
frame = InferenceFrame(q, dq, quat_wxyz, omega, command)

amp.prepare(frame)
try:
    while running:
        # q、dq、quat_wxyz、omega、command 均为复用的 float32 数组。
        target_position = amp.step(frame)
        send_joint_target(target_position)
finally:
    amp.close()

这个例子里的 AMP 领域代码只决定三件事:单帧观测如何排列、历史在何时推进、模型输出 如何变成关节目标。OpenVINO/ONNX Runtime/RKNN 的选择、缺库回退、模型预热、性能监控 和资源关闭都由公共框架处理。若换成带相位或 VAE 速度输出的 AMP,只需扩展 builder 和 decoder,无需修改 Policy 或任何后端。

后端选择与回退

默认 Registry 注册 RKNN、OpenVINO 和 ONNX Runtime 工厂,但注册顺序不决定选择顺序; 真正顺序来自 ModelSpec.artifacts

每个工厂先执行只属于该后端的可用性检查:

后端 主要检查 缺失时的行为
ONNX Runtime onnxruntime、模型文件、provider 给出 pip install onnxruntime 提示并回退
OpenVINO openvino、模型文件、设备初始化 给出安装提示或初始化原因并回退
RKNN RKNN 文件/转换结果、Lite2、Rockchip target 给出官方 wheel 提示并回退

自动模式下,运行时会汇总跳过原因,并只对相同回退路径告警一次。显式指定一个不可用 后端时会抛出 BackendUnavailableError,错误中保留安装建议。

OpenVINO 设备边界

OpenVINO 的 Intel GPU 插件在某些 OpenCL 环境中会枚举到 NVIDIA GPU,但不能可靠执行。 框架会拒绝显式使用这种非 Intel GPU;AUTO 遇到它时忽略该设备并安全选择 CPU 或其他 受支持设备。NVIDIA GPU 应使用 ONNX Runtime CUDA/TensorRT,不能把 OpenVinoArtifact(device="GPU") 当作 CUDA 后端。

热路径性能约束

推理框架遵循以下约束:

  1. 输入数组由策略长期持有,每帧只原地更新。
  2. 不在控制周期中做深拷贝,不重复创建 Session、InferRequest 或输入列表。
  3. ONNX Runtime 首次成功推理后建立 I/O Binding,后续复用调用方输入和输出缓冲。
  4. OpenVINO 持有一个 InferRequest,用共享内存 Tensor 绑定稳定、连续的 NumPy 输入, 并缓存输出视图。
  5. RKNN 复用按逻辑输入顺序排列的 Python 列表;SDK 返回值立即映射为具名输出。
  6. 性能监控默认关闭。关闭时 Policy.step() 不调用性能时钟;开启后才记录四段耗时。
  7. 模型加载、预热、哈希和转换都位于 steady-state 控制循环之外。

调用方替换了输入数组对象时,ONNX Runtime 和 OpenVINO 会重新绑定;仅修改原数组内容 不会触发重新绑定。输入必须保持 C contiguous,dtype 和 shape 应与模型一致。

历史帧

历史帧的存储机制统一使用 HistoryBuffer,但历史语义留在具体 InputBuilder 或策略中。 这样既避免每个策略重复实现环形缓冲,又不会把深度帧采样、动作历史、重置规则等业务 语义塞进公共框架。

HistoryBuffer 的约定:

  • 构造时一次性申请物理存储。
  • fill() 用当前观测填满历史,适合 reset/prepare。
  • append() 通过 np.copyto 写入下一槽,不申请新数组。
  • write_into() 按 oldest-to-newest 写入调用方提供的连续缓冲。
  • shape 和 dtype 不匹配时立即报错,不在热路径中隐式转换。

环形存储转为连续模型输入必然需要一次线性写入;目标数组由策略复用,因此不会产生 逐帧临时大数组。策略负责决定 advance=False 时是否保持历史不变。

性能监控

from bxi_example_py_elf3.inference import InferenceRuntime, RuntimeOptions

runtime = InferenceRuntime(
    options=RuntimeOptions(
        monitor_enabled=True,
        warmup_runs=2,
    )
)

监控按策略名称保存最近一段样本,阶段为 inputbackendoutputtotalsummary() 才会复制快照并计算 P50/P95/P99,统计工作不放在每帧控制路径中。性能监控 用于定位问题;正式部署如果不需要统计,应保持关闭。

RKNN 按需转换

rknn-toolkit-lite2 只负责设备推理;ONNX 转换由主机侧 rknn-toolkit2 完成。框架默认 绝不转换,也不会导入 Toolkit2 或计算模型哈希。只有显式设置环境变量才授权启动转换。

最简单的设置把环境变量值直接写成目标芯片:

BXI_RKNN_CONVERT_ON_LOAD=rk3588 ros2 run ...

复杂设置使用同一个环境变量中的 JSON;它只覆盖本次部署,类配置仍是默认值:

export BXI_RKNN_CONVERT_ON_LOAD='{
  "target": "rk3588",
  "do_quantization": true,
  "dataset": "/data/calibration.txt",
  "config": {"optimization_level": 3},
  "force_rebuild": false
}'

支持 target/target_platformdo_quantizationdatasetconfigforce_rebuild。环境变量未设置或值为 0falsenooff 时禁用。

转换缓存包含 ONNX 身份和 SHA256、目标芯片、Toolkit2 版本、量化设置、数据集身份和 构建参数。生成过程使用文件锁,拿锁后重新检查缓存,先导出同目录临时文件,再用 os.replace() 原子替换。因此多进程启动不会重复构建,转换失败也不会破坏已有 RKNN。 侧边文件为 <model>.rknn.build.json<model>.rknn.build.lock

x86_64 主机可以安装 Toolkit2 并生成 RKNN 缓存,但通常不能用 RKNN Lite 执行模型。 转换完成后,如果当前平台没有 Lite2 或不是目标 Rockchip SoC,运行时会继续回退到 OpenVINO/ONNX Runtime。在 Rockchip 上只有 Lite2、没有 Toolkit2 时,显式请求转换会 给出 Toolkit2 安装提示;已有且有效的 RKNN 文件仍可直接加载。

增加新后端

新后端只需要三个局部类型,不需要修改策略或 ModelSpec

  1. 继承 ModelArtifact,用 ClassVar backend 声明唯一名称。
  2. 实现 InferenceBackend.run(),并按需提供 shape、metadata、warmup 和 close。
  3. 实现 BackendFactory.availability()open(),注册到 BackendRegistry
from dataclasses import dataclass
from typing import ClassVar

from bxi_example_py_elf3.inference.model import ModelArtifact


@dataclass(frozen=True, slots=True)
class MyArtifact(ModelArtifact):
    backend: ClassVar[str] = "my_backend"
    device: str = "default"

后端专属选项属于自己的 Artifact,不应继续向 ModelSpec 增加固定字段。这保证未来加入 TensorRT、厂商 NPU 或远程推理时,核心模型描述仍保持稳定。

一键 benchmark

本地 benchmark 会自动发现 src/ 下所有 ONNX 模型、识别多输入和动态 batch,并枚举 可用的 ONNX Runtime provider、OpenVINO 设备与 RKNN。每个“模型 × 后端”在独立子进程 执行;即使 GPU 驱动触发 native abort,也只标记该组合失败,不会中止整批测试。

# 默认正式测试,并生成带硬件信息的 JSON 报告
python3 tools/benchmark/backend_benchmark.py

# 快速检查全部模型
python3 tools/benchmark/backend_benchmark.py --quick

# 更稳定的跨平台对比
python3 tools/benchmark/backend_benchmark.py --warmup 500 --iterations 10000

# 外部模型目录
python3 tools/benchmark/backend_benchmark.py /opt/models

输出包含冷启动 setup、P50/P95/P99、mean、吞吐率、输出缓冲稳定性以及与参考后端的 数值误差。安装加速 provider 后,ONNX Runtime CPU 保留为数值参考,CUDA/TensorRT 等 provider 单独列出。OpenVINO 非 Intel GPU 会明确跳过;AUTO[CPU] 等标签显示实际设备。

RKNN 转换 benchmark 仍需显式授权,生成物写入本地 benchmark cache,不污染模型资产:

BXI_RKNN_CONVERT_ON_LOAD=rk3588 \
python3 tools/benchmark/backend_benchmark.py --rknn-target rk3588

benchmark 工具位于纳入 Git 的 tools/benchmark/,可以随仓库部署到不同平台;只有 tools/benchmark/results/ 报告和 tools/benchmark/cache/ 转换缓存被忽略。跨机器比较 时应保存 JSON 报告,并保持相同的 warmup、iterations、电源模式和系统负载。

推荐阅读

Clone this wiki locally