diff --git a/docs/sphinx/source/en/2-user_guide/2-algorithms/4-flash_sac.md b/docs/sphinx/source/en/2-user_guide/2-algorithms/4-flash_sac.md index c507edcc2..17f6e7f05 100644 --- a/docs/sphinx/source/en/2-user_guide/2-algorithms/4-flash_sac.md +++ b/docs/sphinx/source/en/2-user_guide/2-algorithms/4-flash_sac.md @@ -30,6 +30,28 @@ playback video. See {doc}`/en/1-getting_started/3-evaluation_and_playback`. - `algo.algo_params.actor_num_blocks=2` - `algo.algo_params.critic_num_blocks=2` +## Learner acceleration controls + +The default single-NVIDIA learner enables all three acceleration layers. They +remain separate configuration choices: + +- `algo.algo_params.use_compile` controls CUDA Inductor compilation. Set it to + `false` for the eager per-update compatibility path; this also disables the + dependent whole-cycle graph. +- `algo.algo_params.compile_full_objectives` selects complete actor/critic + objectives instead of the smaller loss-tensor compile regions. It does not + control CUDA Graph capture. +- `algo.algo_params.use_whole_cycle_cuda_graph` controls learner-owned + whole-cycle CUDA Graph capture. Set only this field to `false` to retain + Inductor compilation without whole-cycle capture. + +For example, run the eager debugging path with: + +```bash +uv run train --algo flashsac --task g1_walk_flat --sim mujoco \ + algo.algo_params.use_compile=false +``` + FlashSAC requires synchronized collection and the same sole replay path as SAC: bounded host ingress plus one complete replay ring on a CUDA or Apple MPS learner device. CPU and XPU training are unsupported. diff --git a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/4-flash_sac.md b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/4-flash_sac.md index be13c9ff5..b138a57d7 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/4-flash_sac.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/4-flash_sac.md @@ -28,6 +28,24 @@ uv run train --algo flashsac --task go2_joystick_flat --sim mujoco training.no_p - `algo.algo_params.actor_num_blocks=2` - `algo.algo_params.critic_num_blocks=2` +## Learner 加速开关 + +单 NVIDIA learner 默认开启全部三层加速,但它们是三个独立的配置选择: + +- `algo.algo_params.use_compile` 控制 CUDA Inductor 编译。设为 `false` 后走 + eager 的逐 update 兼容路径,同时也会关闭依赖编译的 whole-cycle graph。 +- `algo.algo_params.compile_full_objectives` 选择编译完整 actor/critic objective, + 还是较小的 loss-tensor 区域;它不控制 CUDA Graph 捕获。 +- `algo.algo_params.use_whole_cycle_cuda_graph` 控制 learner 自己捕获并回放整轮 + CUDA Graph。只把此字段设为 `false`,可以保留 Inductor 编译而关闭整轮捕获。 + +例如,使用 eager 调试路径: + +```bash +uv run train --algo flashsac --task g1_walk_flat --sim mujoco \ + algo.algo_params.use_compile=false +``` + FlashSAC 要求同步采集,并与 SAC 共用唯一 replay 路径:有界 host ingress 加 一个驻留在 CUDA 或 Apple MPS learner device 上的完整 replay ring。CPU 与 XPU training 不受支持。 diff --git a/src/unilab/conf/flashsac/config.yaml b/src/unilab/conf/flashsac/config.yaml index f33768a00..1d0a5c9ef 100644 --- a/src/unilab/conf/flashsac/config.yaml +++ b/src/unilab/conf/flashsac/config.yaml @@ -46,6 +46,7 @@ algo: amp_dtype: auto use_compile: true compile_full_objectives: true + use_whole_cycle_cuda_graph: true training: task_name: G1WalkFlat diff --git a/src/unilab/structured_configs.py b/src/unilab/structured_configs.py index 10630582c..b5bcb5ab9 100644 --- a/src/unilab/structured_configs.py +++ b/src/unilab/structured_configs.py @@ -83,6 +83,8 @@ class FlashSACAlgoParams: n_step: int = 1 amp_dtype: str = "auto" use_compile: bool = True + compile_full_objectives: bool = True + use_whole_cycle_cuda_graph: bool = True @dataclass diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index e9dbeb2c7..72bfdd8ea 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -46,6 +46,8 @@ def test_flashsac_config_defaults(): assert cfg.algo_params.normalize_reward is True assert cfg.algo_params.amp_dtype == "auto" assert cfg.algo_params.use_compile is True + assert cfg.algo_params.compile_full_objectives is True + assert cfg.algo_params.use_whole_cycle_cuda_graph is True def test_ppo_config_defaults(): @@ -125,6 +127,9 @@ def test_offpolicy_flashsac_g1_task_overrides(): overrides=["task=g1_walk_flat/mujoco"], ) assert cfg.algo.algo == "flashsac" + assert cfg.algo.algo_params.use_compile is True + assert cfg.algo.algo_params.compile_full_objectives is True + assert cfg.algo.algo_params.use_whole_cycle_cuda_graph is True def test_offpolicy_flashsac_go2_task_overrides():