Skip to content

perf: 降低 command_manager.compute 的 MotionCommand 开销 #1317

Description

@TATP-233

一句话问题

g1_motion_tracking 的 Manager-Based 热循环中,command_manager.compute 仍由 MotionCommand._update_metrics 的多次 NumPy 批量表达和重复临时数组主导,优化后约 MuJoCo 7.25 ms、MJWarp 6.71 ms/step。

为什么现在做

这是 #1316 当前最大的 manager 分区之一;#1292/#1314 已证明五个 reward/body-state copy 不是主要剩余杠杆。现有 src/unilab/tasks/motion_tracking/common/manager_terms.py 已有预分配输出和 Numba 依赖,可在 owner layer 内收敛,不需要新 backend contract。

最小交付结果

MotionCommand 正常 step 的 metrics/command 数值更新改为一个可预热、可复用的 Numba 批量实现,并移除被替代的 runtime NumPy 数学实现。CommandTerm 的 reset/env_ids、resampling、command shape/finite contract 保持不变;测试 oracle 可以独立使用 NumPy。

In scope

  • src/unilab/tasks/motion_tracking/common/ 增加必要的 Numba kernel 与冷路径初始化/输出 buffer。
  • 迁移 MotionCommand._update_metrics 的 anchor/body/joint error 计算,支持全 batch 与 partial-reset 行集合。
  • 保持 BoxMotionCommand 继承路径、指标名称/dtype/数值容差和 sampler 行为。
  • tests/tasks/ 增加 parity、output reuse、partial-reset 和 kernel compilation 测试。
  • 用同机 MuJoCo/MJWarp 8192-env workload 报告 command_manager.computeupdate_stateenv_step_total_ms;MJWarp 同时列 backend_physics_msbackend_host_cache_refresh_ms,后者包含 D2H/PCIe 等待。

Non-goals

Owner / 规模

Owner:MotionCommandcommon/kernels.py;预计 3–7 个文件、≤500 行净手写改动、一个 child PR,PR base 为 dev/issue-1316-update-state-numba

Acceptance criteria

  • 全 batch 和 partial-reset 的 metrics 与当前 contract 在明确 dtype/tolerance 内一致,输出 buffer 地址稳定复用。
  • Numba kernel 在构造/初始化冷路径完成编译,step 热路径无 JIT/asset 解析。
  • MuJoCo、MJWarp 都无稳定回退;结果表记录 transfer-inclusive env_step_total_ms/step_core_ms,并单列 D2H。
  • 最终 head 本地 make test-all 通过,PR body 记录实际命令和非 main base 的本地 gate 路由。

Stop conditions

  • 需要公共 manager/backend contract、第二套 runtime 数学实现,或超出 7 文件/500 LOC 预算。
  • partial-reset 或 BoxMotionCommand 语义无法保持,或任一目标 backend 出现稳定回退。

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