README.md

---
license: apache-2.0
library_name: pytorch
tags:
  - robotics
  - VLA
  - vision-language-action
  - robot-manipulation
  - bridgedata
  - Qwen3-VL
  - imitation-learning
datasets:
  - bridgedata-v2
language:
  - en
metrics:
  - mse
pipeline_tag: robotics
---

# FrozenVLA-2B — Lightweight Vision-Language-Action for Robot Manipulation

A 50M-parameter trainable head on top of frozen **Qwen3-VL-2B-Instruct**, trained on **BridgeData v2** (2,617 episodes) for general-purpose robotic manipulation.

---

## 🧠 Architecture

```
┌──────────────────────────────────────────┐
│         Qwen3-VL-2B-Instruct (FROZEN)    │
│  ┌─────────────┐    ┌─────────────────┐  │
│  │ Vision Enc. │ -> │   LLM (28L)     │  │
│  └─────────────┘    │  hidden=2048    │  │
│                     └────────┬────────┘  │
│                              │           │
│                     last_token_hidden    │
│                         (2048-dim)       │
└──────────────────────────────┼───────────┘
                               │
                    ┌──────────▼──────────┐
                    │   MLP Projector     │  ← Trainable (1.3M)
                    │  2048 → 512 → 512   │
                    │  GELU + Dropout     │
                    └──────────┬──────────┘
                               │
                    ┌──────────▼──────────┐
                    │   Action Head       │  ← Trainable
                    │    512 → 7          │
                    └──────────┬──────────┘
                               │
                    EEF delta: [dx, dy, dz, ax, ay, az, gripper]
```

| Component | Params | Status |
|-----------|--------|--------|
| Qwen3-VL-2B (Vision + LLM) | 2,127,532,032 | ❄️ Frozen |
| MLP Projector | 1,314,824 | 🔥 Trained |
| Action Head | 513 | 🔥 Trained |
| **Total** | **2,128,847,369** | 1.3M trainable |

---

## 📊 Training

| Item | Detail |
|------|--------|
| **Dataset** | BridgeData v2 (LeRobot format) |
| **Episodes** | 2,617 |
| **Action format** | EEF delta 7D: [dx, dy, dz, ax, ay, az, gripper] |
| **Action normalization** | Z-score (mean/std per dimension) |
| **Hardware** | Alibaba Cloud PAI · NVIDIA A10 24GB |
| **Framework** | PyTorch 2.6.0 · Transformers 4.51+ · CUDA 12.6 |
| **Epochs** | 4 (of 5 planned) |
| **Effective batch size** | 64 (32 × gradient_accumulation=2) |
| **Optimizer** | AdamW (lr=1e-4, wd=1e-2) |
| **Schedule** | Cosine annealing · 500 warmup steps |
| **Precision** | bfloat16 · gradient clip=1.0 |
| **Image size** | 448×448 (Qwen3-VL default) |
| **Frames** | All frames per episode (frame_sampling=all) |

### Training Loss (per epoch)

| Epoch | Approx MSE Loss | Checkpoint |
|-------|----------------|------------|
| 1 | — | `epoch_1.pt` (15MB) |
| 2 | — | `epoch_2.pt` (15MB) |
| 3 | — | `epoch_3.pt` (15MB) |
| 4 | — | `epoch_4.pt` (15MB) |

### Inference Requirements

| Platform | VRAM | Notes |
|----------|------|-------|
| RTX 4060 Laptop 8GB | ~4 GB | ✅ Verified · sdpa · bf16 |
| A10 24GB | ~4 GB | ✅ Training env |
| T4 16GB | ~4 GB | ✅ Should work |

---

## 🚀 Quick Start

### 1. Clone & Install

```bash
# Requirements
pip install torch>=2.5.0 transformers>=4.51.0 accelerate sentencepiece protobuf Pillow

# Clone this repo
git clone https://huggingface.co/YOUR_USERNAME/frozenvla
cd frozenvla
```

### 2. Download Base Model

```bash
# Qwen3-VL-2B-Instruct (from HuggingFace)
huggingface-cli download Qwen/Qwen3-VL-2B-Instruct --local-dir ./Qwen3-VL-2B-Instruct
```

### 3. Load & Infer

```python
import torch
from PIL import Image
from model import FrozenVLA

# Load model with trained head
model = FrozenVLA(
    llm_name="./Qwen3-VL-2B-Instruct",
    mlp_hidden_dim=512,
    mlp_depth=2,
    action_dim=7,
    attn_implementation="sdpa",  # "flash_attention_2" on A100/H100
)
model.load_trainable("epoch_4.pt")
model = model.to("cuda").eval()

# Cast head to bf16 (to match Qwen3-VL output dtype)
model.mlp_projector = model.mlp_projector.to(dtype=torch.bfloat16)
model.action_head = model.action_head.to(dtype=torch.bfloat16)

# Inference
image = Image.open("robot_view.jpg").convert("RGB")
instruction = "pick up the red block"

with torch.no_grad():
    action = model([image], [instruction])
    # action: (1, 7) EEF delta [dx, dy, dz, ax, ay, az, gripper]
    print(action.float().cpu().numpy())
```

### 4. Deploy Script (One-Click)

```bash
python deploy.py --checkpoint epoch_4.pt --test-image robot_view.jpg
```

---

## 📁 Files

| File | Description |
|------|-------------|
| `epoch_4.pt` | Best trained head (MLP + ActionHead, 15MB) |
| `epoch_1~3.pt` | Intermediate checkpoints |
| `model.py` | Full model architecture (FrozenVLA class) |
| `config.yaml` | Training configuration |
| `deploy.py` | One-click deployment + inference script |

---

## ⚠️ Limitations

- **Action space**: BridgeData EEF delta only — NOT directly compatible with joint-space robots without IK conversion.
- **Domain**: Trained on BridgeData scenes (tabletop manipulation). Zero-shot generalization to novel environments is limited.
- **Single image input**: Uses the current frame only; no temporal context from video history.
- **Language**: English instructions only.

---

## 📝 Citation

```bibtex
@misc{frozenvla-2026,
  title  = {FrozenVLA-2B: Lightweight VLA from Frozen Qwen3-VL on BridgeData},
  author = {},
  year   = {2026},
  url    = {https://huggingface.co/YOUR_USERNAME/frozenvla}
}
```

## 📄 License

Apache 2.0