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

RLinf 在 Isaac Lab 任务上通过强化学习微调预训练视觉—语言—动作(VLA)模型,例如 GR00T 和 OpenVLA。

典型工作流分三个阶段:

  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 模型

  • 通过 FSDP 跨 GPU 和节点分布式训练

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

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

先决条件#

重要

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

安装#

对于 GR00T N1.5,从 Isaac Lab 根目录运行以下命令。N1.7 请使用下方的另一组模型依赖。

# 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 shared dependencies and RLinf for both GR00T generations
# 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
uv pip install --no-deps \
    "git+https://github.com/RLinf/RLinf.git@0f9ea98c7a6d9e3ade24e8f4846c64d3b135dbcc"

# Step 2: Install packages with conflicting constraints (--no-deps to bypass resolver)
uv pip install 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 N1.5 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

已有 RLinf 0.2 安装必须重新执行步骤 1,改用上面指定的共用 RLinf 修订版。两代 GR00T 均使用其中各自对应的模型包和动作注册表。

上述安装的包版本有意与 Isaac Lab 锁文件不同。后续命令请使用 uv run --no-sync,避免 uv 在启动前替换这些与 GR00T 兼容的版本。

N1.5 不使用 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。后面的训练和评估命令无需改变。补丁不适用于 N1.7。

GR00T N1.7#

请使用单独的 Python 3.12 Isaac Lab 环境:N1.5 和 N1.7 安装的是同名 gr00t 包,但 API 不兼容。完成上述步骤 1 后,用以下内容替换步骤 2–4:

uv pip install --no-deps \
    "git+https://github.com/NVIDIA/Isaac-GR00T.git@1a1837f20538b7d7e21f977a11a5aee14f99803c" \
    "transformers==4.57.3" "tokenizers>=0.22,<0.23" "huggingface-hub>=0.34,<1.0"
uv pip install --no-deps --no-build-isolation "flash-attn==2.8.3"

在任务 YAML 中设置 actor.model.model_type: gr00t_n1d7。RLinf 从 model_path 加载模型和处理器,data_config_class 仅用于 N1.5。检查点根目录或其 processor/ 子目录中,必须包含 processor_config.json、statistics.json 和 embodiment_id.json。

显式匹配检查点的语言键和环境的动作顺序。例如,以下条目应放在 env.train.isaaclab 下:

gr00t_mapping:
  language: annotation.human.task_description
  # Keep the checkpoint's video and state mappings here as well.
action_mapping:
  keys: [action.left_arm, action.right_arm, action.left_hand, action.right_hand]

未提供这些条目时,语言键仍为 annotation.human.action.task_description,动作保持模型的字典顺序。设置后,keys 必须按环境顺序列出所有需要输出的动作部分。使用 uv run --no-sync 保留固定依赖版本。

aarch64 上的 OpenMP 预加载#

仅在 DGX Spark 及其他 aarch64 Linux 系统上,预加载 aarch64 OpenMP 库,以便将其载入 Python 进程(见 使用 uv 自动安装(推荐)):

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

评估 :评估预训练基础模型,并录制视频:

uv run --no-sync isaaclab 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

这里,--model_path 指向 HuggingFace 格式的基础模型(包含 config.json);--checkpoint 指向 RLinf 检查点目录(global_step_<N> 文件夹)、仅包含一个 full_weights.pt 的子目录,或该权重文件本身。RLinf 先加载基础模型,再通过其原生检查点钩子应用强化学习微调后的权重。模型和检查点的相对路径会在 Ray 工作进程启动前,基于启动器的工作目录解析。rollout 模型继承 actor 模型设置,显式提供的 rollout 设置优先。RLinf 的训练和评估均依赖 Ray 调度工作进程,单 GPU 运行也如此。

备注

--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

跨重置的动作块#

Isaac Lab 在 step() 内重置已结束的环境。对于绝对关节位置策略,设置 env.train.isaaclab.hold_pose_on_midchunk_reset: true,可在下一个动作块到来前,用返回的关节位置替换旧回合剩余的动作。配置的 states 向量必须仅包含这些关节位置,顺序和单位与环境动作一致。该选项默认关闭,不适用于增量位置、速度或力矩动作。

架构#

┌────────────────────────────────────────────────────────────────┐
│                         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 模型

关键文件#

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

详细配置选项、CLI 参数及添加新任务的方法,请通过统一的 uv run isaaclab train --rl_library rlinf 和 uv run isaaclab play --rl_library rlinf 命令查看。