Skip to content

perf: 降低 observation_manager.compute 的批量 pipeline 开销 #1319

Description

@TATP-233

一句话问题

observation_manager.computeg1_motion_tracking 中约 MuJoCo 4.37 ms、MJWarp 3.85 ms/step,其中 observation term 本身仅约 0.92/0.80 ms,主要开销在通用 copy、clip/scale、NaN 检查和 concat pipeline。

为什么现在做

#1314 已把 reward/body-state 读取优化到次要比例;当前 observation manager 是第三大 manager 分区。需要先在 owner layer 证明哪些批量操作值得 Numba,不能把任意 Python term 或 temporal semantics 强行编译。

最小交付结果

针对当前 g1_motion_tracking 配置中已证明的批量 observation pipeline,减少不必要的 full-batch 临时 copy/concat,并在确有数值循环的地方使用一个预热、复用 buffer 的 Numba 实现;被替代的 runtime NumPy 数学实现移除。任意 term、noise、delay、history 和 NaN policy 的公共 contract 保持原样。

In scope

  • src/unilab/managers/observation_manager.py 的正常 full-batch path 做 profiling-guided、可回退边界清晰的优化。
  • 只处理已确认无 temporal side effect 的 copy/clip/scale/concat 数值段;保留 term callable、noise RNG 顺序、delay/history 和 row-scoped reset 语义。
  • 必要时在 src/unilab/tasks/motion_tracking/common/ 放置 task-owned Numba output writer;不把任意 manager term 下沉到 backend。
  • 增加 observation shape/dtype、输出复用、NaN policy、noise/history/partial-reset parity 测试。
  • 同机 MuJoCo/MJWarp 报告 observation 与 update_state 收益;MJWarp 显式列 D2H/PCIe host-cache 时间,端到端统计包含它。

Non-goals

  • 不修改 command manager、reward/termination 数学、Hydra/config、runner/lifecycle、backend contract 或 Motrix。
  • 不改变随机数调用顺序、观测维度/顺序、NaN policy,也不保留被替代的第二套 runtime 数学路径。
  • 不新增通用 fused executor、编译器、常规 CI 或长期 benchmark 服务。

Owner / 规模

Owner:src/unilab/managers/observation_manager.py 及必要的 motion observation owner;预计 5–12 个文件、≤700 行净手写改动、一个 child PR,PR base 为 dev/issue-1316-update-state-numba

Acceptance criteria

  • actor/critic observation 的 shape、顺序、dtype、数值、noise/delay/history/NaN contract 与基线一致。
  • 热路径不引入额外 full-batch copy;任何 Numba kernel 在冷路径预热并复用输出。
  • MuJoCo/MJWarp 无稳定回退;报告 update_state_msenv_step_total_msstep_core_msbackend_physics_msbackend_host_cache_refresh_ms,明确 PCIe 已计入端到端而未重复计入 manager 阶段。
  • 最终 head 本地 make test-all 通过,PR body 记录本地 gate 和验证命令。

Stop conditions

  • 需要改变公共 observation contract、RNG/lifecycle 或引入第二个 execution path;或预计超过 12 文件/700 LOC。
  • 只有单一 backend 受益、任一 parity 失败,或优化无法区分于 measurement noise。

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