VLA-Scratch

VLA-Scratch is a modular, performant and efficient stack with minimal dependencies that makes training, evaluating, and serving Vision-Language-Action (VLA) models fast and approachable.

🚀 Quickstart

We use uv to manage dependencies. Run GIT_LFS_SKIP_SMUDGE=1 uv sync to set up the environment.

Verify your installation with the following commands:

# Training
uv run torchrun --standalone --nnodes=1 --nproc_per_node=8 \
    scripts/train_policy.py \
    policy=pi-qwen \
    data=libero-spatial \
    lr.base=5e-5 \
    +lr.vlm_bridge=1e-5 \
    +lr.action_expert=5e-5 \
    wandb.mode=online

# Evaluation
uv run scripts/eval_policy.py \
    checkpoint_path=hf:elijahgalahad/libero-spatial-qwen \
    data=libero-spatial \
    data.video_backend=pyav \
    merge_policy_cfg=true

# Serving
uv run scripts/serve_policy.py \
    checkpoint_path=hf:elijahgalahad/libero-spatial-qwen \
    data=libero-spatial \
    merge_policy_cfg=true

See scripts/README.md for more training commands. See examples/libero for LIBERO evaluation instructions.

✨ Key-Features

• Explicit Data Model for Composable Modules
  • TensorClass is used to define explicit data boundaries between modules, making our codebase fully typed and modular.
  • This allows heterogeneous dataset co-training, and clearer data flow between datasets, policies, and transforms.
• Dedicated Tuning for a First-Class Performance Stack
  • We rewrite the forward pass of VLMs to eliminate all host-device syncs.
  • Instead of relying on generic libaries like accelerate, vla-scratch leverage native torch operations like FSDP2 and gradient checkpointing for dedicated performance tuning.
• Rich Feature Set Out-of-the-Box
  • Multi-source dataset co-training: VQA and robotic datasets co-training.
  • Multi VLM backbone support: Qwen3-VL, PaliGemma 1/2, SmolVLM.
  • Simulation-ready checkpoints and serving scripts.
• Clarity-Focused Hydra Workflow for Seamless Experimentation
  • Every policy and data config is registered with Hydra, so experiments are overrideable with minimal boilerplate.
  • Training, eval, and serving scripts share a common config grammar, so switching between workflows is seamless.

🗂️ Codebase Structure

🧭 Developer Guide

2) Add implementation and configs (e.g. dataset.py, config.py).

3) Register the dataset config with Hydra ConfigStore in config.py under the data group. | | Add a new VLM backbone | 1) Create a new folder under vla_scratch/policies/modules/vlm_bridge/.

2) Implement preprocessing in processor.py, define the policy_input contract between preprocessing and encoding, and implement the main encoding in bridge.py.

3) Wire the backbone into vla_scratch/policies/pi/policy.py. | | Add a new loss term | 1) Add the loss inside compute_loss in vla_scratch/policies/pi/policy.py

2) Return it in log_dict so it is logged to Weights & Biases. | | Format code | Run uvx ruff format. |

🛠️ Troubleshooting (To be Continued)

RTX 5090 Compatibility

Issue
PyTorch stable has not yet fully supported CUDA 12.8 for RTX 5090
(Blackwell-generation GPU, compute capability sm_120).
This may lead to kernel launch failures or CUDA unavailable errors.

Solution
Use PyTorch-Nightly instead of stable PyTorch.

Add the following to your <pyproject.toml>:

```
[[tool.uv.index]]
name = "pytorch-nightly-cu128"
url = "https://download.pytorch.org/whl/nightly/cu128"
default = true

[[tool.uv.index]]
url = "https://pypi.org/simple"

[tool.uv.sources]
torch = { index = "pytorch-nightly-cu128" }
torchvision = { index = "pytorch-nightly-cu128" }
triton = { index = "pytorch-nightly-cu128" }

```