RL Post-Training for VLA Models

RLinf is a flexible and scalable open-source RL infrastructure designed for Embodied and Agentic AI. This integration enables reinforcement learning fine-tuning of Vision-Language-Action (VLA) models (e.g., GR00T, OpenVLA) on Isaac Lab simulation tasks.

The typical workflow follows three stages:

1. Data collection — Collect demonstration data from the Isaac Lab environment (e.g., via teleoperation or scripted policy).
2. Base model training — Train a VLA base model (e.g., GR00T) on the collected demonstrations using supervised learning.
3. RL fine-tuning — Fine-tune the pretrained VLA model on the Isaac Lab task using RLinf with PPO / Actor-Critic / SAC.

Overview

The RLinf integration allows Isaac Lab users to:

• Fine-tune pretrained VLA models on Isaac Lab tasks using PPO / Actor-Critic / SAC
• Leverage RLinf’s FSDP-based distributed training across multiple GPUs/nodes
• Define observation/action mappings from Isaac Lab to GR00T format via a single YAML config
• Register Isaac Lab tasks into RLinf without modifying RLinf source code

Architecture

┌────────────────────────────────────────────────────────────────┐
│                         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       │
└────────────────┴──────────────────────┴────────────────────────┘

Data flow:

1. EnvWorker runs Isaac Lab simulation and converts observations to RLinf format
2. RolloutWorker runs VLA model inference (e.g., GR00T) to produce actions
3. Actions are converted back to Isaac Lab format and stepped in the environment
4. ActorWorker updates the VLA model with PPO/actor-critic loss via FSDP

Prerequisites

• Isaac Lab installed and configured
• Isaac-GR00T repo (for VLA inference and data transforms)
• A pretrained VLA checkpoint in HuggingFace format. A pretrained GR00T checkpoint for assemble_trocar is available and can be downloaded via:

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

• Multi-GPU setup recommended (FSDP requires at least 1 GPU)

Installation

From the Isaac Lab root directory:

# 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

The packages installed in Step 2 intentionally differ from the versions in the Isaac Lab lockfile. Use uv run --no-sync for the commands below so that uv does not replace these GR00T-compatible versions before launching.

Skipping flash-attn

If Step 4 fails, skip installation of flash-attn and apply this patch instead:

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

The patch switches GR00T to PyTorch SDPA, so flash-attn is no longer required. The training and evaluation commands below work unchanged.

OpenMP preload on aarch64

On DGX Spark and other aarch64 Linux systems only, preload the aarch64 OpenMP library so it can be loaded into the Python process (see Python environment with Isaac Sim):

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

Do not set this aarch64 path on x86_64 Linux. If it was inherited from a previous setup, run unset LD_PRELOAD before launching Isaac Lab.

Quick Start

Training — RL fine-tuning of a pretrained VLA model:

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

Evaluation — Evaluate a pretrained (base) model with video recording:

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

Evaluation — Evaluate an RL-finetuned checkpoint with video recording:

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

Here --model_path points to the HuggingFace-format base model (with config.json), and --checkpoint points to the RLinf checkpoint directory (the global_step_<N> folder). The script loads the model architecture from the base model and overlays the RL-finetuned weights (full_weights.pt) from the checkpoint.

Checkpoints

Checkpoints are saved every save_interval epochs (default: 2) to:

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

The placeholders are configurable in the task YAML (source/isaaclab_tasks/isaaclab_tasks/contrib/assemble_trocar/config/isaaclab_ppo_gr00t_assemble_trocar.yaml):

• <experiment_name> — runner.logger.experiment_name (default: test_gr00t)
• <N> — increments every runner.save_interval epochs

The exact path is printed at startup as [INFO] Logging to: .... To resume training, pass the global_step_<N> directory via --checkpoint. For playback, --checkpoint also accepts latest and best; both select the newest saved RLinf checkpoint.

Configuration

All configuration lives in a single YAML file loaded by Hydra. The key configuration block is the env.train.isaaclab section, which defines how Isaac Lab observations are converted to GR00T format:

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

Key Files

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

For detailed configuration options, CLI arguments, and how to add new tasks, use the unified ./isaaclab.sh train --rl_library rlinf and ./isaaclab.sh play --rl_library rlinf commands.