perf: 定位 Env overhead 并行化实测天花板 — Numba prange 融合 kernel 在 Xeon 8568Y+ 到 31x

[TATP-233]

Context

Follow-up to #663。跑了一组独立于 UniLab 的 standalone probe（benchmark/benchmark_env_overhead_parallel.py, 812 行，无 UniLab 依赖）来回答：把 NpEnv.step 里 obs/reward/termination 那段 update_state 并行化，实测天花板在哪。

Probe 设置对齐真实训练：num_envs=8192/32768, act_dim=29, obs_dim=98/101, dtype=float32，15 个 active reward terms（对齐 motion-tracking scale），4 个 noise buffer，per-reward log mean，2% _reset_done_envs 拷贝。

三条路径对比：numpy 基线 / ThreadPoolExecutor + numpy shards / Numba prange 融合 kernel。

结论

1. Numba prange 融合 kernel 在三台机器上都拿到显著加速

Xeon 8568Y+ (160T, 2-socket)

配置	numpy 基线	Numba 最佳	Speedup
8k envs, 双 socket (tbb)	4.55 ms	0.312 ms @ 32t	14.6x
8k envs, `numactl -N 0`	4.55 ms	0.260 ms @ 24t	18.0x
32k envs, 双 socket	19.4 ms	0.671 ms @ 48t	28.8x
32k envs, `numactl -N 0`	19.4 ms	0.664 ms @ 32t	31.4x

AMD Ryzen 9950X3D (16C/32T)

配置	numpy 基线	Numba 最佳	Speedup
8k envs (omp)	2.24 ms	0.139 ms @ 32t	16.14x
32k envs (omp)	8.84 ms	0.334 ms @ 16t	26.45x

Apple M-series (18C, workqueue)

配置	numpy 基线	Numba 最佳	Speedup
8k envs	1.78 ms	0.255 ms @ 16t	6.96x
32k envs	9.40 ms	0.567 ms @ 16t	16.56x

单核 fusion 贡献 2.3-3.8x（消掉每个 reward term 的中间 numpy array），并行再拿 ~5-10x。大 batch 一致地给出更高 speedup——每个 parallel section 摊到更多 env，barrier / launch cost 摊平。

2. False sharing 是并行的真正天花板

kernel 里 15 个 reward log 累加不能写成 shared scalar 的 `log_partial_sum[k] += ...`——所有 thread `+=` 同一 cache line 会造成持续 ping-pong，Xeon 上 `nthreads=32` 只到 5.5x。必须用 per-thread scratch `log_thread_sum[tid, k]`，kernel 结束后主线程 `axis=0` sum。同一份代码这个改动之后到 14.6x（8k）/ 31.4x（32k）。

3. update_state 只是 Env overhead 的一半，reset_done 是并列瓶颈

来自 #663 中 b44ae656 之后的完整细分（Xeon 8568Y+ / 8192 envs）：

Task	Backend	update_state	reset_done	主瓶颈
g1_motion_tracking	mujoco	35.1 ms	12.6 ms	update_state 主导
g1_motion_tracking	motrix	40.7 ms	16.4 ms	update_state 主导
g1_walk_flat	mujoco	8.4 ms	8.6 ms	两者相当
g1_walk_flat	motrix	10.3 ms	12.4 ms	reset_done 略大

`reset_done` 的开销分布：

• Terminal-observation copy / scatter：sub-ms，可忽略
• `reset call` 8-15.8 ms，主要在：
  • `g1_walk_flat`：几乎全在 backend `set_state`（7.2-11.7 ms）
  • `g1_motion_tracking`：reset plan + backend `set_state` + reset observation（其中 full-batch body-pose getter ~3.2-3.7 ms）

只上 Numba 优化 update_state 时：`g1_motion_tracking` 的 overhead 削减 71%（40.7→1.3ms + 16.4→16.4ms），`g1_walk_flat` 只削减 54%（10.3→0.3 + 12.4→12.4ms）。walk_flat 的 reset_done 会立即成为新瓶颈。

4. ThreadPool + numpy shards 是死路

Xeon 上 K=2 只 1.7x，K=4 已经打平，K≥8 全崩（GIL 争用 + Python log dict + 双路 NUMA）。Ryzen 稍好（K=4 @ 32k 到 2.95x），但仍远不及 Numba，也无法组合使用。

5. Distribute / aggregate 可忽略

`pool.submit` ~10 μs，aggregate（log 归并 + 2% final_observation 拷贝）~30-50 μs，总占比 <2%。不用为分片调度做优化。

6. RNG 是独立的天花板

Noise 生成（`rng.standard_normal` × 4）单线程占 5-6 ms，在 update_state 里占 55-80%（Xeon 上 80%）。Numba prange 无法帮它并行，需要独立方案（in-kernel numba RNG / 多 BitGenerator / 复用 buffer）。

推荐路径

主线 A（P0）：Numba prange 融合 kernel — `update_state` 层，跨 task 通用

对 `g1_motion_tracking` 直接把 overhead 打掉 71%，对 `g1_walk_flat` 也能砍掉 update_state 全部时间。关键实现要点：

• 融合 reward dispatch + obs concat + termination 到一个 `@njit(parallel=True)` prange loop
• reward log 累加必须 per-thread scratch，`(nthreads, num_terms)` shape，最后 `axis=0` sum
• noise buffer 从 kernel 外传入（RNG 独立处理）
• kernel 内部通过 `numba.set_num_threads(24-32)` 控制并发，不要全进程 `numactl -N 0`（那会伤 Physics，Xeon 8568Y+ Physics 已经吃到 22-66ms/步）

主线 B（P1）：reset_done 拆解优化 — 按 task 分头做，与 A 并行推进

• `g1_motion_tracking`：row-wise body-pose getter，只对 done 的 env 拉数据（3.2-3.7 ms → sub-ms）
• `g1_walk_flat`：sparse `backend.set_state` — 需要 backend 层增加 index-based 写接口（7-12 ms → 期望 sub-ms）

P2（可分阶段叠加）

• RNG 并行化（独立 ~5ms 收益，对所有 task 有效）
• Async pipeline（`step t+1 physics ‖ step t update_state`，乘性 ~2x）

不做

• `ThreadPoolExecutor + numpy shards`（Xeon 无效，Mac/Ryzen 上限低，无法组合 Numba）
• 对 distribute/aggregate 做优化（<2% 占比）
• 全进程 `numactl -N 0` pin（伤 Physics）

预期收益（Xeon 8568Y+ / 8192 envs / motrix）

阶段	动作	motion_tracking overhead	walk_flat overhead	motion_tracking throughput
现状	—	60.3 ms	27.9 ms	84k env/s
P0	Numba prange `update_state`	17.7 ms (-71%)	12.7 ms (-54%)	~145k env/s (+72%)
P0+P1a	motion_tracking row-wise body-pose	~15 ms	—	~155k env/s
P0+P1b	walk_flat sparse `set_state`	—	~5-7 ms	walk_flat +20-30%

CPU% 预期从当前 13-77% 提到 60-80% 区间（依 task 和 backend 不同）。P0 完成后 Physics 可能成为新瓶颈，届时再评估。

探针复现

• 代码：benchmark/benchmark_env_overhead_parallel.py（branch: `bench/env-overhead-parallel-probe`，812 行，无 UniLab 依赖）
• 运行：`uv run --with numpy --with numba python benchmark/benchmark_env_overhead_parallel.py`
• 大 batch：`PROBE_NUM_ENVS=32768 uv run --with numpy --with numba python benchmark/benchmark_env_overhead_parallel.py`
• Xeon 单 socket 测试：`PROBE_NUM_ENVS=32768 numactl -N 0 --membind 0 uv run --with numpy --with numba python benchmark/benchmark_env_overhead_parallel.py`

Refs: #663

[TATP-233]

审阅意见（按风险/优先级排序）：

1. P0「跨 task 通用 Numba kernel」定义过宽。 现有 update_state 不是统一纯函数：g1_motion_tracking 会读 motion sampler、backend body state、更新 sampler、写 info 并跑 reward dispatch；g1_walk_flat 又有自己的 termination / curriculum / reward / obs 结构。Numba 不能直接覆盖这些 Python dict、reward function dispatch、backend getter。建议改成「task-owned typed snapshot + task-specific fused kernel」，不要放在 NpEnv 层承诺通用内核。

2. g1_walk_flat sparse set_state 的表述需要修正。 SimBackend.set_state(env_indices, qpos, qvel, ...) 已经是 index-based contract，MuJoCo / Motrix 也都按 env ids/reset slice 走。这里更准确的下一步是细分并优化现有 set_state 实现，例如 mask 构造、reset、randomization、set_dof、forward_kinematic 各自成本，而不是新增 index-based 写接口。

3. Numba 进入训练热路径前需要明确依赖策略。 当前项目运行依赖没有 numba，仓库里只是 benchmark 脚本可选 import。若 P0 作为主线，需要先定义 optional extra / fallback、Python 3.10-3.13 支持、Apple/Linux runtime 行为、线程数配置方式。

4. RNG 结论需要收窄。 当前 env 噪声路径使用的是 np.random.uniform，不是 standard_normal；而 flashsac/g1_walk_flat offpolicy 配置里 noise_config.level: 0.0，因此不能把该 task 的 update_state 热点泛化归因为 RNG。RNG 并行化可以作为 noisy task 专项，但「对所有 task 有效」需要补仓库内分项证据或收窄表述。

5. g1_motion_tracking row-wise body pose 是低风险高确定性的优先项。 reset observation 目前先取 full batch body pose 再按 env_ids 切片；后端 contract 已有 get_body_pose_w_rows，MuJoCo / Motrix 都实现了。这个可以先做，不必等待 Numba 主线。

6. 外部 probe 需要入库或固化 artifact。 当前核心结论来自 /tmp/probe_numba_env_overhead.py，后续不可复现、不可 CI 化。建议放到 benchmark/，输出 JSON/CSV，并记录命令、依赖和环境。

建议执行顺序：

1. 先做 g1_motion_tracking reset observation 的 row-wise body pose 读取。
2. 给 backend set_state 增加内部细分 timing，再决定 g1_walk_flat reset 优化方向。
3. 再做 task-specific Numba fused kernel prototype，而不是一开始承诺跨 task 通用 kernel。

[TATP-233]

相关现象：

1. perf: improve off-policy collector active throughput #651
2. Work: 160核服务器性能调优总结 #659
3. 并行加速 Env overhead（obs/reward/state-update）以提升 collector CPU 利用率与吞吐 #663

[TATP-233]

结论补充：如果只参考当前 benchmark/numba_profile 的形态，做 task-specific Numba fused kernel、不动整体架构，收益是确定但有上限的。

• 严格按当前 profile，只覆盖 g1_motion_tracking 的 reward + termination micro-slice：端到端通常不会是数量级提升，更可能是约 10%-25% 级别。
• 如果仍不改架构，但把 task 内 update_state 的纯数组段尽量融合（reward / termination / 部分 transform/obs/log scratch）：g1_motion_tracking 可接近 perf: 定位 Env overhead 并行化实测天花板 — Numba prange 融合 kernel 在 Xeon 8568Y+ 到 31x #665 投影的 ~1.7x throughput（约 +70%）。
• 对 g1_walk_flat 这类 reset 占比高的 task，Numba 打掉 update_state 后，reset_done/backend set_state 会马上成为瓶颈，整体更像 1.2x-1.5x。

所以：task-specific Numba 是 P0，值得做；但它主要解决的是“已经 materialize 成 ndarray 后的纯数组计算”。极致性能目标下，后续大头收益必须来自架构/contract 层：

1. backend/reset path：细分并优化 set_state、DR payload、forward/kinematics、reset observation；
2. backend state snapshot：减少多次 getter、Python/backend 边界和拷贝；
3. runner async pipeline：physics step(t+1) 与 update_state(t) 并行；
4. 如需跨 task 通用 fused kernel，则需要 reward/obs 的 typed registry、IR 或 codegen，而不是 NpEnv 层直接通用 Numba 化。

建议表述收窄为：先做 g1_motion_tracking task-owned typed snapshot + task-specific fused kernel，拿确定收益；同时给 reset/backend getter 加 timing，下一阶段直接攻 backend snapshot/reset 和 async pipeline。

[TATP-233]

补充远端实测结论（unilab@134.199.200.57:/home/unilab/unilabsim/UniLab，perf/g1-walk-numba-accel / bc34074f，motrixsim，collector A/B 不含 learner）：

• 验证：相关 benchmark/env tests 36 passed。
• 专项 hot-slice（只含 reward dispatch + termination）确实有明显收益：
  • G1Joystick / walk flat：3.80x - 6.16x。
  • G1MotionTracking：5.36x - 11.04x。
• 但真实 collector active-window 端到端收益明显更小，使用各 env 下专项最优线程后：
  • G1Joystick / walk flat：1024/2048/4096 envs 分别 1.13x / 1.01x / 1.04x。
  • G1MotionTracking：1024/2048/4096 envs 分别 1.08x / 1.11x / 1.04x。
• collector 里 update_state 的收益能看到，但不足以主导总耗时：
  • walk flat：update_state 为 1.31x / 1.57x / 1.57x，节省 1.161 / 2.214 / 2.960 ms。
  • motion tracking：update_state 为 1.04x / 1.12x / 1.03x，节省 0.309 / 1.293 / 0.640 ms。
• 关键原因：当前 Numba 覆盖的是局部 reward+termination slice；真实 collector 还包含 backend step / step_core、reset、action selection、replay write、bookkeeping 等未加速部分。motion tracking 中 _update_relative_transforms 也不在 reward+termination kernel 内；把它加回 synthetic 口径后，总加速只有 2.17x / 2.14x / 2.56x。

结论：task-specific Numba 是正确且可测的局部优化，但按这组远端实测，不处理 backend/getter、reset、observation/state assembly、actor/replay 等真实 collector 瓶颈时，collector 可见收益上限目前约 1.13x。如果目标是极致性能，后续必须扩大优化面，不能只停留在 reward+termination 的 task-specific Numba。

本地报告：agent/remote_numba_benchmark_report_2026-07-07.md。

[TATP-233]

补充：按远端 best-thread collector A/B 的 numpy baseline 拆分，如果暂不考虑物理引擎本身 step 时间，下一步最大优化上限仍在 update_state，不是 reward+termination 已覆盖的那一小段。

口径：bc34074f，motrixsim，collector active-window，不含 learner；比例用 baseline 计算；non-physics = total collector active ms - physics ms。

Task	envs	total ms	physics ms	non-physics ms	update_state ms	update / total	update / non-physics
G1MotionTracking	1024	24.243	5.502	18.741	8.062	33.3%	43.0%
G1MotionTracking	2048	36.293	9.651	26.642	12.301	33.9%	46.2%
G1MotionTracking	4096	55.477	16.351	39.125	21.390	38.6%	54.7%
G1Joystick / walk flat	1024	17.008	5.559	11.449	4.860	28.6%	42.4%
G1Joystick / walk flat	2048	27.750	12.097	15.653	6.124	22.1%	39.1%
G1Joystick / walk flat	4096	36.864	17.347	19.517	8.118	22.0%	41.6%

对照其他非 physics 项：motion tracking 的 reset 约占 total 12.4%-20.0%，walk flat 的 action selection 约占 total 15.5%-17.5%；它们都值得后续看，但单项上限低于 update_state。

原因：当前 Numba 主要覆盖 reward dispatch + termination；但真实 update_state 还包含 backend/state 读取、state assembly，以及 motion tracking 的 _update_relative_transforms 等未覆盖路径。motion tracking synthetic 中把 relative transform 加回后，总加速已从 reward+termination 的 5.36x-11.04x 降到 2.14x-2.56x，说明瓶颈已经从 reward+termination 外溢到更完整的 update path。

结论：如果继续实验，优先做 update_state 全路径剖分和融合实验，目标是 state/backend getter、relative transform、observation/state assembly，而不是继续只扩 reward+termination kernel。按现有数据，这一方向覆盖总体 collector 的 22.0%-38.6%，覆盖非 physics 时间的 39.1%-54.7%，是排除 physics 后最大优化上限。

[TATP-233]

G1MotionTracking 合并版结论（commit b96224a6，远端 unilab@134.199.200.57:/home/unilab/unilabsim/UniLab）：

• 实现范围：G1MotionTracking numba 路径覆盖 full update_state array path，包括相对坐标/anchor transform、reward、termination、actor/critic obs assembly；复用的 numba 几何/坐标系 helper 已抽到 src/unilab/utils/numba_geometry.py。
• 验证：uv run pytest tests/envs/test_g1_motion_tracking_numba.py tests/benchmark/test_g1_motion_tracking_numba_benchmark.py tests/utils/test_numba_geometry.py tests/utils/test_utils_package_policy.py -q，远端结果 19 passed in 29.89s。
• Benchmark：uv run python -m benchmark.benchmark_g1_motion_tracking_numba，覆盖 profiles ppo_default/sac_default/full_supported，envs 512..32768，threads 1,2,4,8（脚本对 motion tracking 当前默认实测到 8T），iters 80，warmup 8；默认包含 e2e collector。
• parity：reward/update_state reward max abs diff 1.27e-07，actor obs max abs diff 1.79e-07，critic obs max abs diff 7.15e-07，termination mismatch 0。

Synthetic full update_state array path 关键结果：

profile        envs   threads  update_state ms      speedup  reward+term ms
ppo_default    512    8        1.207 -> 0.212      5.69x    0.390 -> 0.094
ppo_default    32768  8        83.135 -> 7.089     11.73x   25.868 -> 2.545
sac_default    512    8        1.165 -> 0.164      7.09x    0.395 -> 0.101
sac_default    32768  8        85.725 -> 6.583     13.02x   25.212 -> 2.579
full_supported 512    8        1.213 -> 0.165      7.37x    0.462 -> 0.101
full_supported 32768  8        86.455 -> 6.799     12.72x   28.517 -> 2.876

E2E collector A/B（case sac/g1_motion_tracking/motrixsim，collector only，不含 learner）：

envs   threads  collector speedup  collector step ms       physics ms          update_state ms       update_state speedup
1024   8        1.21x              22.663 -> 18.676       6.263 -> 5.671     7.842 -> 5.723      1.37x
2048   8        1.04x              32.801 -> 31.660       10.049 -> 9.936    12.646 -> 7.363     1.72x
4096   8        1.26x              49.718 -> 39.370       14.966 -> 14.753   21.600 -> 10.987    1.97x
8192   8        1.22x              91.510 -> 75.195       28.497 -> 29.253   41.155 -> 19.913    2.07x
16384  8        1.35x              183.774 -> 135.996     61.474 -> 55.740   85.956 -> 37.494    2.29x
32768  8        1.42x              367.199 -> 258.044     115.981 -> 111.623 174.844 -> 68.883   2.54x

结论：G1MotionTracking 的 full update_state synthetic hot slice 在大 env 下为确定高收益（sac_default 32768 env 13.02x），e2e collector 总体为 1.04x-1.42x，collector 内 update_state 为 1.37x-2.54x。以 32768 env 为例，baseline collector step 中 physics 占 31.6%，update_state 占 47.6%，other 占 20.8%；优化后 update_state 从 174.844 ms 降到 68.883 ms，但仍是非 physics 最大项。因此若不考虑物理 step，本任务下一阶段最大空间仍在继续拆解/优化 update_state 内剩余路径，并同步关注 collector bookkeeping/other 的波动。

[TATP-233]

G1Joystick 合并版结论（commit b96224a6，远端 unilab@134.199.200.57:/home/unilab/unilabsim/UniLab）：

• 实现范围：G1WalkEnv.update_state() 在 numba_acceleration=true 时走 joystick_numba.py 的 full update_state hot slice，一次性生成 reward、termination、actor/critic obs；不再额外跑 numpy _compute_obs。
• 不静默回退：numba 缺失、active unsupported reward term、或 noise_config.level > 0 均显式报错。
• 验证：uv run pytest tests/envs/locomotion/g1/test_joystick_numba.py tests/benchmark/test_g1_joystick_numba_benchmark.py tests/envs/test_g1_motion_tracking_numba.py tests/utils/test_numba_geometry.py tests/utils/test_utils_package_policy.py -q，远端结果 25 passed in 61.33s。

Synthetic full update_state hot slice：

• 完整 all-profile run 覆盖 profiles ppo_default/sac_default/full_supported，envs 512..32768，threads 1..64，iters 80，warmup 8。
• parity：reward max abs diff 5.96e-08，actor/critic obs max abs diff 3.73e-09，termination mismatch 0。
• 大 env 关键结果：sac_default 32768 env 22.061 ms -> 1.277 ms（64T，17.27x）；ppo_default 32768 env 22.042 ms -> 1.380 ms（64T，15.98x）；full_supported 32768 env 38.374 ms -> 2.411 ms（64T，15.91x）。

E2E collector A/B（uv run python -m benchmark.benchmark_g1_joystick_numba --profiles sac_default --e2e，case sac/g1_walk_flat/motrixsim，collector only，不含 learner）：

envs   threads  collector speedup  collector step ms       physics ms          update_state ms      update_state speedup
1024   8        1.06x              15.865 -> 14.998       6.857 -> 6.907     4.900 -> 3.201     1.53x
2048   16       1.08x              23.227 -> 21.538       10.315 -> 10.471   6.154 -> 3.364     1.83x
4096   16       1.13x              34.119 -> 30.291       18.229 -> 19.221   8.337 -> 3.624     2.30x
8192   16       1.23x              67.979 -> 55.356       35.889 -> 34.335   11.740 -> 4.161    2.82x
16384  64       1.22x              127.600 -> 104.740     72.713 -> 67.753   20.911 -> 5.808    3.60x
32768  64       1.21x              239.446 -> 197.546     131.389 -> 130.204 46.942 -> 6.248    7.51x

结论：G1Joystick full update_state synthetic hot slice 在大 env 数下有明确高收益（sac_default 32768 env 约 17x），e2e collector 也吃到收益但被 physics 和 other 稀释，总体为 1.06x-1.23x；collector 内 update_state 为 1.53x-7.51x。以 32768 env 为例，baseline collector step 中 physics 占 54.9%，update_state 占 19.6%，other 占 25.5%；numba 后 update_state 降到 6.248 ms，下一步若不考虑 physics，本任务的剩余主要空间转向 other/collector bookkeeping，而不是继续只盯 reward/obs kernel。