G1 motion tracking update_state 热点优化 A/B 复测：实际收益、归因与剩余瓶颈

[TATP-233]

关联 #1292、roadmap #1304 和 integration PR #1314。本文按 #1292 的思路重新做 update_state A/B：阶段耗时用无埋点运行，热点 term 用完整 callable 埋点运行，并单独核对真正的获益来源。

范围声明：本轮只覆盖 MuJoCo 和 MJWarp；Motrix 已按 roadmap 的 scope amendment 排除，因此本文不对 Motrix 作任何性能声明。

结论

在 SAC g1_motion_tracking、8192 envs 下，相对优化前 commit：

• MuJoCo update_state 从 29.147 ms 降至 19.230 ms，节省 9.917 ms / 34.0%，即 1.52×；完整 env.step 下降 10.7%。
• MJWarp update_state 从 28.150 ms 降至 17.801 ms，节省 10.349 ms / 36.8%，即 1.58×；完整 env.step 下降 18.1%。
• 真正的两个获益点是：
  1. 四个 fixed body rewards 改为 Numba kernel：MuJoCo / MJWarp 每步分别节省约 4.23 / 3.92 ms；
  2. 四组 full-batch getter/gather 收敛为一次 backend-owned fused body-state copy，并且 update_state 中两次 refresh 都受益：分别节省约 6.07 / 6.62 ms。
• termination/anchor_pos 的完整 wrapper 看起来有 8–10× 加速，但不能把这个数字归因于 termination Numba 算术。它是本轮执行顺序中第一个读取 robot body state 的 term，因而把第一次 full-batch refresh 记在了自己名下；扣掉 refresh 后没有观察到独立、显著的 termination 算术收益。
• 五个目标 term 优化后合计仅占当前 update_state 约 6–7%，继续微调它们已经不是最优杠杆。剩余主要瓶颈转移到 command transform、observation 后处理，以及 MuJoCo 的 physics。

本文的精细复测与 PR #1314 中 collector benchmark 的统计口径不同：这里报告 NpEnv.step 内置阶段计时的“三轮 step median 再取 median”；PR body 报告 benchmark 脚本的 run mean 与 env/s。因此两组数值不应逐项混用。

阶段与端到端 A/B

Backend	update_state 前 → 后	update_state 收益	env.step 前 → 后	env.step 收益	reset_done 前 → 后
MuJoCo	29.147 → 19.230 ms	-9.917 ms / -34.0% / 1.52×	93.857 → 83.815 ms	-10.042 ms / -10.7%	8.257 → 8.323 ms
MJWarp	28.150 → 17.801 ms	-10.349 ms / -36.8% / 1.58×	60.877 → 49.880 ms	-10.997 ms / -18.1%	14.974 → 14.695 ms

reset_done 基本不变，符合本轮修改集中在正常 step 热路径的预期。上表为无 Python term wrapper 的独立运行，避免 profiler 自身改变 canonical 阶段结果。

三次无埋点独立运行的原始 per-run median

Backend / metric	优化前三轮 (ms)	优化后三轮 (ms)
MuJoCo update_state	29.377074 / 28.972903 / 29.146988	19.366008 / 19.152474 / 19.229526
MuJoCo env.step	93.857309 / 92.385072 / 94.424855	83.310558 / 85.062404 / 83.814829
MuJoCo reset_done	8.256776 / 8.244145 / 8.310594	8.380513 / 8.292273 / 8.322613
MJWarp update_state	28.075597 / 28.149614 / 28.735984	17.843535 / 17.800564 / 17.717718
MJWarp env.step	60.877409 / 60.951091 / 60.146760	49.954170 / 49.880364 / 47.911239
MJWarp reset_done	14.974384 / 14.973129 / 15.038048	14.694527 / 14.690537 / 14.722034

完整热点 term A/B

这里的“完整 term”是一次完整的 term_cfg.func(...) 调用，不只计 Numba kernel：它会包含该调用触发的 getter/cache refresh、body subset copy 和 kernel；不包含 manager 通用的 shape/finite 检查、weight 与聚合。每轮取 100 次调用的 mean，再取三轮 median。

Backend / 完整 term	优化前 → 优化后 (ms)	表观加速
MuJoCo termination/anchor_pos	4.746 → 0.572	8.30×
MuJoCo reward/motion_body_pos	1.246 → 0.105	11.91×
MuJoCo reward/motion_body_ori	1.083 → 0.351	3.08×
MuJoCo reward/motion_body_lin_vel	1.259 → 0.088	14.25×
MuJoCo reward/motion_body_ang_vel	1.256 → 0.087	14.41×
MJWarp termination/anchor_pos	5.253 → 0.534	9.84×
MJWarp reward/motion_body_pos	1.164 → 0.092	12.59×
MJWarp reward/motion_body_ori	1.058 → 0.387	2.73×
MJWarp reward/motion_body_lin_vel	1.172 → 0.083	14.05×
MJWarp reward/motion_body_ang_vel	1.173 → 0.080	14.72×

聚合结果：

Backend	四个 body rewards 前 → 后	四 rewards 收益	五个目标完整 term 前 → 后
MuJoCo	4.862 → 0.632 ms	-4.231 ms / 7.70×	9.608 → 1.205 ms / 7.97×
MJWarp	4.561 → 0.643 ms	-3.918 ms / 7.09×	9.845 → 1.184 ms / 8.31×

收益归因：为什么 anchor_pos 的 8–10× 不是 termination kernel 的 8–10×

termination/anchor_pos 恰好是第一个访问 motion command robot body cache 的 term。完整 wrapper 的计时因此包含第一次 full-batch body-state refresh。用同条件的 100-step diagnostic run 把 refresh 再分成 first_read 与 force 后得到：

Backend	首次 refresh 前 → 后	anchor_pos wrapper 扣除首次 refresh 后
MuJoCo	4.594 → 0.481 ms	0.052 → 0.082 ms
MJWarp	5.235 → 0.450 ms	0.031 → 0.077 ms

扣除值只是两个嵌套 perf_counter_ns wrapper 的差值，不适合作为微基准；它只用于确认量级。证据说明 anchor_pos 的绝大部分表观下降来自 fused refresh，不能据此声称 termination 算术本身 8–10×。若要单独判断这个极短 kernel，需要专门的 kernel microbenchmark；它不会改变本帖的完整 term 与阶段结论。

每个正常 update_state 实际发生两次有效 refresh：首次读取一次，以及 command.post_compute(force=True) 一次。三轮埋点结果按每步折算为：

Backend	两次 refresh 合计前 → 后	每步节省
MuJoCo	6.886 → 0.813 ms	6.073 ms
MJWarp	7.388 → 0.765 ms	6.624 ms

因此，四个 rewards 的约 4 ms 与 fused refresh 的约 6–7 ms，基本解释了 update_state 的约 10 ms 端到端下降。注意首次 refresh 已经包含在 anchor_pos 完整 wrapper 中，不能再把“五个目标 term 节省”和“两次 refresh 节省”直接相加，否则会重复计数。

当前剩余瓶颈

以下为优化后完整 manager callable 的 coarse instrumentation；占比以无埋点 update_state median 为分母，因此只应按量级读，不应把各 wrapper 小数位当成无扰动真值。

当前 update_state 分区	MuJoCo	MJWarp
command_manager.compute	7.25 ms / 37.7%	6.71 ms / 37.7%
command_manager.post_compute	4.23 ms / 22.0%	4.21 ms / 23.6%
observation_manager.compute	4.37 ms / 22.7%	3.85 ms / 21.6%
所有 reward + termination	2.89 ms / 15.0%	2.60 ms / 14.6%
未覆盖 residual	约 0.53 ms / 2.7%	约 0.46 ms / 2.6%

进一步拆分：

• observation getter terms 总和只有 MuJoCo 0.92 ms、MJWarp 0.80 ms；post/copy/noise/isfinite/concat residual 约为 3.46 / 3.04 ms，所以继续逐个优化 getter 的上限较低。
• command.post_compute 的第二次 fused copy 已只需约 0.28–0.33 ms；剩余约 3.9 ms 主要是 relative transforms / body-frame 转换。
• 五个目标 term 优化后仅 1.205 / 1.184 ms，占 update_state 约 6.3% / 6.7%。在有限改动范围内，这一组已不再是首要优化对象。
• 完整 env.step 中，MuJoCo 最大单项仍是 physics，优化后约 49.10 ms；MJWarp 的主要分区为 update_state 17.80 ms、reset_done 14.69 ms、physics 13.57 ms。

下一轮若继续优化，优先级应是 command transform、observation 的通用 post/concat pipeline；MuJoCo 若追求更大的完整 step 收益，则 physics 端的杠杆更大。它们是新的 scope，不属于 #1304 的五个 term/body-state copy 收敛。

精确复现信息

A/B commits

角色	Commit	说明
优化前	d95c85f0b782f59e0174d93e4f7c21aa46b58f7a	#1304 子项进入前的基线
优化后	9cca010d2c97910a6704a3da80b0d7f347103591	PR #1314 当前 head；包含 #1309、#1310、#1312

优化后分支为 dev/issue-1304-motion-numba-body-state，目标分支为 dev/issue-1042-manager-based-api；截至本帖创建时 PR #1314 仍等待 maintainer 审批和手动 merge。

硬件 / 系统 / runtime

项目	本次环境
Host	yves-9950x3d2
OS / kernel	Ubuntu 24.04.4 LTS；Linux 7.0.0-30-generic x86_64
CPU	AMD Ryzen 9 9950X3D2，1 socket，16 cores / 32 threads，192 MiB L3
RAM	60 GiB
GPU	NVIDIA GeForce RTX 4090，49140 MiB；driver 595.84
Python / uv	Python 3.13.14；uv 0.12.5
Numeric / sim	NumPy 2.4.4；Numba 0.67.0；MuJoCo 3.10.0；Warp 1.14.0；mujoco-warp 3.10.0.3
Numba runtime	未设置外部 NUMBA_*、OMP 或 BLAS thread 环境变量；优化后代码默认 workqueue threading layer、8 threads
MuJoCo	SAC config 启用 adaptive chunk sizing；本次 logger 报告 adaptive chunk cache hit，default chunk size 12；A/B 复用相同 host cache 条件

MJWarp 在上述 RTX 4090 上运行；Numba rewards 与 fused host-copy 是 CPU 路径。复测时建议固定 CPU governor、关闭其他重负载并记录温度/频率；跨机器绝对耗时不能直接比较，A/B 应在同一台机器、同一依赖环境中完成。

Workload 与统计规则

• Hydra owner：conf/sac 下 task=g1_motion_tracking/{mujoco,mjwarp}；通过正式 backend owner 选择后端，不用 training.sim_backend 单独切换。
• num_envs=8192。
• env.seed=123；random action RNG seed 20260826，action 范围 [-1, 1]，float32。
• 每个进程先 env.init_state()；profiler wrapper 在 init 之后安装，因此初始化与 Numba 编译不进入统计。
• 每个进程 warmup 10 steps，随后清空 term 统计，再 measure 100 steps。
• 每个 commit/backend 分别运行：无埋点阶段测试 3 个独立进程；完整 term 埋点测试 3 个独立进程。
• 阶段值：每个进程 100 steps 取 median，再对三次独立运行取 median。
• term 值：每个进程 100 次完整 callable 取 mean，再对三次独立运行取 median。
• update_state_ms / env_step_total_ms / reset_done_ms 直接读取 NpEnv.step 已有的 perf_counter timing。
• 无埋点和有埋点结果来自不同进程；canonical 阶段表只使用无埋点结果。不要把 instrumented update_state 当成阶段 A/B。

复跑命令

先准备两个 worktree，并用优化后的 lock/extras 建立一套共同依赖环境。PYTHONPATH 强制每次加载指定 worktree 的 src，--no-sync 防止在切换 A/B 源码时重写环境。

git clone https://github.com/unilabsim/UniLab.git unilab-update-state-ab
git -C unilab-update-state-ab worktree add ../unilab-before d95c85f0b782f59e0174d93e4f7c21aa46b58f7a
git -C unilab-update-state-ab worktree add ../unilab-after 9cca010d2c97910a6704a3da80b0d7f347103591

cd unilab-after
uv sync --extra mujoco --extra mjwarp

BEFORE="$(cd ../unilab-before && pwd)"
AFTER="$(pwd)"
ENV_PROJECT="$AFTER"

把下方两个折叠块分别原样保存为 /tmp/profile_update_state.py 和 /tmp/aggregate_update_state.py，然后运行：

run_ab_case() {
  local tree="$1"
  local prefix="$2"
  local backend="$3"
  local suffix

  for suffix in "" _r2 _r3; do
    PYTHONPATH="$tree/src" uv run --project "$ENV_PROJECT" --no-sync python \
      /tmp/profile_update_state.py \
      --repo "$tree" --backend "$backend" \
      --num-envs 8192 --warmup 10 --measure 100 --action-seed 20260826 \
      --out "/tmp/${prefix}_${backend}_phase${suffix}.json"
  done

  for suffix in _r2 _r3 _r4; do
    PYTHONPATH="$tree/src" uv run --project "$ENV_PROJECT" --no-sync python \
      /tmp/profile_update_state.py \
      --repo "$tree" --backend "$backend" \
      --num-envs 8192 --warmup 10 --measure 100 --action-seed 20260826 \
      --instrument \
      --out "/tmp/${prefix}_${backend}_terms${suffix}.json"
  done
}

run_ab_case "$BEFORE" base mujoco
run_ab_case "$AFTER" final mujoco
run_ab_case "$BEFORE" base mjwarp
run_ab_case "$AFTER" final mjwarp

uv run --project "$ENV_PROJECT" --no-sync python /tmp/aggregate_update_state.py

Profiler SHA-256：7de572e10c83af040b4bb08f2824ce95dfc48124cf890192b2b9fdc143233386
Aggregator SHA-256：500d5f239782776013611fd4224c2166ac8df76af26907efc7259cd8f958dea0

建议同时保存以下命令输出，便于判断跨机器差异：

git -C "$BEFORE" rev-parse HEAD
git -C "$AFTER" rev-parse HEAD
uname -a
lscpu
free -h
nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader
uv --version
uv run --project "$ENV_PROJECT" --no-sync python -c \
  'import sys,numpy,numba,mujoco,warp,mujoco_warp; print(sys.version); print(numpy.__version__,numba.__version__,mujoco.__version__,warp.__version__,mujoco_warp.__version__)'
env | rg '^(NUMBA|OMP|OPENBLAS|MKL|BLIS|VECLIB|NUMEXPR)_' | sort

下面嵌入本次实际使用的脚本；脚本只做 runtime wrapper，不修改仓库源码。

/tmp/profile_update_state.py（点击展开）

#!/usr/bin/env python3
"""One-shot A/B profiler for the G1 motion-tracking Manager-Based env.

This file intentionally lives outside the repository.  It wraps complete manager
term callables at runtime, so no profiling hooks or benchmark code are committed.
The wrapper is installed only after ``init_state``; initialization/JIT work is
therefore outside the measured window.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from collections import defaultdict
from pathlib import Path
from typing import Any, Callable

import numpy as np


class TimedCallable:
    """Callable proxy that records the complete underlying term invocation."""

    def __init__(self, name: str, fn: Callable[..., Any], env: Any, stats: dict[str, dict[str, list[float]]]):
        self.name = name
        self.fn = fn
        self.env = env
        self.stats = stats

    def _phase(self) -> str:
        return "reset" if bool(getattr(self.env, "_autoreset_reset_active", False)) else "step"

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        t0 = time.perf_counter_ns()
        try:
            return self.fn(*args, **kwargs)
        finally:
            elapsed = (time.perf_counter_ns() - t0) / 1.0e6
            self.stats.setdefault(self.name, {"step": [], "reset": []})[self._phase()].append(elapsed)

    def reset(self, *args: Any, **kwargs: Any) -> Any:
        reset_fn = getattr(self.fn, "reset", None)
        if reset_fn is None:
            return None
        return reset_fn(*args, **kwargs)


class TimedMethod:
    """Instance-bound method proxy used for coarse residual-breakdown timings."""

    def __init__(self, fn: Callable[..., Any], env: Any, name: str, stats: dict[str, dict[str, list[float]]]):
        self.fn = fn
        self.env = env
        self.name = name
        self.stats = stats

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        # CommandManager receives env_ids as a keyword on reset.  The env flag
        # also covers observation calls nested inside the autoreset lifecycle.
        env_ids = kwargs.get("env_ids")
        phase = "reset" if env_ids is not None or bool(getattr(self.env, "_autoreset_reset_active", False)) else "step"
        t0 = time.perf_counter_ns()
        try:
            return self.fn(*args, **kwargs)
        finally:
            elapsed = (time.perf_counter_ns() - t0) / 1.0e6
            self.stats.setdefault(self.name, {"step": [], "reset": []})[phase].append(elapsed)


class TimedCopyCallable:
    """Proxy for MotionCommand's bound full-batch body-state copy contract."""

    def __init__(self, fn: Callable[..., Any], env: Any, stats: dict[str, dict[str, list[float]]]):
        self.fn = fn
        self.env = env
        self.stats = stats

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        t0 = time.perf_counter_ns()
        try:
            return self.fn(*args, **kwargs)
        finally:
            elapsed = (time.perf_counter_ns() - t0) / 1.0e6
            self.stats.setdefault("command/body_state_copy", {"step": [], "reset": []})["step"].append(elapsed)


class TimedRobotRefresh:
    """Time only refresh calls that actually copy/gather body state."""

    def __init__(self, fn: Callable[..., Any], command: Any, env: Any, stats: dict[str, dict[str, list[float]]]):
        self.fn = fn
        self.command = command
        self.env = env
        self.stats = stats

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        force = bool(kwargs.get("force", False))
        env_ids = kwargs.get("env_ids")
        active = force or self.command._robot_cache_step != self.env.common_step_counter
        if not active:
            return self.fn(*args, **kwargs)
        t0 = time.perf_counter_ns()
        try:
            return self.fn(*args, **kwargs)
        finally:
            elapsed = (time.perf_counter_ns() - t0) / 1.0e6
            phase = "reset" if env_ids is not None else "step"
            self.stats.setdefault("command/robot_state_refresh", {"step": [], "reset": []})[phase].append(elapsed)
            detail = "force" if force else "first_read"
            self.stats.setdefault(
                f"command/robot_state_refresh/{detail}", {"step": [], "reset": []}
            )[phase].append(elapsed)


def _stats_dict() -> dict[str, dict[str, list[float]]]:
    return defaultdict(lambda: {"step": [], "reset": []})


def _compose_env(repo: Path, backend: str, num_envs: int) -> Any:
    # Make the selected worktree's source tree win over the installed editable
    # package.  The caller supplies PYTHONPATH as an additional guard.
    src = str(repo / "src")
    if src not in sys.path:
        sys.path.insert(0, src)

    from hydra import compose, initialize_config_dir
    from hydra.core.global_hydra import GlobalHydra
    from unilab.base import registry
    from unilab.base.config_adapter import BackendAdapter

    registry.ensure_registries()
    GlobalHydra.instance().clear()
    with initialize_config_dir(config_dir=str(repo / "conf" / "sac"), version_base="1.3"):
        # The mjwarp owner inherits the same SAC motion-tracking env declaration
        # as the MuJoCo owner.  Keeping the owner explicit mirrors issue #1292.
        cfg = compose(
            config_name="config",
            overrides=[
                f"task=g1_motion_tracking/{backend}",
                "env.seed=123",
                "hydra.run.dir=.",
                "hydra.output_subdir=null",
                "hydra/job_logging=disabled",
                "hydra/hydra_logging=disabled",
            ],
        )
    override = BackendAdapter(cfg, root_dir=repo, algo_name="sac").build_task_env_cfg_override()
    return registry.make(
        "G1MotionTrackingSAC",
        sim_backend=backend,
        env_cfg_override=override,
        num_envs=num_envs,
    )


def _install_term_wrappers(env: Any, stats: dict[str, dict[str, list[float]]]) -> None:
    for manager_name in ("reward_manager", "termination_manager"):
        manager = getattr(env, manager_name)
        for term_name, term_cfg in zip(manager._term_names, manager._term_cfgs, strict=False):
            full_name = f"{('reward' if manager_name == 'reward_manager' else 'termination')}/{term_name}"
            term_cfg.func = TimedCallable(full_name, term_cfg.func, env, stats)

    # Observation terms are useful for explaining the residual update_state
    # cost.  They are also wrapped as complete calls (not individual NumPy ops).
    obs_manager = env.observation_manager
    for group_name, term_names in obs_manager._group_obs_term_names.items():
        cfgs = obs_manager._group_obs_term_cfgs[group_name]
        for term_name, term_cfg in zip(term_names, cfgs, strict=False):
            term_cfg.func = TimedCallable(f"obs/{group_name}/{term_name}", term_cfg.func, env, stats)


def _install_coarse_wrappers(env: Any, stats: dict[str, dict[str, list[float]]]) -> None:
    # These are diagnostic-only wrappers.  They are not used for the canonical
    # uninstrumented update_state numbers reported by the caller.
    for obj, attr, name in (
        (env.command_manager, "compute", "manager/command_compute"),
        (env.command_manager, "post_compute", "manager/command_post_compute"),
        (env.observation_manager, "compute", "manager/observation_compute"),
        (env.observation_manager, "compute_group", "manager/observation_compute_group"),
        (env, "_map_observations", "manager/map_observations"),
    ):
        setattr(obj, attr, TimedMethod(getattr(obj, attr), env, name, stats))

    # Current commits bind one public fused copy callable; the pre-optimization
    # commit executes four np.take calls inside _refresh_robot_state instead.
    command = env.command_manager.get_term("motion")
    if hasattr(command, "_copy_robot_body_state"):
        command._copy_robot_body_state = TimedCopyCallable(command._copy_robot_body_state, env, stats)
    command._refresh_robot_state = TimedRobotRefresh(
        command._refresh_robot_state, command, env, stats
    )


def _clear_stats(stats: dict[str, dict[str, list[float]]]) -> None:
    for phases in stats.values():
        phases["step"].clear()
        phases["reset"].clear()


def _summary(values: list[float]) -> dict[str, float | int]:
    if not values:
        return {"count": 0, "mean_ms": 0.0, "median_ms": 0.0, "p95_ms": 0.0}
    arr = np.asarray(values, dtype=np.float64)
    return {
        "count": int(arr.size),
        "mean_ms": float(arr.mean()),
        "median_ms": float(np.median(arr)),
        "p95_ms": float(np.percentile(arr, 95)),
    }


def run_once(args: argparse.Namespace) -> dict[str, Any]:
    repo = Path(args.repo).resolve()
    env = _compose_env(repo, args.backend, args.num_envs)
    stats = _stats_dict()
    try:
        env.init_state()
        action_dim = int(env.action_space.shape[0])
        action_rng = np.random.default_rng(args.action_seed)

        if args.instrument:
            _install_term_wrappers(env, stats)
            _install_coarse_wrappers(env, stats)

        # Match issue #1292: random actions, ten warmup calls, then one hundred
        # measured calls.  Action generation is outside env.step's timer.
        for _ in range(args.warmup):
            actions = action_rng.uniform(-1.0, 1.0, size=(args.num_envs, action_dim)).astype(np.float32)
            env.step(actions)
        _clear_stats(stats)

        update_state_values: list[float] = []
        total_values: list[float] = []
        reset_values: list[float] = []
        backend_values: dict[str, list[float]] = defaultdict(list)
        done_counts: list[float] = []
        for _ in range(args.measure):
            actions = action_rng.uniform(-1.0, 1.0, size=(args.num_envs, action_dim)).astype(np.float32)
            state = env.step(actions)
            timing = state.info.get("timing", {})
            update_state_values.append(float(timing.get("update_state_ms", np.nan)))
            total_values.append(float(timing.get("env_step_total_ms", np.nan)))
            reset_values.append(float(timing.get("reset_done_ms", np.nan)))
            done_counts.append(float(timing.get("reset_done_count", 0.0)))
            for key, value in timing.items():
                if str(key).startswith("backend_"):
                    backend_values[str(key)].append(float(value))

        result: dict[str, Any] = {
            "repo": str(repo),
            "backend": args.backend,
            "num_envs": args.num_envs,
            "warmup": args.warmup,
            "measure": args.measure,
            "instrument": bool(args.instrument),
            "action_seed": args.action_seed,
            "env_seed": 123,
            "timing": {
                "update_state_ms": _summary(update_state_values),
                "env_step_total_ms": _summary(total_values),
                "reset_done_ms": _summary(reset_values),
                "reset_done_count": _summary(done_counts),
                "backend": {key: _summary(values) for key, values in sorted(backend_values.items())},
            },
        }
        if args.instrument:
            result["terms"] = {
                key: {phase: _summary(values) for phase, values in phases.items()}
                for key, phases in sorted(stats.items())
            }
            # The sum below is a sum of complete term calls.  Each reward and
            # termination term is called exactly once per measured update_state.
            step_terms = [
                values["step"]
                for key, values in stats.items()
                if key.startswith("reward/") or key.startswith("termination/")
            ]
            if step_terms:
                counts = [len(values) for values in step_terms]
                common = min(counts)
                per_step = np.zeros(common, dtype=np.float64)
                for values in step_terms:
                    per_step += np.asarray(values[:common], dtype=np.float64)
                result["term_sum_step_ms"] = _summary(per_step.tolist())
            else:
                result["term_sum_step_ms"] = _summary([])

        return result
    finally:
        env.close()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--repo", required=True)
    parser.add_argument("--backend", choices=("mujoco", "mjwarp"), required=True)
    parser.add_argument("--num-envs", type=int, default=8192)
    parser.add_argument("--warmup", type=int, default=10)
    parser.add_argument("--measure", type=int, default=100)
    parser.add_argument("--action-seed", type=int, default=20260826)
    parser.add_argument("--instrument", action="store_true")
    parser.add_argument("--out", default=None)
    args = parser.parse_args()
    result = run_once(args)
    payload = json.dumps(result, ensure_ascii=False, separators=(",", ":"))
    if args.out:
        Path(args.out).write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n")
    print(payload)


if __name__ == "__main__":
    main()

/tmp/aggregate_update_state.py（点击展开）

from __future__ import annotations

import json
from pathlib import Path
from statistics import median
from typing import Any


PHASE_FILES = {
    ("mujoco", "before"): [
        "/tmp/base_mujoco_phase.json",
        "/tmp/base_mujoco_phase_r2.json",
        "/tmp/base_mujoco_phase_r3.json",
    ],
    ("mujoco", "after"): [
        "/tmp/final_mujoco_phase.json",
        "/tmp/final_mujoco_phase_r2.json",
        "/tmp/final_mujoco_phase_r3.json",
    ],
    ("mjwarp", "before"): [
        "/tmp/base_mjwarp_phase.json",
        "/tmp/base_mjwarp_phase_r2.json",
        "/tmp/base_mjwarp_phase_r3.json",
    ],
    ("mjwarp", "after"): [
        "/tmp/final_mjwarp_phase.json",
        "/tmp/final_mjwarp_phase_r2.json",
        "/tmp/final_mjwarp_phase_r3.json",
    ],
}

TERM_FILES = {
    ("mujoco", "before"): [
        "/tmp/base_mujoco_terms_r2.json",
        "/tmp/base_mujoco_terms_r3.json",
        "/tmp/base_mujoco_terms_r4.json",
    ],
    ("mujoco", "after"): [
        "/tmp/final_mujoco_terms_r2.json",
        "/tmp/final_mujoco_terms_r3.json",
        "/tmp/final_mujoco_terms_r4.json",
    ],
    ("mjwarp", "before"): [
        "/tmp/base_mjwarp_terms_r2.json",
        "/tmp/base_mjwarp_terms_r3.json",
        "/tmp/base_mjwarp_terms_r4.json",
    ],
    ("mjwarp", "after"): [
        "/tmp/final_mjwarp_terms_r2.json",
        "/tmp/final_mjwarp_terms_r3.json",
        "/tmp/final_mjwarp_terms_r4.json",
    ],
}

TARGETS = [
    "termination/anchor_pos",
    "reward/motion_body_pos",
    "reward/motion_body_ori",
    "reward/motion_body_lin_vel",
    "reward/motion_body_ang_vel",
]


def load_many(files: list[str]) -> list[dict[str, Any]]:
    return [json.loads(Path(name).read_text()) for name in files]


def med(values: list[float]) -> float:
    return float(median(values))


def change(before: float, after: float) -> dict[str, float]:
    return {
        "before_ms": before,
        "after_ms": after,
        "saved_ms": before - after,
        "reduction_pct": (before - after) / before * 100.0,
        "speedup_x": before / after,
    }


result: dict[str, Any] = {}
for backend in ("mujoco", "mjwarp"):
    out: dict[str, Any] = {}
    phase_sets = {name: load_many(PHASE_FILES[(backend, name)]) for name in ("before", "after")}
    term_sets = {name: load_many(TERM_FILES[(backend, name)]) for name in ("before", "after")}

    for metric in ("update_state_ms", "env_step_total_ms", "reset_done_ms"):
        before = med([run["timing"][metric]["median_ms"] for run in phase_sets["before"]])
        after = med([run["timing"][metric]["median_ms"] for run in phase_sets["after"]])
        out[metric] = change(before, after)
        out[metric]["before_run_medians_ms"] = [
            run["timing"][metric]["median_ms"] for run in phase_sets["before"]
        ]
        out[metric]["after_run_medians_ms"] = [
            run["timing"][metric]["median_ms"] for run in phase_sets["after"]
        ]

    terms: dict[str, Any] = {}
    for target in TARGETS:
        before = med(
            [run["terms"][target]["step"]["mean_ms"] for run in term_sets["before"]]
        )
        after = med(
            [run["terms"][target]["step"]["mean_ms"] for run in term_sets["after"]]
        )
        terms[target] = change(before, after)

    def per_run_target_sum(run: dict[str, Any], targets: list[str]) -> float:
        return sum(run["terms"][target]["step"]["mean_ms"] for target in targets)

    for name, targets in (("four_body_rewards", TARGETS[1:]), ("target_five", TARGETS)):
        before = med([per_run_target_sum(run, targets) for run in term_sets["before"]])
        after = med([per_run_target_sum(run, targets) for run in term_sets["after"]])
        terms[name] = change(before, after)

    before_all = med([run["term_sum_step_ms"]["mean_ms"] for run in term_sets["before"]])
    after_all = med([run["term_sum_step_ms"]["mean_ms"] for run in term_sets["after"]])
    terms["all_reward_termination"] = change(before_all, after_all)
    before_other = med(
        [
            run["term_sum_step_ms"]["mean_ms"] - per_run_target_sum(run, TARGETS)
            for run in term_sets["before"]
        ]
    )
    after_other = med(
        [
            run["term_sum_step_ms"]["mean_ms"] - per_run_target_sum(run, TARGETS)
            for run in term_sets["after"]
        ]
    )
    terms["other_reward_termination"] = change(before_other, after_other)
    out["terms"] = terms

    coarse: dict[str, Any] = {}
    coarse_paths = {
        "command_compute": "manager/command_compute",
        "command_post_compute": "manager/command_post_compute",
        "observation_compute": "manager/observation_compute",
    }
    for label, key in coarse_paths.items():
        before = med([run["terms"][key]["step"]["mean_ms"] for run in term_sets["before"]])
        after = med([run["terms"][key]["step"]["mean_ms"] for run in term_sets["after"]])
        coarse[label] = change(before, after)

    def obs_term_sum(run: dict[str, Any]) -> float:
        return sum(
            phases["step"]["mean_ms"]
            for key, phases in run["terms"].items()
            if key.startswith("obs/")
        )

    before_obs_terms = med([obs_term_sum(run) for run in term_sets["before"]])
    after_obs_terms = med([obs_term_sum(run) for run in term_sets["after"]])
    coarse["observation_term_calls"] = change(before_obs_terms, after_obs_terms)
    before_obs_pipeline = med(
        [
            run["terms"]["manager/observation_compute"]["step"]["mean_ms"]
            - obs_term_sum(run)
            for run in term_sets["before"]
        ]
    )
    after_obs_pipeline = med(
        [
            run["terms"]["manager/observation_compute"]["step"]["mean_ms"]
            - obs_term_sum(run)
            for run in term_sets["after"]
        ]
    )
    coarse["observation_post_concat_residual"] = change(
        before_obs_pipeline, after_obs_pipeline
    )

    def refresh_per_update(run: dict[str, Any]) -> float:
        stat = run["terms"]["command/robot_state_refresh"]["step"]
        return stat["mean_ms"] * stat["count"] / run["measure"]

    before_refresh = med([refresh_per_update(run) for run in term_sets["before"]])
    after_refresh = med([refresh_per_update(run) for run in term_sets["after"]])
    coarse["two_robot_state_refreshes"] = change(before_refresh, after_refresh)
    coarse["two_robot_state_refreshes"]["calls_per_update"] = 2.0

    current_parts = []
    for run in term_sets["after"]:
        current_parts.append(
            run["terms"]["manager/command_compute"]["step"]["mean_ms"]
            + run["terms"]["manager/command_post_compute"]["step"]["mean_ms"]
            + run["terms"]["manager/observation_compute"]["step"]["mean_ms"]
            + run["term_sum_step_ms"]["mean_ms"]
        )
    coarse["current_accounted_total_ms"] = med(current_parts)
    coarse["current_accounted_share_pct"] = (
        coarse["current_accounted_total_ms"] / out["update_state_ms"]["after_ms"] * 100.0
    )
    coarse["current_share_pct"] = {
        "command_compute": coarse["command_compute"]["after_ms"]
        / out["update_state_ms"]["after_ms"]
        * 100.0,
        "command_post_compute": coarse["command_post_compute"]["after_ms"]
        / out["update_state_ms"]["after_ms"]
        * 100.0,
        "observation_compute": coarse["observation_compute"]["after_ms"]
        / out["update_state_ms"]["after_ms"]
        * 100.0,
        "all_reward_termination": terms["all_reward_termination"]["after_ms"]
        / out["update_state_ms"]["after_ms"]
        * 100.0,
    }
    out["coarse"] = coarse
    result[backend] = out

print(json.dumps(result, indent=2, ensure_ascii=False))