Python + PyTorch: Stack dominante para VLA

Contenido detallado

🔹 Por qué PyTorch domina la robótica y la investigación

Prácticamente todo VLA abierto — OpenVLA, Octo, π0, ACT — está escrito en PyTorch. La razón es cultural y técnica: la ejecución eager (define-by-run) deja el grafo dinámico, así que inspeccionas tensores en medio del forward, colocas breakpoint() dentro de la policy y se itera en minutos. Para la investigación en manipulación, donde la arquitectura cambia cada semana, esto vale oro.

📊 Por qué ganó el ecosistema

autograd — diferenciación automática inversa, base de cualquier entrenamiento de policy.
Hugging Face + PEFT — VLMs listos para usar (PaliGemma, Qwen2-VL) con LoRA en pocas líneas.
torch.compile — fusión de kernels mediante TorchInductor sin reescribir código.
CUDA/ROCm/MPS — el mismo código se ejecuta en un centro de datos, Jetson y Mac.

⚡ Consejo práctico

JAX (con MJX/Flax) destaca en simulación masivamente paralela y física diferenciable, pero el deploy en un robot real y la mayoría de los checkpoints VLA están en PyTorch. Aprende ambos; empieza por PyTorch.

Modo eager

Define-by-run, depuración nativa.

autograd

Backprop automático.

Ecosistema

HF, PEFT, timm, accelerate.

Portabilidad

Mismo código, varios objetivos.

🔹 Tensores, GPU y mixed precision en modelos 7B

Un VLA de 7B en fp32 ocupa ~28 GB solo en pesos. Entrenar en bf16 reduce a la mitad y además evita el overflow de fp16 (bf16 tiene el mismo exponente que fp32). La regla mental de memoria: pesos + gradientes + estados del optimizador + activaciones. Con Adam en fp32, solo los estados cuestan ~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)

📊 Presupuesto de memoria (modelo 7B)

~14 GB — pesos en bf16
~14 GB — gradientes en bf16
~56 GB — estados Adam en fp32 (m, v + copia maestra)
variable — activaciones, reducibles mediante gradient checkpointing

bf16

Rango de fp32, sin GradScaler.

autocast

Conversiones de tipo selectivas por operación.

non_blocking

Copia H2D asíncrona.

VRAM budget

Pesos+grad+activaciones optimizadas.

🔹 Transformers/HF: cargar VLMs, LoRA, cuantización

El backbone de casi todo VLA es un VLM preentrenado. Con transformers + peft + bitsandbytes cargas el modelo cuantizado en 4-bit (NF4) y entrenas solo adaptadores LoRA — ajuste fino de 7B en una sola 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

✓ Hacer

✓Apunta a las proyecciones q/k/v/o y al MLP del LLM con LoRA r=16–64.
✓Mantener el action head con precisión completa y entrenable.
✓Guardar solo los adaptadores; fusionar (merge_and_unload) en el deploy.

✗ Evitar

✗Cuantizar el encoder de visión y esperar fidelidad espacial.
✗LoRA con r minúscula (r=4) en una tarea de manipulación compleja.
✗Olvidar congelar/normalizar las estadísticas de acción.

LoRA

Adaptadores de bajo rango.

QLoRA / NF4

4-bit + adaptadores bf16.

PEFT

Librería de fine-tuning ligero.

target_modules

Dónde inyectar los adaptadores.

🔹 Pipeline de datos: Dataset, DataLoader, normalización

La diferencia entre un VLA que funciona y uno que no suele ser normalización de acciones. Los action heads aprenden mal cuando las dimensiones tienen escalas dispares (mm de Δpose frente al estado binario del gripper). Estandariza por dimensión (media/desviación o cuantiles 1–99%) y guarda las estadísticas junto al 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

⚡ Consejo práctico

Las augmentations de imagen (color jitter, random crop ligero) ayudan a la robustez visual, pero nunca apliques flips/rotations que inviertan la semántica espacial de la acción: izquierda se convierte en derecha y la policy falla. El video decodificado es el cuello de botella: precodifícalo o usa un backend de video (torchcodec).

Normalización

Por dimensión de acción.

num_workers

Paralelizar IO/decode.

pin_memory

Transferencia H2D rápida.

Augmentation

Visual, sí; cuidado con lo espacial.

🔹 Entrenamiento distribuido: DDP, FSDP, gradient checkpointing

Cuando el modelo no cabe en una GPU, pasas de DDP (replica todo, hace all-reduce de gradientes) para FSDP (Fully Sharded Data Parallel): los pesos, gradientes y estados del optimizador se segmentados entre GPUs y se reúnen solo cuando se ejecuta la capa. Esto permite entrenar VLA 7B+ en pocos nodos.

DDP — cabe en una GPU

Una réplica completa por dispositivo; solo se transfieren los gradientes (all-reduce). Simple y rápido hasta ~1–2B.

FSDP — no cabe

Shard de params/grad/optim; all-gather por capa en el forward, reduce-scatter en el backward.

Gradient checkpointing

Intercambia memoria por capacidad de cómputo: vuelve a calcular activaciones durante el backward. Es crucial para lograr un batch decente en 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 de todo el estado.

Checkpointing

Memoria ↔ cómputo.

torchrun

Lanzador multi-GPU.

🔹 Inferencia: torch.compile, exportación, servir a 10–50 Hz

En el robot, el presupuesto cambia: la latencia es clave. Un loop de control a 30 Hz da ~33 ms por paso, desde el frame de la cámara hasta el comando. torch.compile fusiona kernels y reduce el overhead de Python; action chunking (predecir K pasos a la vez) amortiza la inferencia en varios 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)

📊 A dónde va el tiempo (paso de 33 ms)

captura + preprocesamiento — decodificar y redimensionar frames RGB-D
forward del modelo — objetivo de optimización (compile, cuantización)
posprocesamiento — desnormalizar, mapear al frame del brazo
chunking — volver a inferir cada K pasos, no en cada paso

✓ Hacer

✓Preparar el torch.compile con forwards dummy antes del bucle real.
✓Fijar las formas de entrada para evitar recompilaciones.
✓Medir la latencia p99, no solo el promedio.

✗ Evitar

✗Dejar autograd activado durante la inferencia (sin inference_mode).
✗Volver a inferir en cada paso y desperdiciar el chunk previsto.
✗Olvidar desnormalizar la acción antes de enviarla al robot.

torch.compile

Fusión mediante Inductor.

inference_mode

Sin autograd, más rápido.

Action chunking

Amortiza la latencia.

Export

TorchScript/ONNX para edge.

✅ Resumen del módulo

✓

PyTorch eager — la depuración y la iteración rápida lo convierten en el estándar de investigación VLA.

✓

bf16 + presupuesto de memoria — considera pesos+grad+optimizador+activaciones antes de elegir GPU.

✓

QLoRA — ajuste fino de 7B en 24 GB; guarda solo adaptadores.

✓

Normalización de acciones — guarda siempre las estadísticas con el checkpoint.

✓

FSDP + compile — escalar el entrenamiento y reducir la latencia de inferencia.

Próximo módulo

2.2 — LeRobot: framework end-to-end de Hugging Face, donde todo este stack se convierte en un flujo de grabar → entrenar → deploy.

← Volver a la ruta Próximo módulo →