MODULE 2.1

🐍 Python ecosystem: PyTorch, Transformers, Gymnasium

The tools every VLA engineer needs to master: from tensor to environment. PyTorch for autograd, HuggingFace for pre-trained models, Gymnasium for standardized environments.

Topics

~3h

Estimated

Technical

Practical

Type

🔥 PyTorch for robotics: tensors and autograd

The dominant deep learning framework in robotics. Understanding tensors and the automatic differentiation system is a technical prerequisite for any serious work with VLA.

💡 What it is

PyTorch is the dominant deep learning framework in robotics and VLA. Tensors are multidimensional arrays with native GPU support via CUDA. The system autograd automatically computes gradients through the computational graph — essential for backpropagation during neural policy training. Virtually every relevant VLA (OpenVLA, pi-0, ACT) is implemented in PyTorch.

✓ Do

✓ Use bfloat16 on Ampere+ GPUs to train large VLAs
✓ Move tensors to the device before operating on them: .to(device)
✓ Use torch.compile() to speed up repetitive training loops
✓ Call optimizer.zero_grad() before each backward

✗ Avoid

✗ Accumulate tensors in Python lists — silently leaks GPU memory
✗ Mix float32 e float16 without torch.amp
✗ Call .item() inside the training loop (synchronizes CPU/GPU)
✗ Forget model.eval() / torch.no_grad() during inference

📦 Minimum pattern: training a PyTorch policy

import torch
import torch.nn as nn

# device-agnostic
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = PolicyNet().to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

# mixed precision (Ampere+)
scaler = torch.cuda.amp.GradScaler()

for batch in dataloader:
    obs, action = batch["obs"].to(device), batch["action"].to(device)
    optimizer.zero_grad()
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        pred = model(obs)
        loss = nn.functional.mse_loss(pred, action)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

torch.Tensor

ND array with device, dtype, and grad

nn.Module

Base for every network; forward() required

autograd

Computational graph + backward()

torch.compile

JIT for repetitive loops (2× faster)

🤗 HuggingFace Transformers: Pretrained Models

The library that democratized access to language and vision models—and hosts practically every relevant open-source VLA.

⚡ Why learn first

The HuggingFace Hub is the center of the open-source VLA ecosystem. OpenVLA, SmolVLA, Qwen-VL, PaliGemma — all distributed via Hub with a ready-to-use inference pipeline. LeRobot hosts 58k+ datasets there. Without mastering from_pretrained(), you can’t even run inference.

300k+

models on the Hub

58k+

LeRobot datasets

PEFT

Efficient LoRA fine-tuning

🔧 Inference with OpenVLA via HuggingFace

from transformers import AutoModelForVision2Seq, AutoProcessor
import torch

processor = AutoProcessor.from_pretrained("openvla/openvla-7b", trust_remote_code=True)
model = AutoModelForVision2Seq.from_pretrained(
    "openvla/openvla-7b",
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
    trust_remote_code=True
).to("cuda")

# Processar frame + instrução em linguagem natural
inputs = processor(
    "In: What action should the robot take to pick up the cup?\nOut:",
    image,
    return_tensors="pt"
).to("cuda", dtype=torch.bfloat16)

action = model.predict_action(**inputs, unnorm_key="bridge_orig", do_sample=False)

✓ Best practices

✓ Use safetensors instead of pickle (.bin)
✓ Version models with revision="main" or tag
✓ LoRA for fine-tuning: trains 1-10% of the parameters
✓ Use flash_attention_2 when available

✗ Common pitfalls

✗ Forget trust_remote_code=True in custom VLAs
✗ Download models without checking their commercial use license
✗ Load full model without low_cpu_mem_usage (OOM)
✗ Mix processor from model A with model B

AutoModel

Automatically loads the correct architecture

LoRA/PEFT

Efficient fine-tuning with few parameters

safetensors

Safe, fast format for weights

API Hub

Push/pull models and datasets programmatically

🏋️ Gymnasium: standardized environments

The universal API between policy and environment. MuJoCo, Isaac Sim, and PyBullet implement this interface—understanding obs/act/reward/done is fundamental to any training loop.

🔍 What it is

Gymnasium (successor to OpenAI Gym) defines the standard interface for RL environments. In VLA robotics, it serves as a wrapper for simulators and lets you train policies with a consistent interface. The API is simple but powerful:

Spaces

• observation_space — what the robot sees
• action_space — what it can do
• Box for continuous, Discrete for discrete

Lifecycle

• env.reset() → initial state
• env.step(action) → (obs, reward, done, info)
• Loop up to terminated or truncated

Gymnasium episode loop

`obs, info = env.reset(seed=42)`

Starts an episode and returns the initial observation. The seed ensures reproducibility of the initial state.

`action = policy(obs)`

Policy processes an observation and returns an action in continuous space (e.g., joint velocities, Cartesian delta).

`obs, reward, terminated, truncated, info = env.step(action)`

The environment executes the action and returns a new observation, a reward signal, and termination flags.

`VectorEnv: N ambientes paralelos`

To train faster: gymnasium.vector.make("env-id", num_envs=8) runs 8 simultaneous episodes in a single process.

gym.Env

Base class to subclass

Box(shape)

Continuous n-dimensional space

Wrappers

Transform obs/act without touching env

VectorEnv

N parallel environments in one process

📦 Datasets and DataLoaders for trajectories

Data is the number one bottleneck in robotics. Structuring, loading, and augmenting trajectories efficiently is just as important as the model architecture.

📊 Robotics Dataset Formats

RLDS (Google): TFRecord with episode metadata, the Open X-Embodiment standard with 1M+ trajectories.
LeRobot Dataset: Parquet + MP4 video on HuggingFace Hub. 58k+ datasets, native Python API with from_hub().
Custom HDF5: Used in academic robotics (RoboMimic, R2R2). Hierarchical and efficient for long sequences.

RLDS / Open X

• 1M+ real-world trajectories
• TFRecord (requires TF)
• Rich metadata per step
• Google/DeepMind standard

LeRobot Dataset ★

• 58k+ datasets on Hub
• Parquet + MP4 video
• Clean Python API
• Native training integration

HDF5 / Proprio

• Flexible and hierarchical
• Fast random access
• No cloud dependency
• Common in academia

💡 Tip: DataLoader for trajectories

Trajectories have variable lengths. Use collate_fn customized for padding or truncating sequences when assembling batches. Combine with num_workers=4 e pin_memory=True to maximize GPU throughput.

loader = DataLoader(
    dataset, batch_size=32, num_workers=4,
    pin_memory=True, collate_fn=pad_trajectory_batch
)

RLDS

TFRecord with Open X episodes

LeRobot DS

Parquet+video, 58k datasets

collate_fn

Build batches of variable-length trajectories

augmentation

Robotics-specific crop/jitter

📈 Weights & Biases: experiment tracking

Without structured tracking, VLA training becomes a black box. W&B makes loss, success rate, and hyperparameters visual and reproducible.

💡 W&B + LeRobot: native integration

LeRobot has native W&B support via YAML config. Just set wandb.enable: true in the configuration file — without a single extra line of code.

# config/train_act.yaml
wandb:
  enable: true
  project: "vla-pickup-cup"
  entity: "meu-time"
  notes: "ACT baseline lr=1e-4"

wandb.init()

Starts an experiment run

Sweeps

Automatic hyperparameter search

Tables

Visualize trajectories and predictions

Registry

Checkpoint versioning

🐳 Docker and reproducibility in robotics

The CUDA driver + PyTorch + MuJoCo combination can break in mysterious ways. Containers remove “it works on my machine” from the equation.

🔍 Why Docker is critical in robotics

Docker containerizes the development environment: exact versions of Python, PyTorch, CUDA, and dependencies stay fixed. In robotics, where the combination CUDA driver + PyTorch + MuJoCo can break easily; containers are essential. NVIDIA NGC and LeRobot provide ready-to-use containers with everything configured.

✓ Recommended setup

✓ Start with an NVIDIA NGC image: nvcr.io/nvidia/pytorch:24.01-py3
✓ Multi-stage build: builder → runtime (smaller image)
✓ Mount datasets via volumes; don't copy them into the image
✓ devcontainer for VS Code with GPU passthrough

✗ Common errors

✗ Copy the entire dataset into the Docker image
✗ Use docker run without --gpus all
✗ Ignore CUDA driver vs. PyTorch CUDA compatibility
✗ Commit secrets (HF tokens, W&B API key) in the Dockerfile

🐳 Minimal Dockerfile for VLA with LeRobot

# Base: PyTorch + CUDA oficial NVIDIA
FROM nvcr.io/nvidia/pytorch:24.01-py3

WORKDIR /workspace

# Instalar LeRobot e dependências
RUN pip install lerobot[gymnasium,wandb] huggingface_hub

# Montar datasets externamente (não copiar)
VOLUME ["/data"]

# GPU passthrough: nvidia-docker2 ou --gpus all
CMD ["python", "train.py", "--config", "config/act.yaml"]

💡 devcontainers: GPU in VS Code

Add .devcontainer/devcontainer.json with "runArgs": ["--gpus", "all"] and the whole team opens the project in the same environment, including CUDA, with no manual setup. LeRobot provides a ready-to-use devcontainer in the official repository.

NVIDIA NGC

Optimized PyTorch+CUDA containers

--gpus all

GPU passthrough to container

devcontainer

Consistent environment in VS Code

multi-stage

Builder + runtime = lean image

✅ Module Summary

✓

PyTorch is at the heart of VLA training — tensors, autograd, mixed precision, and torch.compile form the technical foundation for every engineer.

✓

HuggingFace Hub + LoRA — access to 300k+ pretrained models and efficient fine-tuning with 1-10% of the original parameters.

✓

Gymnasium: universal obs/act/reward contract — any simulator (MuJoCo, Isaac Sim, PyBullet) speaks this language; VectorEnv parallelizes episodes.

✓

Docker + W&B = reproducibility — fixed environments with containers and visual experiment tracking are as important as the model.

Next module

2.2 — LeRobot: End-to-End Pipeline (collection, training, and deployment on a real robot)

← Back to the Learning Path Next module →