Python + PyTorch: The Dominant Stack for VLA

Detailed content

🔹 Why PyTorch dominates robotics and research

Almost every open VLA — OpenVLA, Octo, π0, ACT — is written in PyTorch. The reason is cultural and technical: execution eager (define-by-run) keeps the graph dynamic, so you inspect tensors in the middle of the forward pass, add breakpoint() inside the policy and iterates in minutes. For manipulation research, where the architecture changes every week, this is invaluable.

📊 Why the ecosystem won

autograd — reverse-mode automatic differentiation, the basis of any policy training.
Hugging Face + PEFT — Ready-to-use VLMs (PaliGemma, Qwen2-VL) with LoRA in a few lines.
torch.compile — kernel fusion via TorchInductor without rewriting code.
CUDA/ROCm/MPS — same code runs in a data center, on Jetson, and on Mac.

⚡ Practical tip

JAX (with MJX/Flax) shines in massively parallel simulation and differentiable physics, but the deploy on a real robot, and most VLA checkpoints run in PyTorch. Learn both; start with PyTorch.

Eager mode

Define-by-run, native debugging.

autograd

Automatic backpropagation.

Ecosystem

HF, PEFT, timm, accelerate.

Portability

Same code, multiple targets.

🔹 Tensors, GPU, and mixed precision in 7B models

A 7B VLA in fp32 takes ~28 GB for weights alone. Train on bf16 cuts it in half while also avoiding fp16 overflow (bf16 has the same exponent as fp32). The mental rule for memory: weights + gradients + optimizer states + activations. With Adam in fp32, the states alone cost ~8 bytes/parameter.

import torch
device = "cuda"
model = model.to(device)

# autocast: forward em bf16, mestre em fp32
scaler = None  # bf16 dispensa GradScaler (fp16 precisa)
for batch in loader:
    obs = batch["observation.images.top"].to(device, non_blocking=True)
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        out = model(obs, batch["action"].to(device))
        loss = out.loss
    loss.backward()
    optimizer.step(); optimizer.zero_grad(set_to_none=True)

📊 Memory budget (7B model)

~14 GB — weights in bf16
~14 GB — gradients in bf16
~56 GB — Adam states in fp32 (m, v + master copy)
variable — activations, which can be reduced with gradient checkpointing

bf16

fp32 range, without GradScaler.

autocast

Selective casts per op.

non_blocking

Asynchronous H2D copy.

VRAM budget

Weights + gradient + activation optimization.

🔹 Transformers/HF: loading VLMs, LoRA, quantization

The backbone of almost every VLA is a pretrained VLM. With transformers + peft + bitsandbytes load the quantized model in 4-bit (NF4) and train only adapters LoRA — fine-tuning 7B on a single 24 GB GPU.

pip install "transformers>=4.45" peft bitsandbytes accelerate

from transformers import AutoModelForVision2Seq, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model

qcfg = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                          bnb_4bit_compute_dtype=torch.bfloat16)
vlm = AutoModelForVision2Seq.from_pretrained(
    "openvla/openvla-7b", quantization_config=qcfg, device_map="auto")

lora = LoraConfig(r=32, lora_alpha=16, lora_dropout=0.05,
                  target_modules=["q_proj","k_proj","v_proj","o_proj"])
model = get_peft_model(vlm, lora)
model.print_trainable_parameters()   # ~0.5% dos params treináveis

✓ Do

✓Target the LLM’s q/k/v/o projections and MLP with LoRA r=16–64.
✓Keep the action head at full precision and trainable.
✓Save only the adapters; merge (merge_and_unload) in deployment.

✗ Avoid

✗Quantize the vision encoder and expect spatial fidelity.
✗Small-r LoRA (r=4) on a complex manipulation task.
✗Forgetting to freeze/normalize action statistics.

LoRA

Low-rank adapters.

QLoRA / NF4

4-bit + bf16 adapters.

PEFT

Lightweight fine-tuning library.

target_modules

Where to inject the adapters.

🔹 Data pipeline: Dataset, DataLoader, normalization

The difference between a VLA that works and one that doesn't is often action normalization. Action heads learn poorly when dimensions have widely different scales (mm of Δpose vs. binary gripper state). Normalize by dimension (mean/standard deviation or 1–99% quantiles) and save the statistics with the checkpoint.

from torch.utils.data import DataLoader

loader = DataLoader(dataset, batch_size=64, shuffle=True,
                    num_workers=8, pin_memory=True,
                    persistent_workers=True, prefetch_factor=4)

# normalização por dimensão (salvar stats com o ckpt!)
mean, std = stats["action"]["mean"], stats["action"]["std"]
def normalize(a): return (a - mean) / (std + 1e-6)
def denormalize(a): return a * std + mean   # usado na inferência

⚡ Practical tip

Image augmentations (color jitter, light random cropping) help visual robustness, but never apply flips/rotations that invert the action’s spatial semantics—left becomes right and the policy breaks. Decoded video is the bottleneck: pre-decode it or use a video backend (torchcodec).

Normalization

By action dimension.

num_workers

Parallelize IO/decode.

pin_memory

Fast H2D transfer.

Augmentation

Visual, yes; spatial, be careful.

🔹 Distributed training: DDP, FSDP, gradient checkpointing

When the model doesn't fit on a single GPU, you move from DDP (replicates everything, performs gradient all-reduce) for FSDP (Fully Sharded Data Parallel): weights, gradients, and optimizer states are sliced across GPUs and brought together only when the layer executes. That’s what makes it possible to train 7B+ VLAs on a few nodes.

DDP — fits on one GPU

Full replica per device; only gradients are transferred (all-reduce). Simple and fast up to ~1–2B.

FSDP — doesn't fit

Shard of params/grad/optim; all-gather per layer in the forward pass, reduce-scatter in the backward pass.

Gradient checkpointing

Trades memory for compute: recomputes activations during the backward pass. Crucial for a decent batch size on 7B.

# lançar treino com torchrun em 8 GPUs
torchrun --nproc_per_node=8 train_vla.py --fsdp full_shard --bf16

# no código:
model.gradient_checkpointing_enable()
# accelerate config / FSDP policy: transformer_auto_wrap_policy

DDP

Gradient all-reduce.

FSDP

Full state sharding.

Checkpointing

Memory ↔ compute.

torchrun

Multi-GPU launcher.

🔹 Inference: torch.compile, export, serve at 10–50 Hz

On the robot, the budget changes: latency is king. A 30 Hz control loop gives you ~33 ms per step, from the camera frame to the command. torch.compile fuses kernels and reduces Python overhead; action chunking (predicting K steps at once) amortizes inference across several cycles.

model.eval()
policy = torch.compile(model, mode="reduce-overhead")  # CUDA graphs

@torch.inference_mode()
def act(obs):
    with torch.autocast("cuda", dtype=torch.bfloat16):
        chunk = policy(obs)            # prevê K ações de uma vez
    return denormalize(chunk).cpu().numpy()

# loop de controle a 30 Hz com action chunking (executa K antes de re-inferir)

📊 Where the time goes (33 ms step)

capture + preprocessing — decode and resize RGB-D frames
model forward pass — optimization target (compile, quantization)
post-processing — denormalize, map to the arm frame
chunking — re-infer every K steps, not every step

✓ Do

✓Warm up the torch.compile with dummy forwards before the real loop.
✓Fix input shapes to avoid recompilations.
✓Measure p99 latency, not just the average.

✗ Avoid

✗Leave autograd enabled during inference (without inference_mode).
✗Re-infer at every step, wasting the predicted chunk.
✗Forgetting to denormalize the action before sending it to the robot.

torch.compile

Fusion via Inductor.

inference_mode

No autograd, faster.

Action chunking

Amortizes latency.

Export

TorchScript/ONNX for edge.

✅ Module summary

✓

PyTorch eager — debugging and rapid iteration make it the VLA research standard.

✓

bf16 + memory budget — think about weights+grad+optim+activations before choosing a GPU.

✓

QLoRA — fine-tune 7B on 24 GB; save adapters only.

✓

Action normalization — always save statistics with the checkpoint.

✓

FSDP + compile — scale training and cut inference latency.

Next module

2.2 — LeRobot: Hugging Face’s end-to-end framework, where this entire stack becomes a record → train → deploy workflow.

← Back to the track Next module →