Skip to content

perf: 降低 command_manager.post_compute 的 relative transform 开销 #1318

Description

@TATP-233

一句话问题

command_manager.post_compute 在 motion tracking 中仍用多段 NumPy relative transform/body-frame 计算,优化后约 MuJoCo 4.23 ms、MJWarp 4.21 ms/step。

为什么现在做

#1314 的分区埋点显示第二次 body-state copy 已降到约 0.3 ms,剩余约 3.9 ms 是 relative transform、anchor frame position/orientation 和输出写入;这是当前第二大可控 manager 分区。

最小交付结果

将 MotionCommand 的正常 step 与 partial-reset relative/body-frame 变换收敛为预热、复用 buffer 的单一 Numba 数学实现,并删除生产路径中被替代的 NumPy 数学实现。MuJoCo/MJWarp 共用同一 task kernel;不引入 backend 特例。

In scope

  • src/unilab/tasks/motion_tracking/common/ 增加/扩展 fused relative-state kernel,覆盖现有 body_pos_relative_wbody_quat_relative_w、anchor pose 和 robot body-frame 输出。
  • 迁移 MotionCommand._refresh_relative_state 及 row-scoped reset 分支,保持输出 shape、dtype、写入顺序和 BoxMotionCommand 的 post hook。
  • 删除不再被 production call graph 使用的对应 NumPy helper;测试中的独立 oracle 可保留。
  • 增加随机 quaternion、body subset、partial-reset、输出复用和 MuJoCo/MJWarp parity 测试。
  • 报告 command_manager.post_compute 与端到端阶段;MJWarp 必须同时报告 backend_host_cache_refresh_ms(D2H/PCIe)而不重复加到 update_state_ms

Non-goals

  • 不改 command_manager.compute、observation pipeline、reward/termination 公式、backend API、runner/lifecycle 或 Motrix。
  • 不维护第二个 runtime 数学表达,不新增通用编译器或 device-resident observations。

Owner / 规模

Owner:motion_tracking/common/transforms.pyobservations.pymanager_terms.py 与共享 kernel;预计 5–10 个文件、≤650 行净手写改动、一个 child PR,PR base 为 dev/issue-1316-update-state-numba

Acceptance criteria

  • 正常 step 与 partial reset 的所有 relative outputs 在明确 tolerance 内 parity,BoxMotionCommand 仍通过现有 contract。
  • 冷路径完成 Numba 编译;热路径不分配临时 full-batch arrays、不解析 asset。
  • MuJoCo/MJWarp 均无稳定回退;报告 transfer-inclusive env_step_total_ms/step_core_ms 和独立 D2H 项。
  • 最终 head 本地 make test-all 通过,PR body 记录本地 gate。

Stop conditions

  • 需要新增公共 contract/执行路径、必须保留双 runtime 数学实现,或超过 10 文件/650 LOC。
  • 任一 backend、partial reset 或 BoxMotion 数值契约出现稳定回退。

Parent roadmap:#1316

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:benchmarkBenchmark recording and evaluation workflowenhancementNew feature or request

    Type

    No type

    Projects

    No projects

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions