Octo Documentation

When providing observations to Octo, two mask types are used to control attention:

timestep_pad_mask

Indicates which observations in the history window should be attended to. Octo is typically trained with a history window size of 2 (current + previous observation).

• At the start of a trajectory: Since there is no previous observation, set timestep_pad_mask=False for the missing index.
• Window size of 1: If using a window size of 1, timestep_pad_mask should always be [True].
• Automation: If you use the HistoryWrapper from octo/utils/gym_wrappers.py, this mask is added to the observation dictionary automatically.

pad_mask_dict

Indicates which elements within a single timestep should be attended to. This is used for modalities that might be missing in a specific dataset.

• Example (No language): pad_mask_dict["language_instruction"] = False.
• Example (No wrist camera): pad_mask_dict["image_wrist"] = False.
• Default behavior: If a key is missing from the pad_mask_dict, it is treated as False for that key.

To install Octo, create a Conda environment with Python 3.10, install the package in editable mode, and install the required dependencies.

GPU Installation

If you are using a GPU, you must install the appropriate JAX version with CUDA support:

pip install --upgrade "jax[cuda11_pip]==0.4.20" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

TPU Installation

For TPU support, install the JAX TPU version:

pip install --upgrade "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

Verify Installation

You can verify the installation by running the finetuning script on the provided debug dataset:

python scripts/finetune.py --config.pretrained_path=hf://rail-berkeley/octo-small-1.5 --debug

conda create -n octo python=3.10
conda activate octo
pip install -e .
pip install -r requirements.txt

Octo can be finetuned on small target domain datasets to adapt to new sensory inputs, action spaces, or morphologies.

Finetuning Modes

You can specify which parts of the model to freeze using one of three modes:

• head_only: Only the readout heads are finetuned.
• head_mlp_only: Only the head MLP is finetuned.
• full: The entire model is finetuned.

Task Types

You can also specify the task type for finetuning:

• image_conditioned
• language_conditioned
• multimodal

Execution

Use the scripts/finetune.py script. You can combine modes and task types in the config. For example, to finetune the full transformer with image inputs only:

python scripts/finetune.py --config=finetune_config.py:full,image_conditioned

To run advanced finetuning with a specific pretrained path:

python scripts/finetune.py --config.pretrained_path=hf://rail-berkeley/octo-small-1.5

When running inference over a trajectory, you must feed inputs of the correct temporal window size.

1. Create Task: Use model.create_tasks(goals=...) for goal-conditioned tasks or model.create_tasks(texts=...) for language-conditioned tasks.
2. Prepare Observations: Stack a window of images to match the required WINDOW_SIZE. The observation dictionary must include image_primary and a timestep_pad_mask of the same temporal length.
3. Unnormalize Actions: model.sample_actions returns normalized actions. You must unnormalize them using the statistics found in model.dataset_statistics["<dataset_name>">["action"] to get usable control values.

WINDOW_SIZE = 2
# ... load images and task ...

for step in tqdm.trange(len(images) - (WINDOW_SIZE - 1)):
    # Stack images for the temporal window
    input_images = np.stack(images[step:step+WINDOW_SIZE])[None]
    
    observation = {
        'image_primary': input_images,
        'timestep_pad_mask': np.full((1, input_images.shape[1]), True, dtype=bool)
    }

    # Sample normalized actions
    actions = model.sample_actions(
        observation, 
        task, 
        unnormalization_statistics=model.dataset_statistics["bridge_dataset"]["action"], 
        rng=jax.random.PRNGKey(0)
    )
    actions = actions[0] # remove batch dim

For realistic training, use make_interleaved_dataset to combine multiple datasets with specific weights. This method is more powerful than manual shuffling as it can handle interleaved datasets, trajectory transformations (like goal relabeling and windowing), and frame transformations (like image augmentation and resizing) in a single pipeline.

Key configuration areas:

• traj_transform_kwargs: Controls trajectory-level logic like goal_relabeling_strategy, window_size, action_horizon, and subsample_length.
• frame_transform_kwargs: Controls frame-level logic like image_augment_kwargs, resize_size, and num_parallel_calls.
• shuffle_buffer_size: Since make_interleaved_dataset shuffles JPEG-encoded images, you can often use a larger buffer than manual frame shuffling.

from octo.data.oxe import make_oxe_dataset_kwargs_and_weights
from octo.data.dataset import make_interleaved_dataset

dataset_kwargs_list, sample_weights = make_oxe_dataset_kwargs_and_weights(
    "rtx",
    "gs://gresearch/robotics",
    load_camera_views=("primary", "wrist"),
)

dataset = make_interleaved_dataset(
    dataset_kwargs_list,
    sample_weights,
    train=True,
    shuffle_buffer_size=1000,
    batch_size=8,
    traj_transform_kwargs=dict(
        goal_relabeling_strategy="uniform",
        window_size=2,
        action_horizon=4,
        subsample_length=100,
    ),
    frame_transform_kwargs=dict(
        image_augment_kwargs=dict(
            primary=dict(
                augment_order=["random_resized_crop", "random_brightness"],
                random_resized_crop=dict(scale=[0.8, 1.0], ratio=[0.9, 1.1]),
                random_brightness=[0.1],
            )
        ),
        resize_size=dict(
            primary=(256, 256),
            wrist=(128, 128),
        ),
        num_parallel_calls=64,
    ),
    traj_transform_threads=16,
    traj_read_threads=16,
)

iterator = dataset.iterator(prefetch=1)