NVIDIA Isaac Lab是一個適用于機器人學(xué)習(xí)的開源統(tǒng)一框架，基于 NVIDIA Isaac Sim 開發(fā)，其模塊化高保真仿真適用于各種訓(xùn)練環(huán)境，可提供各種物理 AI 功能和由 GPU 驅(qū)動的物理仿真，縮小仿真與現(xiàn)實世界之間的差距。

Isaac Lab 主要通過 NVIDIA GPU 加速高度逼真的物理模擬、實時渲染、深度強化學(xué)習(xí)（RL）模仿學(xué)習(xí)（IL）集成等先進(jìn)技術(shù)，然而對于復(fù)雜的強化學(xué)習(xí)環(huán)境，可能需要在多個 GPU 上擴大訓(xùn)練規(guī)模。Isaac Lab 支持多 GPU 和多節(jié)點功能，與在單個 GPU 上進(jìn)行訓(xùn)練相比，可以更快地加速訓(xùn)練過程并達(dá)到更高的性能水平。

在之前的機器人仿真教程里，我們分別介紹了Isaac Lab 的安裝教程以及Isaac Lab 的可用環(huán)境與強化學(xué)習(xí)腳本使用指南，本篇教程將帶大家了解如何通過多 GPU 和多節(jié)點進(jìn)行擴展訓(xùn)練。

核心要點

Isaac Lab 支持多 GPU 和多節(jié)點強化學(xué)習(xí)。但是此功能目前僅可用于 RL-Games、RSL-RL 和 skrl 庫。多 GPU 和多節(jié)點訓(xùn)練僅支持 Linux 系統(tǒng)，因 NCCL 庫限制，暫不支持 Windows 系統(tǒng)。

一、多 GPU

Isaac Lab 支持以下兩種多 GPU 訓(xùn)練框架。

1. 通過 PyTorch Torchrun 進(jìn)行分布式訓(xùn)練

Torchrun 通過以下方式管理分布式訓(xùn)練：

進(jìn)程管理：為每個 GPU 創(chuàng)建一個獨立進(jìn)程，并將每個進(jìn)程分配至指定的 GPU。

腳本執(zhí)行：在每個進(jìn)程上運行相同的訓(xùn)練腳本（例如 RL Games）。

環(huán)境實例：每個進(jìn)程都會創(chuàng)建獨立的 Isaac Lab 環(huán)境實例。

梯度同步：收集并同步所有進(jìn)程的梯度，在每個訓(xùn)練步驟結(jié)束后將更新后的梯度廣播回各進(jìn)程。

此流程中的關(guān)鍵組件是：

Torchrun：處理進(jìn)程生成、通信和梯度同步。

RL 庫：運行實際訓(xùn)練算法的強化學(xué)習(xí)庫。

Isaac Lab：提供每個進(jìn)程獨立實例化的模擬環(huán)境。

Torchrun 在底層使用 DistributedDataParallel 模塊來管理分布式訓(xùn)練。當(dāng)使用多個 GPU 進(jìn)行訓(xùn)練時，會發(fā)生以下情況：

每個 GPU 運行一個獨立的進(jìn)程

每個進(jìn)程執(zhí)行完整的訓(xùn)練腳本

每個進(jìn)程都獨立擁有：

① Isaac Lab 環(huán)境實例（含 n 個并行環(huán)境）

② 策略網(wǎng)絡(luò)副本

③ 用于收集訓(xùn)練經(jīng)驗的緩沖區(qū)

所有進(jìn)程僅針對梯度更新進(jìn)行同步

2. 通過 JAX 進(jìn)行分布式訓(xùn)練

在使用 JAX 時，我們借助skrl.utils.distributed.jax模塊來實現(xiàn)分布式訓(xùn)練。由于 JAX 等機器學(xué)習(xí)框架通常不會在單個程序調(diào)用中自動啟動多個進(jìn)程，因此 skrl 庫提供了該模塊，用于負(fù)責(zé)進(jìn)程的啟動與管理。

*注：僅 skrl 庫支持 JAX。

二、多節(jié)點

除了將訓(xùn)練規(guī)模擴展到單臺機器上的多個 GPU 之外，還可以跨多個節(jié)點進(jìn)行訓(xùn)練。要跨多個節(jié)點/機器進(jìn)行訓(xùn)練，需要在每個節(jié)點上啟動單獨的進(jìn)程。

運行訓(xùn)練

接下來，我們來演示下多 GPU 多節(jié)點訓(xùn)練操作步驟。本期教程以 NVIDIA RTX 5880 Ada GPU 為底層硬件支持。

一、訓(xùn)練環(huán)境

1. CUDA 安裝：

1.1使用 nvidia-smi 查看當(dāng)前驅(qū)動支持的最高 cuda 版本。如下圖所示，可以看到當(dāng)前顯示最高版本為 CUDA Version：12.8。

圖1

1.2訪問官網(wǎng)（https://developer.nvidia.com/cuda-toolkit-archive），下載并安裝 CUDA 12.8 及以下版本。

根據(jù)提示，選擇對應(yīng)系統(tǒng)版本，獲取 CUDA 工具包安裝程序下載鏈接以及安裝方式。

1.3運行以下命令：

sudosh cuda_12.8.0_570.86.10_linux.run

2. Pytorch安裝

2.1訪問https://pytorch.org/，根據(jù)提示選擇對應(yīng)系統(tǒng)、CUDA版本等下載安裝即可。

2.2如果提示沒有 pip，按照要求 apt 安裝。

2.3執(zhí)行pip3 install torch torchvision。

2.4安裝完成后，查看 pytorch 版本。

二、運行結(jié)果（多節(jié)點訓(xùn)練）

1. 在窗口中再次確認(rèn)環(huán)境變量

exportISAACSIM_PATH="${HOME}/isaacsim"
exportISAACSIM_PYTHON_EXE="${ISAACSIM_PATH}/python.sh"

2. 節(jié)點運行

在節(jié)點 1 運行：

${ISAACSIM_PYTHON_EXE}-m torch.distributed.run --nproc_per_node=1 --nnodes=2 --node_rank=0 --rdzv_id=123 --rdzv_backend=c10d --rdzv_endpoint=192.168.150.161:5555 scripts/reinforcement_learning/rl_games/train.py --task=Isaac-Cartpole-v0 --headless –distributed

在節(jié)點 2 運行：

${ISAACSIM_PYTHON_EXE}-m torch.distributed.run --nproc_per_node=1 --nnodes=2 --node_rank=1 --rdzv_id=123 --rdzv_backend=c10d --rdzv_endpoint=192.168.150.161:5555 scripts/reinforcement_learning/rl_games/train.py --task=Isaac-Cartpole-v0 --headless –distributed

注意，需要調(diào)整的參數(shù)主要是：

--nproc_per_node

每個節(jié)點（機器）上啟動的進(jìn)程數(shù)，通常設(shè)置為該節(jié)點的 GPU 數(shù)量。例如，若單機有 8 塊 GPU，--nproc_per_node=8。

--nnodes

參與訓(xùn)練的物理節(jié)點總數(shù)。例如，--nnodes=4表示使用 4 臺機器組成集群進(jìn)行訓(xùn)練。

--rdzv_endpoint

主節(jié)點的 IP 地址和端口號，格式為host:port。所有節(jié)點通過此端點進(jìn)行通信協(xié)調(diào)，例如--rdzv_endpoint=192.168.1.100:29500。

*如需了解詳細(xì)步驟，可參考官方文檔：

https://isaac-sim.github.io/IsaacLab/main/source/features/multi_gpu.html

推薦硬件配置

以下是麗臺針對個人開發(fā)者/研究人員的機器人仿真及訓(xùn)練環(huán)境推薦配置，能流暢運行 Isaac Lab 以及機器人訓(xùn)練和仿真所需的算力，同時可再與企業(yè)級多卡集群方案形成互補，能夠完全滿足中小規(guī)模的實驗需求。

LEADTEK WS3008

LEADTEK WS3008 產(chǎn)品特性

支持選配 2 張 NVIDIA RTX 5880 Ada GPU

支持單 Intel Xeon W-3400、W-2400 系列處理器

Intel W790 芯片組

支持最高至 350W CPU TDP

支持最多 DDR5-4800MHz x8 內(nèi)存

支持 4 PCIe 5.0 x16 插槽

支持 1 M.2 NVMe PCI-E 4.0 x4

支持 2 個 10GbE BaseT、1 個 2.5GbE BaseT 和 1 個 1GbE BaseT 網(wǎng)口

支持 1 個 IPMI 管理口

支持 2 個 USB-A

支持 1 個 VGA 口，1 個 COM 口

支持 1 個 1200W/1300W 鉑金電源

機箱體積：400.0×278.0×167.6 mm