VLA Robot — Vision-Language-Action Model

A compact, trainable Vision-Language-Action (VLA) model designed to run on consumer GPUs (RTX 5080 / 16 GB VRAM). Uses a SigLIP/ViT vision encoder, a small causal LLM backbone (Phi-3 Mini or Qwen2.5-3B), and an MLP action head. Simulation is handled via MuJoCo + gymnasium-robotics (FetchPickAndPlace) or LIBERO.

Camera Image  →  ViT Encoder  ──┐
                                  ├──→  LLM Backbone (LoRA)  →  Action Head  →  [dx,dy,dz,gripper]
Task String   →  Tokenizer    ──┘

---

Table of Contents

1. Requirements
2. Installation
3. Project Structure
4. Quick Start
5. Collecting Demonstrations
6. Training
7. Running in MuJoCo Simulation
8. Configuration
9. Extending the Codebase

---

Requirements

Component	Minimum
GPU	NVIDIA RTX 5080 (16 GB) or better
CUDA	12.1+
Python	3.10+
RAM	32 GB system RAM
Disk	~20 GB (models + datasets)

---

Installation

1. Clone and create environment

git clone https://github.com/yourname/vla_robot.git
cd vla_robot

python -m venv .venv
source .venv/bin/activate        # Windows: .venv\Scripts\activate

2. Install dependencies

pip install --upgrade pip
pip install -r requirements.txt

3. Install MuJoCo

MuJoCo 3.x is bundled with the mujoco Python package — no separate binary needed:

pip install mujoco

Test it:

python -c "import mujoco; print(mujoco.__version__)"

4. Download a backbone (first run auto-downloads via HuggingFace)

The default backbone is microsoft/phi-3-mini-4k-instruct (~7 GB). Set your HuggingFace token if needed:

export HF_TOKEN=hf_your_token_here

---

Project Structure

vla_robot/
├── configs/
│   ├── default.yaml          # Main hyperparameter config
│   └── small_debug.yaml      # Tiny config for smoke-testing
├── data/
│   ├── demo_collector.py     # Scripted policy → HDF5 demos
│   ├── dataset.py            # PyTorch Dataset for (img, text, action)
│   └── augmentations.py      # Image augmentation helpers
├── models/
│   ├── vision_encoder.py     # ViT / SigLIP wrapper
│   ├── action_head.py        # MLP action decoder
│   ├── vla_model.py          # Full VLA: encoder + LLM + head
│   └── lora_utils.py         # LoRA injection helpers
├── training/
│   ├── trainer.py            # Training loop (BC / supervised)
│   ├── losses.py             # Action regression + auxiliary losses
│   └── callbacks.py          # Logging, checkpointing helpers
├── simulation/
│   ├── env_wrapper.py        # Gymnasium env → VLA-compatible wrapper
│   ├── evaluator.py          # Roll out policy in sim, log success rate
│   └── visualizer.py        # Live render + action overlay
├── scripts/
│   ├── collect_demos.py      # CLI: collect scripted demonstrations
│   ├── train.py              # CLI: launch training
│   ├── evaluate.py           # CLI: evaluate checkpoint in sim
│   └── interactive.py        # CLI: interactive inference in sim
├── tests/
│   ├── test_model.py
│   ├── test_dataset.py
│   └── test_env.py
├── docs/
│   └── architecture.md
├── requirements.txt
└── README.md

---

Quick Start

Smoke test (no GPU required for shape checks)

python -m tests.test_model

Full pipeline in 3 commands

# 1. Collect 200 scripted demonstrations
python scripts/collect_demos.py --env FetchPickAndPlace-v3 --n_demos 200 --out data/demos.hdf5

# 2. Train for 50 epochs
python scripts/train.py --config configs/default.yaml --demos data/demos.hdf5

# 3. Evaluate the best checkpoint
python scripts/evaluate.py --checkpoint checkpoints/best.pt --n_episodes 20

---

Collecting Demonstrations

collect_demos.py runs a scripted oracle policy (included for FetchPickAndPlace) and saves episodes as HDF5:

python scripts/collect_demos.py \
    --env FetchPickAndPlace-v3 \
    --n_demos 500 \
    --out data/demos.hdf5 \
    --render          # optional: show window while collecting

Each episode stores: rgb_obs, instruction, actions, rewards, dones.

---

Training

python scripts/train.py \
    --config configs/default.yaml \
    --demos data/demos.hdf5 \
    --run_name my_first_vla

Key flags:

Flag	Default	Description
--config	configs/default.yaml	YAML config path
--demos	required	Path to HDF5 demo file
--run_name	vla_run	WandB / checkpoint prefix
--resume	None	Resume from checkpoint path
--fp16	True	Mixed precision training
--lora_rank	16	LoRA rank (lower = less VRAM)

Training logs to runs/<run_name>/ and saves checkpoints to checkpoints/.

---

Running in MuJoCo Simulation

Headless evaluation

python scripts/evaluate.py \
    --checkpoint checkpoints/best.pt \
    --env FetchPickAndPlace-v3 \
    --n_episodes 50

Interactive mode (live render)

python scripts/interactive.py \
    --checkpoint checkpoints/best.pt \
    --instruction "pick up the object and place it at the goal"

Press Q to quit, R to reset episode.

---

Configuration

configs/default.yaml controls everything:

model:
  vision_encoder: "google/siglip-base-patch16-224"
  llm_backbone: "microsoft/phi-3-mini-4k-instruct"
  action_dim: 4          # [dx, dy, dz, gripper]
  lora_rank: 16
  lora_alpha: 32

training:
  batch_size: 16
  lr: 2e-4
  epochs: 100
  warmup_steps: 500
  grad_clip: 1.0

data:
  image_size: 224
  history_len: 1         # number of stacked frames

simulation:
  env_id: "FetchPickAndPlace-v3"
  max_episode_steps: 50
  camera: "external_camera_0"

---

Extending the Codebase

Swap in a different LLM

Edit configs/default.yaml:

model:
  llm_backbone: "Qwen/Qwen2.5-3B-Instruct"

Add a diffusion action head

Replace models/action_head.py with a DDPM head and update the loss in training/losses.py to use denoising score matching.

Use LIBERO environments

pip install libero

Then set simulation.env_id: "LIBERO_SPATIAL" in your config and implement simulation/libero_wrapper.py mirroring env_wrapper.py.