Python + PyTorch: Stack Dominante para VLA

Conteúdo detalhado

🔹 Por que PyTorch domina robótica e pesquisa

Praticamente todo VLA aberto — OpenVLA, Octo, π0, ACT — é escrito em PyTorch. A razão é cultural e técnica: execução eager (define-by-run) deixa o grafo dinâmico, então você inspeciona tensores no meio do forward, coloca breakpoint() dentro da policy e itera em minutos. Para pesquisa em manipulação, onde a arquitetura muda toda semana, isso vale ouro.

📊 Por que o ecossistema venceu

autograd — diferenciação automática reversa, base de qualquer treino de policy.
Hugging Face + PEFT — VLMs prontos (PaliGemma, Qwen2-VL) com LoRA em poucas linhas.
torch.compile — fusão de kernels via TorchInductor sem reescrever código.
CUDA/ROCm/MPS — mesmo código roda em data center, Jetson e Mac.

⚡ Dica prática

JAX (com MJX/Flax) brilha em sim massivamente paralelo e física diferenciável, mas o deploy em robô real e a maioria dos checkpoints VLA vivem em PyTorch. Saiba os dois; comece pelo PyTorch.

Eager mode

Define-by-run, debug nativo.

autograd

Backprop automático.

Ecossistema

HF, PEFT, timm, accelerate.

Portabilidade

Mesmo código, vários alvos.

🔹 Tensores, GPU e mixed precision em modelos 7B

Um VLA de 7B em fp32 são ~28 GB só de pesos. Treinar em bf16 corta pela metade e ainda evita o overflow do fp16 (bf16 tem o mesmo expoente do fp32). A regra mental de memória: pesos + gradientes + estados do otimizador + ativações. Com Adam em fp32, os estados sozinhos custam ~8 bytes/parâmetro.

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)

📊 Orçamento de memória (modelo 7B)

~14 GB — pesos em bf16
~14 GB — gradientes em bf16
~56 GB — estados Adam em fp32 (m, v + cópia mestre)
variável — ativações, reduzíveis por gradient checkpointing

bf16

Range do fp32, sem GradScaler.

autocast

Casts seletivos por op.

non_blocking

Cópia H2D assíncrona.

VRAM budget

Pesos+grad+otim+ativações.

🔹 Transformers/HF: carregar VLMs, LoRA, quantização

O backbone de quase todo VLA é um VLM pré-treinado. Com transformers + peft + bitsandbytes você carrega o modelo quantizado em 4-bit (NF4) e treina só adaptadores LoRA — fine-tuning de 7B em uma única GPU de 24 GB.

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

✓ Fazer

✓Mirar q/k/v/o proj e MLP do LLM com LoRA r=16–64.
✓Manter a action head em precisão cheia, treinável.
✓Salvar só os adaptadores; mesclar (merge_and_unload) no deploy.

✗ Evitar

✗Quantizar o encoder de visão e esperar fidelidade espacial.
✗LoRA r minúsculo (r=4) numa tarefa de manipulação complexa.
✗Esquecer de congelar/normalizar estatísticas de ação.

LoRA

Adaptadores de baixo rank.

QLoRA / NF4

4-bit + adaptadores bf16.

PEFT

Lib de fine-tuning leve.

target_modules

Onde injetar os adaptadores.

🔹 Pipeline de dados: Dataset, DataLoader, normalização

A diferença entre um VLA que funciona e um que não é frequentemente normalização de ações. Action heads aprendem mal quando dimensões têm escalas díspares (mm de Δpose vs estado binário do gripper). Padronize por dimensão (média/desvio ou quantis 1–99%) e guarde as estatísticas junto do 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

⚡ Dica prática

Augmentations de imagem (color jitter, random crop leve) ajudam a robustez visual, mas nunca aplique flips/rotations que invertam a semântica espacial da ação — esquerda vira direita e a policy quebra. Vídeo decodificado é o gargalo: pré-decodifique ou use backend de vídeo (torchcodec).

Normalização

Por dimensão de ação.

num_workers

Paralelizar IO/decode.

pin_memory

Transferência H2D rápida.

Augmentation

Visual sim, espacial cuidado.

🔹 Treino distribuído: DDP, FSDP, gradient checkpointing

Quando o modelo não cabe numa GPU, você sai de DDP (replica tudo, faz all-reduce de gradientes) para FSDP (Fully Sharded Data Parallel): pesos, gradientes e estados do otimizador são fatiados entre GPUs e reunidos só quando a camada executa. É o que viabiliza treinar VLA 7B+ em poucos nós.

DDP — cabe numa GPU

Réplica completa por device; só os gradientes trafegam (all-reduce). Simples e rápido até ~1–2B.

FSDP — não cabe

Shard de params/grad/optim; all-gather por camada no forward, reduce-scatter no backward.

Gradient checkpointing

Troca memória por compute: recomputa ativações no backward. Crucial para batch decente em 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

All-reduce de gradientes.

FSDP

Shard total de estado.

Checkpointing

Memória ↔ compute.

torchrun

Lançador multi-GPU.

🔹 Inferência: torch.compile, exportação, servir a 10–50 Hz

No robô, o orçamento muda: latência é rei. Um loop de controle a 30 Hz dá ~33 ms por passo, do frame da câmera ao comando. torch.compile funde kernels e reduz overhead Python; action chunking (prever K passos de uma vez) amortiza a inferência sobre vários ciclos.

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)

📊 Onde o tempo vai (passo de 33 ms)

captura + pré-proc — decodificar e redimensionar frames RGB-D
forward do modelo — alvo de otimização (compile, quantização)
pós-proc — denormalizar, mapear para frame do braço
chunking — re-inferir a cada K passos, não a cada passo

✓ Fazer

✓Aquecer o torch.compile com forwards dummy antes do loop real.
✓Fixar shapes de entrada para evitar recompilações.
✓Medir latência p99, não só a média.

✗ Evitar

✗Deixar autograd ligado na inferência (sem inference_mode).
✗Re-inferir a cada passo desperdiçando o chunk previsto.
✗Esquecer de denormalizar a ação antes de enviar ao robô.

torch.compile

Fusão via Inductor.

inference_mode

Sem autograd, mais rápido.

Action chunking

Amortiza latência.

Export

TorchScript/ONNX p/ edge.

✅ Resumo do módulo

✓

PyTorch eager — debug e iteração rápida fazem dele o padrão de pesquisa VLA.

✓

bf16 + orçamento de memória — pense em pesos+grad+otim+ativações antes de escolher GPU.

✓

QLoRA — fine-tune 7B em 24 GB; salve só adaptadores.

✓

Normalização de ações — guarde estatísticas com o checkpoint, sempre.

✓

FSDP + compile — escalar treino e cortar latência de inferência.

Próximo módulo

2.2 — LeRobot: framework end-to-end da Hugging Face, onde toda essa stack vira fluxo de gravar → treinar → deploy.

← Voltar à trilha Próximo módulo →