Views
No views yet
rail-berkeley/octo-base-1.5
(step 300000), built on the emb-ai/octo-pytorch port.t5-base language encoder (202M params total, 810 MB), so no JAX checkpoint or
separate T5 download is required at load time.| 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 |
timestep_pad_mask) is also parity-checked.1pip install -r requirements-inference.txt # torch-only inference deps
2pip install "octo @ git+https://github.com/emb-ai/octo-pytorch"1import torch
2from huggingface_hub import snapshot_download
3from octo.model.octo_model_pt import OctoModelPt
4
5repo_dir = snapshot_download("theguy21/octo-base-1.5-pytorch")
6model = OctoModelPt.load_pretrained(repo_dir)["octo_model"].cuda().eval()
7
8# 1) build a language-conditioned task
9tasks = model.create_tasks(texts=["pick up the black bowl on the stove"], device="cuda")
10
11# 2) observations: images as float tensors (B, window=2, C, H, W), raw 0..255 scale
12obs = {
13 "image_primary": torch.rand(1, 2, 3, 256, 256).cuda() * 255,
14 "image_wrist": torch.rand(1, 2, 3, 128, 128).cuda() * 255,
15 "timestep_pad_mask": torch.ones(1, 2, dtype=torch.bool).cuda(),
16 "pad_mask_dict": {k: torch.ones(1, 2, dtype=torch.bool).cuda()
17 for k in ("image_primary", "image_wrist", "timestep")},
18}
19
20# 3) sample actions, un-normalized to bridge_dataset action space
21actions = model.sample_actions(
22 obs, tasks,
23 unnormalization_statistics=model.dataset_statistics["bridge_dataset"]["action"],
24 generator=torch.Generator("cuda").manual_seed(0),
25)
26print(actions.shape) # (1, action_horizon=4, action_dim=7)example.py for a runnable script (includes small compatibility shims that
make the port import cleanly on recent jax versions).(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.create_tasks(texts=...).FromJaxModel.load_jax_weights,
then full-model round-trip re-verified against JAX golden outputs).