feat(mujoco_uni): BatchEnvPool 单次 dispatch per-substep control callback 能力 (#1259 M1a)

[TATP-233]

Parent: #1259 (Wave 2, M1a). Capability decision: #1262.
代码落在 upstream unilabsim/mujoco_uni(该仓库已关闭 Issues,故 issue 挂在本仓库,PR 在 mujoco_uni)。M1b adapter 接入另行跟踪。

一句话目标

BatchEnvPool.step 新增可选 control_callback,让调用方在一次 pool dispatch 内完成 N 个 physics substep,同时在每个 substep 前用最新 intermediate state/sensordata 经 Python 回调计算 control,替代 UniLab 目前每 substep 一次 step(nstep=1) 的多 dispatch 循环。

背景

• Manager-Based API 要求 community-compatible per-substep action;MuJoCo backend 当前为 sim_substeps 个 substep 各做一次 pool dispatch,固定成本(pybind 边界、参数校验、state0 scatter、warmstart/warning 重置、state+sensor gather、numpy 分配)造成 排查：mujoco physics_ms 在 MBA 分支系统性偏高（+23-43%）归因 #1257 测得的 23%–43% physics 差值。
• control 依赖上一 substep 后的状态(action term 经 getter 读 joint pos/vel 等),无法预计算成 (nbatch, nstep, ncontrol) 轨迹,因此需要回调而非现有 control trajectory 参数。

API 设计

pool.step(
    initial_state,
    nstep=N,
    control_spec=...,
    control_callback=fn,   # 与 control 互斥
    return_sensor=True,
    post_step_forward_sensor=False,
    chunk_size=...,
)
# fn(step_index: int, state: (nbatch, nstate) f64, sensordata: (nbatch, nsensordata) f64)
#   -> control: (nbatch, ncontrol) f64 C-contiguous

• callback 在每个 substep 之前调用,共 nstep 次;t=0 收到 scatter 后的 state0 与 mjData 当前 sensordata。
• 逐 substep 语义与现状逐调用 step(nstep=1) 完全一致:每 substep 前 zero qacc_warmstart、warning 清零不中断、post_step_forward_sensor 逐 substep 生效。
• GIL:worker dispatch/wait 期间释放,调 callback 时持有;callback 抛异常直接传播,该次 rollout 中止。
• 错误处理:与 control 同时给出、callback 返回 shape/dtype/layout 不符均 fail-fast ValueError。

Acceptance criteria

• 状态依赖 callback(如 ctrl=-k*qvel)与手工 step(nstep=1) 循环逐 substep 数值一致(threaded 与单线程同口径)。
• 错误用例与异常传播有测试。
• mujoco_uni 现有 tests/ 全部通过;不影响既有 step/forward/reset 行为。
• 发布后版本固定为 mujoco-uni-runtime==0.4.0(tag 触发 release.yml 发 sdist)。

Non-goals

• 不改既有 control trajectory 路径语义;不引入 UniLab 侧改动(接入属 M1b child)。
• 不做 chunk tuner、sensor 裁剪等其他优化。

[TATP-233]

M1a 完成:

• 实现:feat: BatchEnvPool per-substep control_callback (UniLab #1259 M1a) unilabsim/mujoco_uni#1 已 rebase 合入 main(e73e89a + c96290d),BatchEnvPool.step 新增 control_callback / callback_sensordata。
• 发布:release 矩阵首轮暴露 0.3.1 引入的 cpu_affinity 测试在非 Linux 收集期失败的存量 bug,经 fix: gate cpu_affinity test collection to Linux unilabsim/mujoco_uni#2 修复后,v0.4.0 tag 重触发发布成功,mujoco-uni-runtime==0.4.0 sdist 已在 PyPI。
• 验证:mujoco_uni tests/ 37 passed;与逐调用 step(nstep=1) 的数值 parity atol=1e-12, rtol=0(threaded/单线程同口径)。
• 收益边界(实测记录于 PR body):pool 级 callback ≈ per-call;收益主要来自跳过逐 substep sensor gather + UniLab 侧 astype/copy;dispatch 同步固定成本保留,残余归 M2/M3 与 U 线。barrier 驻留 worker 方案经 maintainer 决策不采用。

M1b 接入在 #1267 进行。