Skip to content

Inference Architecture

konodoki edited this page Jul 29, 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               # 后端注册、选择和回退
      monitor.py               # 可选耗时统计
      backends/                # ONNX Runtime / OpenVINO / RKNN
    joints/                    # Layout、显式默认值和命令解析器
    platform/                  # 机器人接入边界和关节 I/O 编码器
  policies/                    # ELF3 内置的具体策略
    joints.py                  # 模型关节契约与策略内部顺序
    normal.py
    amp.py
    beyondmimic.py
    depth.py

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

五种关节布局

机器人状态布局 State Layout
  -> 策略观测布局 Policy Observation Layout
  -> 策略动作布局 Policy Action Layout
  -> 完整机器人布局 Robot Layout
  -> 可选的硬件固定布局 Hardware Fixed Layout

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

  • State Layout 可以比策略更多。例如 31 DoF 机器人状态可以包含 29 个策略关节和两个夹爪。
  • Observation Layout 是策略真正读取的关节及顺序。
  • Action Layout 是策略真正输出的关节及顺序,可以与 observation 不同。
  • Robot Layout 是平台当前实际具有的完整关节集合,也是最终 MotorFrame 的布局。
  • Hardware Layout 只用于不携带关节名称、只能接收固定数组顺序的接口。

每周期的实际路径是:

完整 RobotObservation(N)
  -> PolicyJointContract 按名称选择 observation(M)
  -> 策略输出自己的 action MotorFrame(K)
  -> JointCommandResolver:K + 平台显式默认目标 -> 完整 Robot Frame(N)
  -> Transition:两端都先解析成 N,再插值
  -> 具名消息直接发布,或映射到 Hardware Layout

MKN 可以不同,只要 observation 名称存在于 Robot Layout、action 名称是 Robot Layout 的子集,且 N-K 的每个关节都有平台默认目标。

策略类中的 PolicyJointContract 是唯一权威:

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


class WalkPolicy:
    joint_contract = PolicyJointContract(
        observation=ELF3_POLICY_JOINTS,
        action=ELF3_POLICY_JOINTS,
    )

模型 metadata 只能校验该声明,不能在运行时反过来修改契约。声明和 metadata 不一致应在 加载时失败,不能猜测顺序。

调用策略

平台每周期更新一个长期复用的 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_POLICY_JOINTS


class AmpPolicy(JointPolicy):
    joint_contract = PolicyJointContract(
        observation=ELF3_POLICY_JOINTS,
        action=ELF3_POLICY_JOINTS,
    )

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

        count = self.joint_contract.action.dof_num
        self._default = np.asarray(default_position, dtype=np.float32).reshape(count)
        self._scale = np.asarray(action_scale, dtype=np.float32).reshape(count)
        self._kp = np.asarray(kp, dtype=np.float32).reshape(count)
        self._kd = np.asarray(kd, dtype=np.float32).reshape(count)

        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._default)
        self.publish_output(self._target, self._kp, self._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._default)

    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._scale, out=self._scaled)
        np.add(self._default, 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._default, 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()

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

29 策略与 31 关节机器人

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