---

Overview

Nemotron-VLA is a fully-functional Vision-Language-Action model that uses NVIDIA's foundation models as its backbone. It demonstrates how to build a VLA that takes camera images + text instructions → robot actions, using state-of-the-art NVIDIA models for vision and language understanding.

Architecture

   [Camera Image]               [Text Instruction]              [Robot State]
         │                              │                             │
         ▼                              ▼                             ▼
┌──────────────────┐           ┌──────────────────┐          ┌──────────────────┐
│   NVIDIA RADIO   │           │  Nemotron Nano   │          │  State Encoder   │
│  (Vision ViT)    │           │   (Text LLM)     │          │      (MLP)       │
└────────┬─────────┘           └────────┬─────────┘          └────────┬─────────┘
         │ (Frozen)                     │ (Frozen)                    │ (Trainable)
         │                              │                             │
         ▼                              ▼                             ▼
  ┌────────────┐                 ┌────────────┐                     │
  │ Vision Proj│                 │ Text Proj  │                     │
  └──────┬─────┘                 └──────┬─────┘                     │
         │                              │                             │
         └──────────────┬───────────────┘                             │
                        │                                             │
                        ▼                                             ▼
              ┌────────────────────┐                        ┌────────────┐
              │   Cross-Attention  │◄───────────────────────┤ LayerNorm  │
              │       Fusion       │                        └────────────┘
              └─────────┬──────────┘
                        │ (Fused Context)
                        ▼
              ┌────────────────────┐
              │  Diffusion Policy  │
              │       Head         │
              └─────────┬──────────┘
                        │
                        ▼
                 [Robot Action]

Models Used

Getting Started

Option 1: Google Colab (Recommended)

1. Upload nemotron_vla.ipynb to Google Colab
2. Set runtime to GPU → A100 (Colab Pro required)
3. Run all cells sequentially

The notebook will:

• Install all dependencies
• Write helper modules (env.py, models.py, utils.py)
• Collect expert demonstrations from MetaWorld
• Precompute embeddings from NVIDIA models
• Train the VLA
• Evaluate and save videos

Option 2: Local Setup

# Clone and setup
git clone <repo-url>
cd nemotron-vla

# Install dependencies
pip install torch torchvision transformers accelerate
pip install gymnasium metaworld mujoco
pip install causal-conv1d mamba-ssm
pip install imageio[ffmpeg] matplotlib

Project Structure

nemotron-vla/
├── nemotron_vla.ipynb    # Main Colab notebook (run this!)
├── env.py                # MetaWorld environment wrapper
├── models.py             # All model components
│   ├── RADIO loading     #   NVIDIA RADIO vision encoder
│   ├── Nemotron loading  #   Nemotron Nano 9B text encoder
│   ├── StateEncoder      #   Robot state MLP
│   ├── CrossAttentionFusion  # Multi-modal fusion
│   ├── DiffusionPolicyHead   # DDPM action generation
│   └── NemotronVLA       #   Full VLA model
├── utils.py              # Dataset, training, evaluation
└── README.md             # This file

Memory Management Strategy

The key insight enabling training on A100 40GB:

1. Load RADIO → extract vision features for all images → unload (~1.4GB freed)
2. Load Nemotron 9B → extract text embedding → unload (~18GB freed)
3. Train only lightweight fusion + diffusion (~0.8M params, <4GB)

Key Differences from mini-VLA

License

This project uses NVIDIA models under their respective licenses:

• NVIDIA RADIO: License
• NVIDIA Nemotron: NVIDIA Open Model License