theguy21/octo-base-1.5-pytorch

octo-base-1.5 (PyTorch)

A fully self-contained PyTorch conversion of rail-berkeley/octo-base-1.5 (step 300000), built on the emb-ai/octo-pytorch port.

Everything needed for inference ships in this one folder — including the frozen t5-base language encoder (202M params total, 810 MB), so no JAX checkpoint or separate T5 download is required at load time.

Numerical fidelity

Verified against the original Flax/JAX model on identical inputs:

output	max abs diff
readout_action tokens	0.0051
obs_primary tokens	0.0081
obs_wrist tokens	0.0103
task_language tokens	0.0062
diffusion eps prediction	0.0056

Differences are float32 accumulation-order noise between frameworks. Padded-timestep handling (timestep_pad_mask) is also parity-checked.

Quickstart

pip install -r requirements-inference.txt   # torch-only inference deps
pip install "octo @ git+https://github.com/emb-ai/octo-pytorch"

import torch
from huggingface_hub import snapshot_download
from octo.model.octo_model_pt import OctoModelPt

repo_dir = snapshot_download("theguy21/octo-base-1.5-pytorch")
model = OctoModelPt.load_pretrained(repo_dir)["octo_model"].cuda().eval()

# 1) build a language-conditioned task
tasks = model.create_tasks(texts=["pick up the black bowl on the stove"], device="cuda")

# 2) observations: images as float tensors (B, window=2, C, H, W), raw 0..255 scale
obs = {
    "image_primary": torch.rand(1, 2, 3, 256, 256).cuda() * 255,
    "image_wrist":   torch.rand(1, 2, 3, 128, 128).cuda() * 255,
    "timestep_pad_mask": torch.ones(1, 2, dtype=torch.bool).cuda(),
    "pad_mask_dict": {k: torch.ones(1, 2, dtype=torch.bool).cuda()
                      for k in ("image_primary", "image_wrist", "timestep")},
}

# 3) sample actions, un-normalized to bridge_dataset action space
actions = model.sample_actions(
    obs, tasks,
    unnormalization_statistics=model.dataset_statistics["bridge_dataset"]["action"],
    generator=torch.Generator("cuda").manual_seed(0),
)
print(actions.shape)  # (1, action_horizon=4, action_dim=7)

See example.py for a runnable script (includes small compatibility shims that make the port import cleanly on recent jax versions).

Notes

• Images are fed channel-first (B, W, C, H, W) as floats in [0, 255]; normalization happens inside the image tokenizers.
• window_size <= max_horizon (10); pad timesteps via timestep_pad_mask.
• The T5 encoder inside the checkpoint is frozen (as in the original release); it is only used at inference through create_tasks(texts=...).
• For fine-tuning and the full data pipeline, see the upstream repos.

Provenance & credits

• Original model & training: Octo team (rail-berkeley/octo, MIT license).
• PyTorch architecture port: emb-ai/octo-pytorch.
• Weight conversion + numerical verification: this repo (orbax → state_dict mapping via the port's FromJaxModel.load_jax_weights, then full-model round-trip re-verified against JAX golden outputs).