VLA 模型的强化学习后训练

VLA 模型的强化学习后训练#

RLinf 是面向具身智能和智能体 AI 的灵活、可扩展开源强化学习基础设施。通过这一集成，可以在 Isaac Lab 仿真任务上，对 GR00T、OpenVLA 等 视觉-语言-动作（VLA）模型进行强化学习微调 。

典型工作流分三个阶段：

  1. 数据采集 ：从 Isaac Lab 环境采集示范数据，例如使用遥操作或脚本化策略。

  2. 基础模型训练 ：用采集的示范进行监督学习，训练 GR00T 等 VLA 基础模型。

  3. 强化学习微调 ：通过 RLinf 的 PPO / Actor-Critic / SAC，在 Isaac Lab 任务上微调预训练 VLA 模型。

概述#

RLinf 集成使 Isaac Lab 用户能够：

  • 通过 PPO / Actor-Critic / SAC，在 Isaac Lab 任务上微调预训练 VLA 模型

  • 使用 RLinf 基于 FSDP 的多 GPU、多节点分布式训练

  • 通过单个 YAML 配置，定义 Isaac Lab 到 GR00T 格式的观测与动作映射

  • 无需修改 RLinf 源码，即可将 Isaac Lab 任务注册到 RLinf

架构#

┌────────────────────────────────────────────────────────────────┐
│                         RLinf Runner                           │
│                 (EmbodiedRunner / EvalRunner)                  │
├────────────────┬──────────────────────┬────────────────────────┤
│  Actor Worker  │   Rollout Worker     │      Env Worker        │
│  (FSDP)        │  (HF Inference)      │  (IsaacLab Sim)        │
│                │                      │                        │
│ Policy         │  Multi-step rollout  │ IsaacLabGenericEnv     │
│ Update         │  with VLA model      │  ├─ _make_env_function │
│                │                      │  ├─ _wrap_obs          │
│                │                      │  └─ _wrap_action       │
└────────────────┴──────────────────────┴────────────────────────┘

数据流：

  1. EnvWorker 运行 Isaac Lab 仿真，并将观测转换为 RLinf 格式

  2. RolloutWorker 运行 GR00T 等 VLA 模型推理，生成动作

  3. 动作转换回 Isaac Lab 格式，并在环境中执行一步

  4. ActorWorker 通过 FSDP，使用 PPO 或 actor-critic 损失更新 VLA 模型

先决条件#

重要

RLinf 后训练目前只支持 Linux。兼容发行版包括 Ubuntu、Debian、Red Hat 系发行版和 Arch Linux。

  • 已安装并配置 Isaac Lab

  • Isaac-GR00T 仓库，用于 VLA 推理和数据转换

  • HuggingFace 格式的 预训练 VLA 检查点 。已有用于 assemble_trocar 的预训练 GR00T 检查点，可通过以下方式下载：

    hf download --repo-type model nvidia/Assemble_Trocar --local-dir /path/to/local/models

  • 建议使用多 GPU 配置；FSDP 至少需要一块 GPU

安装#

从 Isaac Lab 根目录运行：

# If running Isaac Sim headless for the first time, accept the EULA via env var
# (interactive sessions prompt automatically; headless mode requires this)
export OMNI_KIT_ACCEPT_EULA=yes

# Step 1: Install safe dependencies via the rlinf and video extras
# NOTE: On DGX Spark / aarch64 systems, build decord from source first
# (see "Building decord on DGX Spark / aarch64" below), then run this step.
# --inexact keeps the existing environment (e.g. Isaac Sim) untouched while
# adding the rlinf and video dependencies from the root pyproject.
uv sync --inexact --extra rlinf --extra video

# Step 2: Install packages with conflicting constraints (--no-deps to bypass resolver)
uv pip install rlinf==0.2.0dev2 transformers==4.51.3 \
    "tokenizers>=0.21,<0.22" "huggingface-hub>=0.30,<1.0" --no-deps
# Use the official PyTorch3D v0.7.9 tag instead of the older pipablepytorch3d package.
# GR00T only uses pytorch3d.transforms, so skip the compiled extension.
PYTORCH3D_NO_EXTENSION=1 uv pip install --no-build-isolation \
    "git+https://github.com/facebookresearch/pytorch3d.git@v0.7.9" --no-deps

# Step 3: Install Isaac-GR00T (pinned version)
git clone https://github.com/NVIDIA/Isaac-GR00T.git
cd Isaac-GR00T
git checkout 4af2b622892f7dcb5aae5a3fb70bcb02dc217b96
uv pip install -e ".[base]" --no-deps
cd ../

# Step 4: Install flash-attn (see "Skipping flash-attn" below if this fails)
pip install flash-attn==2.8.3 --no-build-isolation --no-deps

第 2 步安装的包版本有意不同于 Isaac Lab 锁文件。以下命令使用 uv run --no-sync ，避免 uv 在启动前替换这些兼容 GR00T 的版本。

跳过 flash-attn#

如果第 4 步失败，可跳过 flash-attn 安装，改为应用此补丁：

cd Isaac-GR00T
git apply /path/to/IsaacLab/scripts/imitation_learning/locomanipulation_sdg/gr00t/no_flash_attn.patch

该补丁将 GR00T 切换到 PyTorch SDPA，因此不再需要 flash-attn。以下训练和评估命令无需修改。

aarch64 上的 OpenMP 预加载#

只在 DGX Spark 和其他 aarch64 Linux 系统上，预加载 aarch64 OpenMP 库，使其能加载到 Python 进程中（见 包含 Isaac Sim 的 Python 环境 ）：

unset LD_PRELOAD
export LD_PRELOAD=/lib/aarch64-linux-gnu/libgomp.so.1

不要在 x86_64 Linux 上设置此 aarch64 路径。如果继承了旧配置，启动 Isaac Lab 前运行 unset LD_PRELOAD 。

快速开始#

训练 ：对预训练 VLA 模型进行强化学习微调：

uv run --no-sync isaaclab train --rl_library rlinf \
    --config_name isaaclab_ppo_gr00t_assemble_trocar \
    --model_path /path/to/base_model

./isaaclab.sh train --rl_library rlinf \
    --config_name isaaclab_ppo_gr00t_assemble_trocar \
    --model_path /path/to/base_model

评估 ：评估预训练基础模型，并录制视频：

uv run --no-sync isaaclab play --rl_library rlinf \
    --config_name isaaclab_ppo_gr00t_assemble_trocar \
    --model_path /path/to/base_model \
    --video

./isaaclab.sh play --rl_library rlinf \
    --config_name isaaclab_ppo_gr00t_assemble_trocar \
    --model_path /path/to/base_model \
    --video

评估 ：评估强化学习微调后的检查点，并录制视频：

uv run --no-sync isaaclab play --rl_library rlinf \
    --config_name isaaclab_ppo_gr00t_assemble_trocar \
    --model_path /path/to/base_model \
    --checkpoint /path/to/checkpoints/global_step_N \
    --video

./isaaclab.sh play --rl_library rlinf \
    --config_name isaaclab_ppo_gr00t_assemble_trocar \
    --model_path /path/to/base_model \
    --checkpoint /path/to/checkpoints/global_step_N \
    --video

这里 --model_path 指向包含 config.json 的 HuggingFace 格式基础模型， --checkpoint 指向 RLinf 检查点目录，即 global_step_<N> 文件夹。脚本从基础模型加载网络结构，再覆盖检查点中的强化学习微调权重 full_weights.pt 。

备注

--config_path 是可选标志。省略时，脚本自动在 isaaclab_tasks 包中搜索匹配的 YAML 配置文件。

检查点#

检查点每隔 save_interval 个 epoch（默认 2 ）保存到:

logs/rlinf/<timestamp>-IsaacContrib-Assemble-Trocar-G129-Dex3/<experiment_name>/checkpoints/global_step_<N>/

占位符可在任务 YAML（ source/isaaclab_tasks/isaaclab_tasks/contrib/assemble_trocar/config/isaaclab_ppo_gr00t_assemble_trocar.yaml ）中配置：

  • <experiment_name> ： runner.logger.experiment_name （默认 test_gr00t ）

  • <N> ：每隔 runner.save_interval 个 epoch 增加

确切路径在启动时以 [INFO] Logging to: ... 打印。恢复训练时，通过 --checkpoint 传入 global_step_<N> 目录。回放时， --checkpoint 也接受 latest 和 best ，二者都会选择最新保存的 RLinf 检查点。

小技巧

训练吞吐量随并行环境数增加。如果 GPU 尚有空闲显存，可增加任务 YAML 中的 env.train.total_num_envs （默认 4 ）。

小技巧

每个检查点可能有数 GB。为避免占满磁盘，可增加任务 YAML 的 save_interval ，减少训练期间保存的中间检查点。

配置#

全部配置位于 Hydra 加载的 单个 YAML 文件 中。关键配置块是 env.train.isaaclab ，定义如何将 Isaac Lab 观测转换为 GR00T 格式：

isaaclab: &isaaclab_config
  task_description: "assemble trocar from tray"

  # IsaacLab → RLinf observation mapping
  main_images: "front_camera"
  extra_view_images:
    - "left_wrist_camera"
    - "right_wrist_camera"
  states:
    - key: "robot_joint_state"
      slice: [15, 29]
    - key: "robot_dex3_joint_state"

  # GR00T → IsaacLab action conversion
  action_mapping:
    prefix_pad: 15
    suffix_pad: 0

关键文件#

source/isaaclab_rl/isaaclab_rl/entrypoints/backends/
├── train_rlinf.py      # Training entry point
├── play_rlinf.py       # Evaluation entry point
└── cli_args_rlinf.py   # Shared CLI argument definitions

source/isaaclab_contrib/isaaclab_contrib/rl/rlinf/
├── __init__.py
└── extension.py       # Task registration, obs/action conversion

详细配置选项、命令行参数和添加新任务的方法，通过统一的 ./isaaclab.sh train --rl_library rlinf 和 ./isaaclab.sh play --rl_library rlinf 命令查看。

目录