Octo

Load and execute a pretrained transformer-based diffusion model trained on 800K diverse robot episodes from the Open X-Embodiment dataset. The model processes multimodal observations (images from multiple camera views, proprioceptive state) and task specifications (language instructions or goal images) through a causal transformer backbone, then decodes actions via learned action heads (diffusion or L1-based). Inference runs through OctoModel.sample_actions() which handles tokenization, transformer forward pass, and action sampling in a single call.

Adapt a pretrained Octo model to a new robot by freezing the transformer backbone and retraining only the observation tokenizers, task tokenizers, and action heads on your robot's specific sensor/action configuration. The framework provides efficient fine-tuning via gradient-based optimization on small datasets (100s-1000s of trajectories), using callbacks for monitoring and early stopping. Fine-tuning leverages the pretrained transformer's learned representations, reducing sample complexity compared to training from scratch.

Evaluate trained policies on simulation environments (MuJoCo, PyBullet) and real robots using standardized metrics (success rate, trajectory length, task completion time). The system provides evaluation scripts that run policies in closed-loop control, collect rollouts, and compute metrics. Evaluation supports both deterministic (L1 head) and stochastic (diffusion head) policies, enabling comparison of action prediction methods.

Define model architecture, training hyperparameters, and data pipeline via configuration files (YAML or Python configs in scripts/configs/). Configurations specify transformer depth/width, tokenizer types, action head type, learning rate, batch size, and dataset paths. This abstraction enables reproducible experiments and easy hyperparameter sweeps without modifying code.

Encode task specifications as either natural language instructions or goal images, processed through dedicated task tokenizers that convert them into transformer-compatible token sequences. Language tasks use a language tokenizer (e.g., T5-based) to embed instructions like 'pick up the red cube'; visual goals use an image tokenizer to embed a target image showing the desired end state. Both are concatenated with observation tokens in the transformer input sequence, enabling the model to condition action prediction on either modality.

Convert raw sensor observations (RGB images from multiple cameras, proprioceptive state like joint angles/velocities) into fixed-size token sequences via modular observation tokenizers. Image tokenizers use learned or pretrained vision encoders (e.g., ViT, ResNet) to compress images into tokens; proprioception tokenizers embed joint states as learnable embeddings. Multiple camera views are tokenized independently and concatenated, enabling the transformer to attend across all sensor modalities in a unified sequence.

Predict robot actions from transformer outputs using learned action heads that decode token representations into action sequences. Diffusion-based heads use iterative denoising (reverse diffusion process) to sample actions, enabling multi-modal action distributions and better handling of stochastic tasks; L1 regression heads directly predict action means, offering faster inference but assuming unimodal action distributions. Both heads support action chunking (predicting multiple future timesteps) and can be swapped during fine-tuning.

Core transformer architecture (OctoTransformer) processes tokenized observations and task specifications in a causal (autoregressive) manner, where each position attends only to previous tokens in the sequence. The transformer learns to predict the next action token given the history of observations and task context. Architecture uses standard transformer blocks (multi-head self-attention, feed-forward layers) with positional embeddings to encode temporal structure, enabling the model to learn temporal dependencies in robot trajectories.

Load and preprocess robot trajectory data from the Open X-Embodiment dataset (800K episodes across 22+ robot embodiments) using a unified data pipeline. The system handles multiple data formats (HDF5, tfrecord), performs on-the-fly transformations (image resizing, normalization, augmentation), and batches trajectories for training. Dataset loading is abstracted via a modular interface (octo/data/dataset.py) supporting custom observation/action spaces, enabling seamless integration of new robot data.

Apply learned and heuristic augmentations to training data to improve generalization and robustness. Image augmentations include resizing, color jittering, and random crops; task augmentations include paraphrasing language instructions and generating synthetic goal images from trajectory frames. Augmentations are applied on-the-fly during training, reducing memory overhead and enabling diverse data views from limited trajectories.

Provide Gym-compatible wrappers (NormalizeProprio, HistoryWrapper, RHCWrapper) that interface Octo policies with robot environments and simulators. Wrappers handle observation normalization, history buffering (stacking recent observations), and receding horizon control (RHC) where actions are re-planned at each timestep. This abstraction enables drop-in deployment of Octo policies to any Gym-compatible environment without modifying the policy code.

Provide a configurable training loop (scripts/configs/octo_pretrain_config.py) with callbacks for logging, checkpointing, and early stopping. The system tracks training metrics (loss, validation accuracy), saves model checkpoints at regular intervals, and supports distributed training across multiple GPUs. Callbacks enable custom monitoring logic (e.g., periodic evaluation on held-out tasks) without modifying core training code.