UniLab FlashSAC 性能优化

本仓库保存 FlashSAC 的 full-objective compile 优化、补丁、测试和可复现实验脚本。优化目标是在不改变算法更新次数和数学定义的前提下，减少 GPU kernel launch、CPU/GPU 同步、无效梯度和 metrics 拷贝开销。

详细原理、完整路径定义和全部实验数据见 FLASH_SAC_OPTIMIZATION.md。

固定版本

物理训练结果只适用于下面这组版本。不要把 unilab-rl、unisim-core 或 MuJoCo 升级后，仍然把新结果和本仓库的结果直接比较：

组件	固定版本	备注
UniLab	1.3.2	Git 82fa794a88ee9466d8588e7c4a3a00dfeadd451f
unilab-rl 兼容依赖	1.4.0	UniLab owner 环境中的精确 wheel 版本
FlashSAC 优化源码	Git abb3d82	本仓库 src/uni_rl；基准脚本优先使用它
unisim-core	1.7.8	精确 pin，不使用 >=
PyTorch	2.8.0+cu128	CUDA 12.8 wheel
MuJoCo	3.11.0	与 UniLab MuJoCo extra 匹配
GPU	NVIDIA GeForce RTX 4090	结果中的硬件

REPRODUCE_VERSIONS.env 保存了同一组机器可读的值。这里的 unilab-rl==1.4.0 是 UniLab 的兼容依赖；运行优化基准时，脚本会把本仓库 的 src/uni_rl 放在 PYTHONPATH 最前面，因此真正执行的是表中指定的 FlashSAC 优化源码快照。运行前可以检查：

本仓库 pyproject.toml 中的 unilab-rl 版本号 1.2.1 是这份源码快照的 历史发行元数据，不是物理训练使用的 UniLab owner wheel。不要把这两个版本 混为一谈；以 PYTHONPATH 指向的源码和下面的 package_versions 为准。

cd /path/to/UniLab
PYTHONPATH=/path/to/UniLab-FlashSAC-Optmization/src \
uv run --locked --no-sync -- python -c \
  'import importlib.metadata as m, torch, mujoco, unisim, unilab, uni_rl; \
   print("unilab", m.version("unilab")); \
   print("unilab-rl", m.version("unilab-rl")); \
   print("unisim-core", m.version("unisim-core")); \
   print("torch", torch.__version__); \
   print("mujoco", mujoco.__version__); \
   print("unilab", unilab.__file__); \
   print("uni_rl", uni_rl.__file__); \
   print("unisim", unisim.__file__)'

输出应分别为 1.3.2、1.4.0、1.7.8、2.8.0+cu128 和 3.11.0。如果 uni_rl.__file__ 没有指向本仓库的 src/uni_rl/__init__.py，说明没有使用优化源码。

做了什么

• 将 actor/critic 的完整 objective 放进 torch.compile，扩大 Inductor 的融合范围。
• actor 更新时冻结 critic 参数，只保留 action 到 critic 输出的梯度路径。
• finite guard 留在 GPU，避免每次 optimizer step 都同步到 CPU。
• metrics 延迟到一轮 learner update 末尾统一 D2H。
• 使用固定 metrics buffer，避免 Inductor CUDA Graph 或手工 CUDA Graph replay 覆盖返回值。
• 正确处理 Inductor 内部 CUDA Graph 与外层手工 CUDA Graph：纯 compile 路径允许 Inductor Graph，手工 Graph 路径关闭嵌套 capture。

推荐默认配置：

algo_params:
  use_compile: true
  compile_full_objectives: true
  use_cuda_graph_critic: false
  use_cuda_graph_actor: false

手工 CUDA Graph 只建议用于 batch 和输入 shape 固定、追求极限速度且能承担更复杂 graph 生命周期管理的场景。

不同配置的修改方式

下面的配置都作用于 algo.algo_params。可以临时在命令行覆盖，也可以把 相同字段写入 UniLab 的 src/unilab/conf/flashsac/config.yaml。每次修改后都 应保留输出 JSON，便于比较同一统计口径下的指标。

1. 推荐：full-objective compile

这是本报告的优化主路径：编译完整 actor/critic objective，不启用外层手工 CUDA Graph。

algo_params:
  use_compile: true
  compile_full_objectives: true
  use_cuda_graph_critic: false
  use_cuda_graph_actor: false

命令行写法：

algo.algo_params.use_compile=true \
algo.algo_params.compile_full_objectives=true \
algo.algo_params.use_cuda_graph_critic=false \
algo.algo_params.use_cuda_graph_actor=false

2. 只编译局部 loss

用于判断“扩大 compile 边界”本身带来的收益：

algo_params:
  use_compile: true
  compile_full_objectives: false
  use_cuda_graph_critic: false
  use_cuda_graph_actor: false

3. eager 兼容路径

用于排查编译器或 CUDA Graph 问题。它仍保留本仓库的 deferred metrics 等 安全优化，但不使用 Inductor：

algo_params:
  use_compile: false
  compile_full_objectives: false
  use_cuda_graph_critic: false
  use_cuda_graph_actor: false

4. 外层手工 CUDA Graph

只有输入 shape、optimizer state 和 staging buffer 全部稳定时才尝试：

algo_params:
  use_compile: true
  compile_full_objectives: true
  use_cuda_graph_critic: true
  use_cuda_graph_actor: true

手工 Graph 由 learner 负责 capture/replay；本仓库会关闭 Inductor 内部的 CUDA Graph Trees，避免嵌套 capture。它适合单独做峰值实验，不是报告中真实 MuJoCo 结果的默认路径。

如果已经把上游 Issue #94 的独立 whole-cycle 开关合入 unilab_rl，推荐使用 下面的等价配置关闭外层整轮 Graph，同时保留 Inductor；它需要对应的新核心 版本，不能直接用于本仓库的旧 API：

algo_params:
  use_compile: true
  compile_full_objectives: true
  use_whole_cycle_cuda_graph: false

结果摘要

两种 update 口径均使用 policy_frequency=2：

• 1 actor / 2 critic：每 iter 有 2 次 critic、1 次 actor、1 次 temperature update。
• 4 actor / 8 critic：每 iter 有 8 次 critic、4 次 actor、4 次 temperature update。

真实 MuJoCo 训练

RTX 4090、G1 Walk Flat、256 env、batch 256、1000 iter，统计最后 500 iter。单位为 ms：

update 口径	路径	learner mean	median	p90	p95	learner 降幅	iter mean
1 actor / 2 critic	deferred eager	10.118	9.965	10.760	11.092	—	11.009
1 actor / 2 critic	full-objective compile	4.323	4.397	4.874	5.088	57.3%	5.341
4 actor / 8 critic	deferred eager	36.885	36.183	39.010	40.171	—	37.760
4 actor / 8 critic	full-objective compile	13.304	13.145	14.344	14.888	63.9%	14.235

这里的 learner 是一个 iter 内全部 update 的总耗时，不是单次 critic 或 actor update。iter 是完整 runner iteration 墙钟时间。collector 与 learner 采用 double buffer，会发生重叠，因此 iter != learner + collector。

update 数量扩大 4 倍时，真实训练的优化后 learner mean 从 4.323 ms 增加到 13.304 ms，只增加 3.08 倍。原因是每 iter 固定开销只支付一次，并且 metrics D2H 和最终 CUDA synchronize 也只做一次；详细定量分析见优化报告。

复现实验

1. 目录和环境

假设目录为：

/path/to/UniLab
/path/to/UniLab-FlashSAC-Optmization

UniLab 环境需能导入 torch、unisim、MuJoCo 和 TensorBoard。先检查：

cd /path/to/UniLab
uv run --project /path/to/UniLab --locked --no-sync -- python -c \
  'import torch, unisim; print(torch.__version__); print(torch.cuda.get_device_name(0)); print(unisim.__file__)'

不要直接执行 uv run /path/to/UniLab-FlashSAC-Optmization/...。这种写法可能让 uv 选择优化仓库自己的 .venv，导致 ModuleNotFoundError: unisim。下面所有命令都显式使用 UniLab 项目环境。

2. 真实物理训练

下面命令真正调用 UniLab 的 train_flashsac.py 并创建 MuJoCo 环境，同时保留生产 Rich 实时面板。默认训练 1000 iter，统计最后 500 iter：

uv run --project /path/to/UniLab --locked --no-sync -- python \
  /path/to/UniLab-FlashSAC-Optmization/experiments/benchmark_flashsac_training.py \
  --unilab-root /path/to/UniLab \
  --backend mujoco --iterations 1000 --summary-last 500 \
  --num-envs 256 --batch-size 256 --replay-buffer-n 32 \
  --learning-starts 8 \
  --updates-per-step 2 --policy-frequency 2 \
  --compile \
  --output /path/to/UniLab-FlashSAC-Optmization/results/physical_u2_compile.json

四组对照只需修改以下参数：

实验	参数
1 actor / 2 critic，优化后	--updates-per-step 2 --policy-frequency 2 --compile
1 actor / 2 critic，deferred eager	--updates-per-step 2 --policy-frequency 2 --no-compile
4 actor / 8 critic，优化后	--updates-per-step 8 --policy-frequency 2 --compile
4 actor / 8 critic，deferred eager	--updates-per-step 8 --policy-frequency 2 --no-compile

每组使用不同的 --output 文件名。JSON 中：

• rows_all：全部共同 timing 样本；
• rows：最后 500 个参与统计的样本；
• learner_train_ms、collector_cycle_ms、iter_ms：mean、median、p90、p95；
• training_config：每 iter 的 critic、actor、temperature update 数量；
• device.package_versions：实际启动环境发现的 unilab、unilab-rl、 unisim-core 和 mujoco 版本；device.flashsac_source：实际使用的优化源码路径。

验证指标和通过标准

真实训练必须同时查看三组时间，每组报告 mean、median、p90、p95：

JSON 字段	含义	推荐用途
learner_train_ms	一个 iter 内全部 critic/actor/temperature update 的 learner 总耗时	判断算法更新本身是否变快
collector_cycle_ms	collector 一个周期的耗时，包含等待 learner inference 的影响	判断采集端是否被 learner 拖慢
iter_ms	完整 runner iteration 墙钟时间	最终端到端指标

本报告的 4090 基线为：

口径	路径	learner mean / p95	iter mean / p95
1 actor / 2 critic	deferred eager	10.118 / 11.092 ms	11.009 / 12.032 ms
1 actor / 2 critic	full-objective compile	4.323 / 5.088 ms	5.341 / 6.213 ms
4 actor / 8 critic	deferred eager	36.885 / 40.171 ms	37.760 / 41.041 ms
4 actor / 8 critic	full-objective compile	13.304 / 14.888 ms	14.235 / 15.886 ms

复现实验时，至少应满足：

1. JSON 中 summary_last=500，且 iterations=1000；
2. training_config 明确记录 critic、actor、temperature 的 update 数；
3. 三组 timing 字段都存在，并同时给出 mean、median、p90、p95；
4. iter_ms 是判断端到端性能的主指标，不能只报告单次 update；
5. 数值检查和 replay 稳定性检查均返回 passed=true；
6. 在本固定环境下，4 actor / 8 critic 的 full-objective compile 应保持约 14.235 ms/iter mean、15.886 ms/iter p95 的量级，明显低于 20 ms。

3. 数值和稳定性检查

uv run --project /path/to/UniLab --locked --no-sync -- python experiments/validate_numerics.py
uv run --project /path/to/UniLab --locked --no-sync -- python experiments/validate_replay_stability.py