FlashSAC

FlashSAC 通过 src/unilab/scripts/train_flashsac.py 运行,拥有独立的配置树。使用 --algo flashsac 选择它;默认值内联在 src/unilab/conf/flashsac/config.yaml 中,实现位于 unilab.rl.algos.flash_sac 下。

它与 SAC 共用 off-policy runner 设计,但默认网络并不相同:actor 使用 block-based 结构,critic 使用 distributional(categorical)Q 变体。

快速开始

uv run train --algo flashsac --task g1_walk_flat --sim mujoco
uv run train --algo flashsac --task go2_joystick_flat --sim mujoco training.no_play=true

关键字段

对于 off-policy 回放路径(src/unilab/scripts/train_flashsac.py / CLI --algo flashsac),设 置 training.export_onnx=false 可在仍然录制回放视频的同时跳过 policy.onnx 导出。 参见 评估与回放。

  • algo.algo_log_name=flash_sac

  • algo.num_envs=1024

  • algo.max_iterations=5000

  • algo.tau=0.01

  • algo.save_interval=1000

  • 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 调试路径:

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 不受支持。

日志根目录为 logs/flash_sac/<task>/。

Checkpoint 兼容性

每个 FlashSAC checkpoint 都会记录 actor_semantics 载荷 (schema_version/kind/builder/kwargs/obs_dim/action_dim/ action_transform,由 unilab.rl.algos.flash_sac.actor_spec 写入)。训练 续训与 play 两条加载路径都会在恢复任何权重之前对其 fail-closed 校验,内置 actor 与经 algo.algo_params.actor_config 注入的自定义 actor 同样适用。

pre-contract checkpoint(缺失或为 null 的 actor_semantics 条目)对所有 actor 一律不再支持,包括内置 actor:加载会 fail closed,错误信息说明最低 支持的格式。唯一恢复路径是使用当前版本重新训练策略。支持基线从 1.3.5 之后 首个包含该变更的版本开始,不回 port 到任何 1.3.5 patch release。

多卡数据并行

FlashSAC 与 SAC 共用同一套多卡数据并行机制:CUDA_VISIBLE_DEVICES 下每个 rank 各跑一 套独立的 learner+collector;启动时广播完整模型状态,稳态在每个实际 optimizer step 前分别平均 actor / critic / temperature 梯度。仅 rank 0 保存 checkpoint。用法与限制见 SAC 的”多卡数据并行”小节。

export CUDA_VISIBLE_DEVICES=<gpu-a>,<gpu-b>
uv run train --algo flashsac --task g1_walk_flat --sim mujoco \
  training.no_play=true