Repository navigation
perf(env): collapse duplicate action/control scalar syncs on collector hot path - #2486
Merged
Merged
Conversation
…r hot path Issue #2482 partial groundwork. CUDA-carrier runs without MPS expose every blocking scalar read on the apply_action path to dual-CUDA-context switching overhead. The collector step previously performed two hard bool() reads of 0-d finite verdicts per step: - TorchEnv._validate_control ran torch.isfinite(control).all() and read the scalar every step. - ActionManager.process_action(reuse_recorded_verdict=True) re-read the verdict that step() had already recorded. The reuse path now records without reading, and JointPositionAction keeps one device-resident verdict for its affine target. TorchEnv._validate_control publishes the manager verdict at the control-validation boundary, so the contract stays fail closed with one small verdict read instead of a full control reduction read. The verdict is cleared per step and never leaks. Benchmark parity was confirmed locally (RTX 4090, no-MPS, mujoco SAC g1_walk_flat, 2048 envs, 3x200 iters, tail-100; A/B by stashing only actions.py): baseline 64.9k steps/s / 5.65 ms apply_action vs patched 65.2k / 5.57 ms. The remaining no-MPS apply_action penalty is attributed outside this change (learner action publication non_blocking=False and host-bridge write_control hard sync) and is tracked for follow-up work.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
摘要
#2482 的部分前置清理。无 MPS 的 CUDA 载体下，collector/learner 双 CUDA context 切换会放大
apply_action路径上的每一次 blocking scalar read。本 PR 移除其中重复的两处：TorchEnv._validate_control：此前每步对全 control 执行torch.isfinite(control).all()后立即bool()读取（硬同步）。ActionManager.process_action(reuse_recorded_verdict=True)：step()已记录 device verdict，manager 重用路径又对同一个 0-d bool 做bool()（第二次硬同步）。修复方式：reuse 路径只记录不读取；
JointPositionAction为 affine 目标保留一个 device-resident verdict；TorchEnv._validate_control在 control-validation 边界发布 manager verdict —— fail-closed 合同不变，但每步只读一次小 verdict 而非全 control reduction 读。verdict 每步清空，不跨步泄漏。保留的行为边界（调研发现，未改变）
clip生效时+Infraw action 会先被 clamp 成有限值（如 1.0），随后 raw-based 的validate_actions()仍按 raw Inf 抛错 —— 现状合同为"raw 严格，applied 宽松"。本 PR 沿用该语义，未引入新漏洞。Benchmark（无回归）
本机 RTX 4090、无 MPS、mujoco SAC g1_walk_flat、seed=1、2048 envs、3 blocks × 200 iters、tail-100；A/B 成对（仅 stash
envs/mdp/actions.py切换）：吞吐 +0.5%、apply_action −0.08 ms，均在噪声内且方向一致（非回归）。
范围说明（不关闭 #2482）
issue 的验收目标（88k steps/s、apply_action 0.4 ms）未由本 PR 达成，剩余 no-MPS 惩罚归因于本 PR 范围之外：
non_blocking=False（double_buffer_runner.py:1149）；write_control的每步硬synchronize（unisimtensor_io.py）。这两项需后续 PR 处理；本 PR 定位为该路径的前置清理，保持 benchmark 归因干净。
验证
make check✅（ruff / mypy / pyright）make test✅ 2718 passedmake test-all✅（34/34 module-mode + 35/35 script-mode）tests/base/test_torch_env.py2 个用例（manager verdict fail-closed、verdict 不跨步泄漏）Closes 部分推进 #2482（issue 保持 open，剩余两项在范围说明中列出）。