调研报告：MJLab / RSL-RL 多卡训练的数据并行与同步语义

[TATP-233]

摘要

直接回答最关心的几个问题：

• MJLab 的官方多卡入口是 单机一张 GPU 一个进程：torchrunx.Launcher 负责 spawn，RSL-RL 根据 WORLD_SIZE / RANK / LOCAL_RANK 初始化 NCCL process group。
• 这是同步数据并行，但 MJLab 的 num_envs 语义与 Holosoma 不同：每个 rank 都运行完整的 num_envs，不是把一个固定全局环境数拆到多卡。若每卡环境数为 N、rollout 长度为 T、GPU 数为 W，每 iteration 的全局新样本数是 W * N * T。
• 每个 rank 有自己的环境、policy copy 和 RolloutStorage；GAE、advantage normalization、mini-batch 切分都在本地完成。不 all-gather rollout，也不跨 rank 交换 observation / action / transition。
• PPO 训练阶段同步的是 梯度：每个本地 mini-batch backward 后，把 actor、critic（以及启用时的 RND predictor）现有梯度 flatten 成一个 tensor，执行一次阻塞式 all_reduce(SUM) / W，然后每个 rank 本地 clip 和 optimizer.step()。
• 参数/模型状态只在 learn() 开始时由 rank 0 广播一次；稳态没有周期性参数平均。各 rank 依靠相同初始模型、相同平均梯度和相同 optimizer state 保持可训练参数一致。
• PPO 每 iteration 的梯度同步次数是 num_learning_epochs * num_mini_batches。默认 5 * 4 = 20 次；默认 adaptive-KL 还会在每个 mini-batch 额外做一次标量 KL all-reduce 和一次 learning-rate broadcast，因此默认合计是 20 次大梯度 all-reduce + 20 次 KL all-reduce + 20 次 LR broadcast。
• RSL-RL 当前内置算法只有 PPO 和 Student-Teacher Distillation。仓库里没有 SAC、TD3、replay buffer 或其他 off-policy RL pipeline；Distillation 不是 off-policy RL。

这是一份源码审计，不是多卡性能 benchmark。审计时本地与远端 main 一致：

• mujocolab/mjlab@0fb8a68，tag v1.6.0；
• leggedrobotics/rsl_rl@c281c32，即 v5.4.2-1-gc281c32。

MJLab 固定依赖 rsl-rl-lib==5.4.2：pyproject.toml#L34-L51。本地 RSL-RL 相比 v5.4.2 只多一个 console log 字段顺序修复，本文涉及的 runner、PPO、Distillation、storage 和 distributed 代码完全相同。

1. 启动方式与进程拓扑

MJLab 先解析 CUDA_VISIBLE_DEVICES 和 --gpu-ids。单卡直接调用训练函数；多卡时创建 torchrunx.Launcher(hostnames=["localhost"], workers_per_host=num_gpus)，每卡拉起一个 worker，并明确令 backend=None，让 RSL-RL 自己初始化 process group：train.py#L182-L227。因此 MJLab 提供的这个 CLI launcher 是单机多卡入口。

每个 worker 使用 cuda:LOCAL_RANK，同时令环境和 agent seed 等于 base_seed + global_rank：train.py#L49-L69。环境构造时该 seed 会设置 Python、NumPy、Torch 和 Warp RNG：manager_based_rl_env.py#L177-L205、random.py#L9-L24。不同 rank 因此得到不同 rollout；模型也会先被不同 seed 初始化，但随后会被 rank 0 的模型状态覆盖。

RSL-RL 在 runner 构造期间检查 WORLD_SIZE > 1，读取 local/global rank，校验 device 后调用：

torch.distributed.init_process_group(
    backend="nccl", rank=gpu_global_rank, world_size=gpu_world_size
)

证据：on_policy_runner.py#L207-L249。模型没有包 DistributedDataParallel；后续所有 collective 都由算法代码显式调用。

2. Rollout 如何做数据并行

MJLab 明确把多卡定义为 “data-parallel, not work-splitting”：每张卡都创建完整的 num_envs，所以：

N = 每个 rank 的 num_envs
T = num_steps_per_env
W = world_size

local rollout size  = N * T
global rollout size = W * N * T

官方说明见 distributed_training.rst#L38-L51。每个进程各自持有环境、policy copy 和 N * T 的经验 buffer：distributed_training.rst#L61-L88。RSL-RL 实际也按 [T, N, ...] 在每个 rank 上单独分配 RolloutStorage：rollout_storage.py#L125-L164。

一次 iteration 的时序是：

每个 rank 独立采 T 个环境步
        -> 本地计算 GAE / returns / advantage normalization
        -> 本地生成 mini-batch 并 backward
        -> 每个 mini-batch 做一次跨 rank 梯度平均
        -> 各 rank 本地执行相同 optimizer.step()

rollout 期间没有 collective；不同速度的 rank 会在 update 阶段遇到第一个 collective 时隐式等待。runner 的采集、return 计算和 update 调用位置见 on_policy_runner.py#L79-L108。

Feed-forward PPO 的 batch 细节

设 M = num_mini_batches、E = num_learning_epochs。每个 rank 的本地 mini-batch size 是：

b_local = floor(N * T / M)
b_effective_global = W * b_local
optimizer updates per iteration = E * M

各 rank 分别调用本地 randperm，没有全局 shuffle：rollout_storage.py#L224-L258。有两个容易忽略的精确语义：

1. 如果 N*T 不能被 M 整除，只使用前 M * floor(N*T/M) 个 flatten 后的样本，末尾 remainder 不参加更新；
2. randperm 在 epoch loop 外只生成一次，所以同一 iteration 的 E 个 learning epoch 复用相同的样本划分和顺序，不会每个 epoch 重洗牌。

RNN 路径按环境维切分，单 rank 每个 mini-batch 使用 floor(N/M) 条环境流；N % M 个尾部环境不会进入更新，然后再按 done 切成 padded trajectories：rollout_storage.py#L260-L329。

GAE 与 advantage normalization 是 rank-local

每个 rank 只对自己的 storage 计算 GAE；默认 advantage 标准化也只使用本 rank 的 N*T 个样本。如果打开 per-mini-batch normalization，则统计范围更小，只是当前 rank 的当前 mini-batch：ppo.py#L163-L188、ppo.py#L200-L214。

所以在各 rank mini-batch 大小相等时，平均 local mean gradient 等价于对各 rank 已经完成本地预处理的 mini-batch 并集求 mean gradient；但它不严格等价于先拼成一个全局 rollout，再用全局 mean/std 做一次 advantage normalization。

对 MJLab scaling 文档的一点校正

文档一方面说每个 iteration 的速度大致不变，另一方面又说 GPU 增多而不调整 max_iterations 会让训练时间按比例增长：distributed_training.rst#L49-L58。从上述公式能直接推出的是：固定 iteration 数时，总样本数和 aggregate GPU compute 约按 W 增长；wall-clock 是否增长取决于通信和实际吞吐，不会仅由样本公式自动按 W 增长。若目标是保持相同总样本数，应把 iteration 数约除以 W；这不等同于“保持相同 wall-clock”。

3. PPO：梯度同步还是参数同步，多久一次

答案是两者都有，但发生在不同阶段：

对象	机制	周期
actor + critic model state	rank 0 broadcast_object_list(state_dict)	learn() 开始时一次
actor + critic 梯度	flatten 后 all_reduce(SUM) / W	每个 PPO mini-batch 一次
optimizer moments	无 collective	每个 rank 对相同平均梯度本地 step
模型参数	稳态无周期 broadcast / parameter averaging	依靠同步梯度保持一致
adaptive-KL	标量 all_reduce(SUM) / W	每个 mini-batch 一次
learning rate	rank 0 标量 broadcast	adaptive schedule 下每个 mini-batch 一次

启动广播在 on_policy_runner.py#L64-L74 和 ppo.py#L451-L463。这次广播使用完整 actor/critic state_dict，所以既包含 trainable parameters，也包含它们当时已有的 registered buffers。

每个 mini-batch 先在本地完成 PPO backward；若启用 RND，也先完成 RND backward。随后 reduce_parameters() 把 actor、critic 和 RND 中所有 grad is not None 的梯度拼成一个大 tensor，一次阻塞式 all-reduce 求平均，再拷回各 parameter gradient；之后才分别做 gradient clipping 和本地 optimizer step：ppo.py#L280-L307、ppo.py#L465-L488。

adaptive schedule 会先对 local KL mean 做全局平均，只让 rank 0 决定新 LR，再把 LR 广播回所有 rank：ppo.py#L233-L259。

因此精确通信次数为：

U = E * M

fixed LR:     U 次 gradient all-reduce / iteration
adaptive KL:  U 次 gradient all-reduce
            + U 次 KL scalar all-reduce
            + U 次 LR scalar broadcast / iteration

RSL-RL 和 MJLab 的默认 E=5, M=4，所以 U=20：ppo.py#L34-L52、mjlab rl/config.py#L41-L61。这些 collective 本身就是 rendezvous；没有另加每 iteration 参数广播或显式 barrier。

4. 可训练权重同步，不代表完整训练状态全局一致

这是实现里最值得注意的边界。正常前提下，actor/critic 的 trainable parameters 与 optimizer moments 会保持一致；但以下状态没有周期同步：

Observation normalizer

actor/critic 在每个本地 env step 后只用本 rank observation 更新 EmpiricalNormalization：ppo.py#L128-L137、mlp_model.py#L170-L177。mean、variance 和 count 是 buffers，不进入梯度 all-reduce：normalization.py#L15-L66。

这不是冷门配置：MJLab 的 G1 velocity 等任务明确为 actor 和 critic 开启 observation normalization：g1/rl_cfg.py#L10-L44。所以多卡时“可训练权重相同”，但各 rank 的完整 policy state 可能因 normalizer buffer 不同而不同；最终 rank 0 checkpoint 保存的是 rank 0 那份统计量。

同理，若 RSL-RL 的 CNN 配置使用 BatchNorm，其 running mean/variance 也不会周期同步；CNN 确实支持 norm="batch"：cnn.py#L16-L50、cnn.py#L89-L100。

RND extension

RSL-RL 新建 RND 时会各自初始化 predictor 和 frozen target，但 startup broadcast 只追加 rnd.predictor.state_dict()，不广播 target：rnd.py#L92-L124、ppo.py#L451-L463。在 MJLab 的 rank-specific seed 下，fresh run 的各 rank target 因而不同；同步后的 predictor gradient 实际上是针对不同 random target 的梯度平均。

RND 的 state normalizer 和 intrinsic-reward normalizer 也只看本地数据，没有 collective：rnd.py#L126-L159、rnd.py#L184-L189。当前 MJLab 的 typed PPO config 没有暴露 RND 配置，因此这主要是 RSL-RL 自身的 distributed edge case。

MJLab curriculum

环境 reset 时直接调用本进程的 curriculum_manager.compute()，manager 也只是本地执行各 curriculum term：manager_based_rl_env.py#L580-L606、curriculum_manager.py#L95-L117。没有 curriculum metric all-reduce 或 rank-0 broadcast，所以不同 rollout 可能让各 rank curriculum 状态逐渐分叉。

日志

只有 rank 0 创建 writer；其他 rank 的 logger 被禁用：logger.py#L51-L78。loss、reward、episode extras 没有全局 reduce，展示的是 rank 0 本地视角；只有 collection size 和 FPS 用 world_size 乘回全局口径：logger.py#L175-L226。

5. Checkpoint / resume 的一致性

RSL-RL 只在有 writer 的 rank 0 定期保存。PPO checkpoint 包含 actor、critic、optimizer，以及启用时的完整 RND state 和 RND optimizer：on_policy_runner.py#L127-L161、ppo.py#L360-L395。

MJLab 的 runner 另外保存并恢复环境 common_step_counter，用来延续依赖全局步数的 curriculum；它没有保存所有环境、episode、RNG、rollout buffer 或每个 rank 的完整 curriculum state：runner.py#L67-L80、runner.py#L83-L140。

resume 时 MJLab 的每个 worker 都执行 runner.load()，然后进入 learn() 再由 rank 0 广播 actor/critic model state：train.py#L167-L177。只要所有 rank 解析到同一 checkpoint，model/optimizer 会从同一份 rank-0 快照开始；随后 observation normalizer、BatchNorm、curriculum 等 rank-local state 又可能继续分叉。

6. Distillation 的多卡实现

RSL-RL 的第二个内置算法是 Student-Teacher Distillation。它复用同一 runner 的 rollout：student 在各 rank 本地行动，teacher 给出监督 action，数据写入各 rank 自己的 RolloutStorage。这是一种在线收集监督样本的 behavior cloning，不是 SAC 意义上的 off-policy RL，也没有 replay buffer。

多卡语义是：

对象	机制	周期
student + teacher model state	rank 0 broadcast	learn() 开始时一次
teacher gradient	无，teacher 只提供 target	不更新
student gradient	flat all_reduce(SUM) / W	每累计 gradient_length 个 timestep batch 一次
student normalizer	rank-local update	每个 env step，无 collective

一次 distillation storage.generator() 按时间步 yield batch，每个 batch 包含本 rank 的全部 N 个环境：rollout_storage.py#L211-L222。算法将连续 L=gradient_length 个 timestep loss 累加后 backward、同步 student gradient、clip、step：distillation.py#L124-L167。因此每次 student optimizer step 覆盖 W * N * L 个样本；需要注意，代码对这 L 个 timestep loss 求和而不是再除以 L。每 iteration 的同步/step 次数是：

floor(num_learning_epochs * T / gradient_length)

计数器跨 learning epoch 累加；如果 num_learning_epochs * T 不能整除 gradient_length，最后不足 L 个 timestep 的 loss 不会 backward，storage 随后被清空。默认是 num_learning_epochs=1, gradient_length=15，但 T 由 runner 配置提供：distillation.py#L31-L44。

startup 广播 student/teacher 完整 state；稳态 all-reduce 只包含 student gradient：distillation.py#L277-L306。

7. SAC / off-policy RL 是否存在

当前 production package 只 export PPO 和 Distillation：algorithms/__init__.py#L6-L11。官方 overview 同样只列这两个算法，并明确把 Distillation 定义为 student-teacher behavior cloning：overview.rst#L22-L35。仓库没有 SAC、TD3、DQN、replay buffer 或 off-policy runner/storage。

RSL-RL 确实允许通过配置解析外部自定义 class，但那只是扩展点，不等于内置 SAC 支持；当前 runner loop、storage 和 checkpoint contract 都是 PPO/Distillation 语义。若要加入 SAC，需要另行定义 replay 所有权、跨 rank 采样、actor/critic/temperature/target 的同步周期和 checkpoint 边界，不能直接把现有 PPO 的 reduce_parameters() 当成完整 off-policy distributed implementation。

MJLab 当前 typed runner 配置也只构造 RslRlOnPolicyRunnerCfg + PPO，没有 SAC 训练入口：mjlab rl/config.py#L128-L145。

8. 与 Holosoma FastSAC 的关键差异及对 UniLab 的启示

和已有的 Holosoma FastSAC 多卡调研 #975 对照：

维度	MJLab + RSL-RL PPO	Holosoma FastSAC
环境数语义	每 rank 完整 num_envs，全局随 W 放大	配置为全局环境数，再除到各 rank
数据容器	rank-local、每 iteration 清空的 rollout	rank-local replay shard
数据通信	不交换 rollout	不交换 replay transition
主同步	每个 PPO mini-batch 平均梯度	每个实际 actor/critic/alpha update 平均梯度
参数广播	startup 一次	startup 一次
observation normalizer	rank-local，不同步	用 collective 更新全局统计
curriculum	rank-local	rank 0 状态广播
算法性质	on-policy PPO	off-policy SAC

两者共同点是 replicated model + rank-local data + synchronous gradient averaging，都不是共享数据池、周期参数平均或 parameter server。但不能只因为通信原语相同就直接复用 execution path：PPO 每轮固定 rollout 后多 epoch 更新；SAC 则长期维护 replay、target network、temperature，并可能每个 env step 做多次不同频率的 optimizer update。

如果 UniLab 后续单独规划多 learner / 多 GPU，至少需要先明确：

1. num_envs 是全局预算还是每 rank 预算；
2. rollout/replay 是 rank-local shard、中心化，还是允许跨 rank 采样；
3. advantage/observation normalizer、BatchNorm、curriculum 是否要求全局一致；
4. 每类 optimizer、target network、temperature 的同步周期；
5. checkpoint 是否只代表 rank 0，以及日志要显示 rank-0 视角还是全局聚合。

一句话定性：MJLab 负责一进程一卡、独立仿真与 seed/device 隔离；RSL-RL 负责 startup 模型状态广播和每个优化 mini-batch 的同步梯度平均。可训练权重同步，但数据、归一化统计、curriculum 和日志并不天然全局同步。当前 RSL-RL 没有 SAC。