在足式机器人策略训练中，环境并行度与网络规模的增长往往使单卡显存与计算吞吐量成为瓶颈。本页面向具备多 GPU 集群或单机多卡环境的高级开发者，系统阐述如何在 RSL-RL 算法库及 Isaac Lab 仿真环境中启用分布式训练。内容覆盖从入口脚本、进程组初始化、梯度同步到日志归约的完整链路，适用于 On-Policy Runner、AMP Runner 以及 Distillation Runner 三种训练管线。

分布式训练架构概览

项目采用 PyTorch DDP (Distributed Data Parallel) 范式的轻量实现：不依赖 DistributedDataParallel 包装器，而是在算法内部显式调用 torch.distributed 原语完成参数广播与梯度归约。每个 GPU 运行独立的 Isaac Sim 实例与环境集合，通过 NCCL 后端进行低延迟 GPU 间通信。

架构上存在三个层级协同工作：

1. 入口脚本层 (train.py)：解析 --distributed 标志，为各进程分配独立 GPU 与随机种子；
2. Runner 层 (OnPolicyRunner / AMPRunner / DistillationRunner)：读取 WORLD_SIZE / RANK / LOCAL_RANK 环境变量，初始化进程组；
3. 算法层 (PPO / PPOAMP / Distillation)：在 update() 中执行梯度 all_reduce，并在训练起点调用 broadcast_parameters() 确保所有进程从同一参数出发。

flowchart TB
    subgraph Entry["入口脚本层 (train.py)"]
        A1["解析 --distributed"] --> A2["按 local_rank 分配 cuda:x"]
        A2 --> A3["seed += local_rank"]
    end

    subgraph Runner["Runner 层"]
        B1["_configure_multi_gpu()"] --> B2["init_process_group(nccl)"]
        B2 --> B3["构造 multi_gpu_cfg"]
    end

    subgraph Algorithm["算法层"]
        C1["broadcast_parameters()"] --> C2["各卡独立 Rollout"]
        C2 --> C3["backward()"]
        C3 --> C4["reduce_parameters()<br/>all_reduce + average"]
        C4 --> C5["optimizer.step()"]
    end

    A3 --> B1
    B3 --> C1

Sources: train.py, on_policy_runner.py

环境变量与进程组初始化

分布式训练需要由外部启动器（如 torchrun、mpirun 或 Slurm）预先注入标准环境变量。OnPolicyRunner._configure_multi_gpu() 在初始化时读取这些变量，并通过严格的校验逻辑确保配置一致性：

当校验通过后，Runner 调用 torch.distributed.init_process_group(backend="nccl") 建立 GPU 间通信域，并执行 torch.cuda.set_device(self.gpu_local_rank) 将当前进程绑定到对应 GPU。若使用 CPU 设备（--device cpu）与 --distributed 同时传入，train.py 将在启动阶段抛出 ValueError 并终止，因为 NCCL 后端不支持 CPU 设备。

Sources: on_policy_runner.py, train.py

训练启动流程

在分布式模式下，每个进程执行完全一致的 train.py，但操作不同的 GPU 设备。以下是启动到首次参数同步的时序流程：

sequenceDiagram
    participant Launcher as torchrun / mpirun
    participant P0 as Rank 0 (cuda:0)
    participant P1 as Rank 1 (cuda:1)
    participant PN as Rank N (cuda:N)

    Launcher->>P0: WORLD_SIZE=2, RANK=0, LOCAL_RANK=0
    Launcher->>P1: WORLD_SIZE=2, RANK=1, LOCAL_RANK=1
    P0->>P0: env_cfg.sim.device = "cuda:0"
    P1->>P1: env_cfg.sim.device = "cuda:1"
    P0->>P0: seed += 0
    P1->>P1: seed += 1
    P0->>PN: init_process_group(nccl)
    P1->>PN: init_process_group(nccl)
    Note over P0,PN: 各自独立创建环境并 Rollout
    P0->>P1: broadcast_parameters() (rank 0 → all)
    Note over P0,PN: 所有卡参数严格一致后开始训练

每个进程的环境种子会根据 local_rank 进行偏移，从而在不同 GPU 上产生差异化的环境随机化，增加策略探索的多样性。

Sources: train.py, on_policy_runner.py

算法层的梯度同步机制

分布式训练的核心在于保证所有 GPU 上的模型参数在每次更新后保持一致。项目在各算法类内部实现了显式的梯度同步，而非使用 PyTorch 的 DistributedDataParallel 包装器。

参数广播 (broadcast_parameters)

在 OnPolicyRunner.learn() 与 AMPRunner.learn() 的起始阶段，会调用 self.alg.broadcast_parameters()。该方法将 rank 0 上的 policy.state_dict()（以及 RND predictor 的 state_dict，若启用）通过 torch.distributed.broadcast_object_list 发送到所有进程，随后各进程 load_state_dict，确保训练起点完全一致。

Sources: on_policy_runner.py, ppo.py, distillation.py

梯度归约 (reduce_parameters)

在每个 mini-batch 的 loss.backward() 之后，算法调用 reduce_parameters()。其执行逻辑如下：

1. 将 policy（及 RND）的所有可训练参数的 grad 展平并拼接为单一长向量；
2. 调用 torch.distributed.all_reduce(..., op=ReduceOp.SUM) 对梯度求和；
3. 按 gpu_world_size 取平均；
4. 将归约后的梯度写回对应参数的 grad.data。

这一设计使得各进程在独立计算梯度后，最终执行的是全局平均梯度对应的参数更新，从而保证多卡等价于单卡大 batch 训练。

Sources: ppo.py, distillation.py

自适应学习率同步

对于使用 adaptive learning rate schedule 的配置，PPO 与 PPOAMP 会在每次更新前基于 KL 散度调整学习率。在分布式场景下，KL 散度同样通过 all_reduce 求全局平均后，由 rank 0 计算新的学习率，再通过 torch.distributed.broadcast 广播给所有进程，确保各卡学习率严格一致。

Sources: ppo.py, ppo_amp.py

Runner 与算法支持矩阵

以下表格总结了不同训练管线对分布式训练的支持情况，以及各自在梯度同步中覆盖的模型组件：

注意：由于 PPOAMP 的 reduce_parameters() 与 broadcast_parameters() 直接继承自 PPO，AMP Discriminator 的梯度与参数未纳入跨卡同步。若需严格的多卡 AMP 训练一致性，请留意此行为。

Sources: on_policy_runner.py, amp_runner.py, distillation_runner.py

日志与指标统计

在多卡训练中，只有 rank 0 进程执行日志写入与控制台打印，以避免重复记录。Logger 与 LoggerAMP 通过以下逻辑控制：

self.disable_logs = is_distributed and gpu_global_rank != 0

为了正确反映全局训练吞吐，collection_size 的计算会乘以 gpu_world_size：

collection_size = self.cfg["num_steps_per_env"] * self.num_envs * self.gpu_world_size
fps = int(collection_size / (collect_time + learn_time))
self.tot_timesteps += collection_size

这意味着日志中显示的 total_fps 与 Total steps 均为所有 GPU 的汇总值，与单卡训练的统计口径保持一致。

Sources: logger.py, logger.py, amp_logger.py

启动命令示例

假设当前已在 Isaac Lab / RoboLab 环境中完成依赖安装，以下是单机双卡与多机多卡的典型启动方式。

单机多卡 (torchrun)

torchrun --standalone --nnodes=1 --nproc_per_node=2 \
    scripts/rsl_rl/train.py --task=Isaac-Velocity-Flat-Anymal-C-v0 \
    --distributed --headless

• --nproc_per_node：等于单机的 GPU 数量；
• --distributed：告知 train.py 启用多 GPU 配置逻辑；
• 无需手动指定 --device，脚本会自动将各进程绑定到 cuda:{local_rank}。

多机多卡 (torchrun)

# Node 0 (主节点)
torchrun --nnodes=2 --node_rank=0 --nproc_per_node=4 \
    --master_addr=<主节点IP> --master_port=29500 \
    scripts/rsl_rl/train.py --task=Isaac-Velocity-Flat-Anymal-C-v0 \
    --distributed --headless

# Node 1
torchrun --nnodes=2 --node_rank=1 --nproc_per_node=4 \
    --master_addr=<主节点IP> --master_port=29500 \
    scripts/rsl_rl/train.py --task=Isaac-Velocity-Flat-Anymal-C-v0 \
    --distributed --headless

Sources: train.py

配置检查清单与故障排查

在首次启用分布式训练时，建议按以下清单逐项确认：

Sources: train.py, on_policy_runner.py

延伸阅读

分布式训练建立在标准的 On-Policy 训练循环之上，若尚未熟悉单卡 Runner 的执行逻辑，建议先阅读 On-Policy Runner 训练循环。对于需要结合模仿学习的场景，可参考 AMP Runner 与模仿学习训练。完成训练后，模型保存与加载流程与单卡一致，详情见 策略测试与 Sim2Sim 部署 以及 MuJoCo Sim2Sim 部署与真机迁移。