OpenTau: VLA Training Infrastructure for Real-World Robotics

OpenTau is an open-source training toolchain designed for training Vision-Language-Action (VLA) models. It is built as a fork of the LeRobot repository, meaning it is compatible with any LeRobot-compliant policy and dataset.

Key capabilities include:

• Co-training on heterogeneous datasets.
• Discrete actions for fast VLM convergence.
• Knowledge Insulation (KI) between the VLM backbone and the action expert.
• Support for various policy architectures: $\pi_{0.5}$, $\pi_{0.6}$, $\pi_{0.7}$ (hierarchical), and cosmos3 (frozen backbone with custom action expert).
• Multi-node and multi-GPU training support.
• Integrated training profilers and unused-parameter auditors.

OpenTau provides several advanced training techniques for Vision-Language-Action (VLA) models, specifically designed to implement the methodologies found in the $\pi$-series papers:

• Co-training: Training on an adjustable mixture of heterogeneous datasets.
• Discrete Actions: Enables faster VLM convergence in $\pi_{0.5}$ models.
• Knowledge Insulation: Separates the VLM backbone from the action expert.
• Response Prediction: Supports training on VQA datasets and high-level planning during inference in $\pi_{0.5}$.
• Overfitting Prevention: Includes dropout layers within the VLM.
• RL Pipeline: Implements the reinforcement learning pipeline described in $\pi^*_{0.6}$.
• Scalability: Supports multi-node and multi-GPU training.
• Evaluation: Includes simulation environments for model evaluation.

In pi05_ttt sequence training, the window is anchored at the last timestep (the 'live' frame). While observations use existing history windows, actions require new temporal offsets to provide a chunk of actions for every timestep $t$ in the sequence.

For a sequence of length $T$, a chunk size $H$, and a temporal stride stride, the action offsets are calculated as a flat list of $T \times H$ offsets:

# timestep t, chunk position h  ->  one flat list of T * H offsets
[(-(T - 1 - t) * stride + h) / action_freq
 for t in range(T) for h in range(H)]

These are fetched as a tensor of shape (T * H, dim) and then reshaped to (T, H, dim) during item assembly.

OpenTau uses Draccus for configuration management. The central orchestration class is opentau.configs.train.TrainPipelineConfig. This class manages training settings, policy configurations, and environment setups. You can load configurations directly from pretrained checkpoints to ensure experiment reproducibility.

# Example concept: loading a config from a checkpoint
# (Actual implementation depends on the Draccus/OpenTau integration)
config = opentau.configs.train.TrainPipelineConfig.from_checkpoint("path/to/checkpoint")

Metadata defines the structure and statistics of a dataset, managed by LeRobotDatasetMetadata and DatasetMetadata. It is stored in JSON (info.json, stats.json) and JSONL (tasks.jsonl) files. Key components include:

• Info: Contains feature shapes, data types, FPS, and robot_type (a LeRobot v2 field).
• Stats: Contains mean, standard deviation, min, and max values for normalization.
• Tasks: Natural language descriptions of the tasks in the dataset.

The server communicates via WebSocket using MessagePack. The server uses policy.sample_actions to predict a temporal chunk of actions rather than a single step.

Request Formats

• Single observation: A dictionary containing:
  • images: JPEG bytes per camera name.
  • state: A list of floats.
  • prompt: A string.
• Batch: {"batch": true, "items": [ { ...single_observation_fields... }, ... ]}

Response Formats

• Single: A single action chunk as nested lists: [[float, ...], ...] with shape (T, action_dim), where T is the predicted horizon (e.g., n_action_steps).
• Batch: A list of chunks: [ chunk_0, chunk_1, ... ], where each chunk corresponds to a row in the input items list.

OpenTau supports various Vision-Language-Action (VLA) models. To select a model, use the --policy.type CLI flag or set the policy.type key in your JSON configuration file. The available models are categorized into two backbone families:

• PaliGemma family: Uses the paligemma_with_expert.py wrapper (PaliGemma VLM + Gemma action expert). Includes pi0, pi05, pi05_mem, and pi07_paligemma.
• Gemma 3 family: Uses the gemma3_with_expert.py wrapper (Gemma 3 4B VLM + Gemma action expert) at 448×448 resolution. Includes pi06 and pi07.

Models emit actions using flow matching (continuous), FAST discrete-action tokens (autoregressive), or both.

# Example using CLI flag
python train.py --policy.type=pi05

{
  "policy": {
    "type": "pi06"
  }
}

When training on datasets with multiple cameras (like LIBERO), training time per step increases significantly because the frozen VLM prefix is recomputed for every timestep in a sequence.

For example, on LIBERO, a step might take ~5.3s compared to ~2.97s on droid_100 due to the two-camera setup. To reduce this overhead and enable longer context, implement caching for the frozen VLM prefix so it is not recomputed for every timestep within a single training step.

A single policy.sample_actions call consists of a fixed backbone prefill cost and a variable cost based on the number of denoising steps. The latency follows an affine relationship:

L(N) ≈ F + N·s

Where:

• F (Fixed Cost): Includes the Qwen3-VL backbone prefill (over images + prompt), CPU-side preprocessing, and rope/noise setup.
• N: The number of denoising steps (num_steps).
• s (Step Cost): The cost of one action-expert forward pass (e.g., ~0.9B parameters) cross-attending the cached backbone KV.

Performance Tip: torch.compile provides significant speedups primarily by optimizing the expert loop (s), which is often launch- and fusion-bound due to many small kernels across multiple layers and steps. The compute-bound backbone prefill (F) sees much smaller improvements.