Octo

核心思路

设计哲学：让通用策略像 BERT 之于 NLP

Octo 的目标做第一个真正开源、轻量级、资源友好的通用机器人策略。让通用机器人策略像 BERT 之于 NLP 一样——成为可被广泛使用的、低门槛的策略初始化（policy initialization）。开发者拿到预训练权重，用少量领域数据微调，就能获得强泛化能力。

要解决的问题

已有的通用策略（如 RT-X）通常闭源、参数量大、对硬件要求高，难以被普通研究者复现和改进。Octo 把「通用策略」拉到了单卡可玩的级别。

技术方案特点

模型架构：纯 Transformer + 模块化注意力

Octo 是一个基于 Transformer 的扩散策略，核心是 OctoTransformer。它把整个输入序列组织成三个区域：

关键设计 — Blockwise Causal Attention（分块因果注意力）：

• Task tokens：只互相 attend，作为前缀
• Observation tokens：causal 地 attend 到所有 task + 同时刻及之前的 obs tokens
• Readout tokens：只读取前面的信息，不回写。不同 readout 互相独立

Transformer 规模

输入输出格式

输入（Observations） — 字典形式，历史窗口 history_window=2（当前帧 + 上一帧）：

{
    "image_primary": (batch, 2, 256, 256, 3),   # 第三人称/工作区相机
    "image_wrist":   (batch, 2, 128, 128, 3),   # 腕部相机（可选）
    "timestep_pad_mask": (batch, 2),            # True=该帧有效
}

Task（指令） — 支持三种条件方式：

• 语言指令：T5-Base 编码器（冻结），max_length=16
• 目标图像：goal image（可选）
• 多模态：语言 + 图像各 50%

输出（Actions）：(batch, action_horizon=4, action_dim=7) — 一次预测未来 4 步动作（action chunking）。

动作表达方式：Diffusion Action Head

Octo 默认使用 DiffusionActionHead：

• Transformer 只前向一次，得到每个时刻的 action embedding
• 再用一个条件扩散模型生成动作（扩散网络是带残差连接的 MLP，非 U-Net）
• cosine beta schedule，20 步扩散，推理用 DDPM 采样
• 训练目标：预测噪声 ε（MSE loss）

仓库还提供其他 head 供微调时替换：

• ContinuousActionHead / L1ActionHead：tanh 压缩 + 回归（常用于 ALOHA 双臂 14 维）
• DiscreteActionHead：动作离散化为 256 bin（类 RT-1 风格）

训练数据：Open X-Embodiment 800K

代码框架

仓库结构

octo/
├── octo/                          # 核心库
│   ├── model/
│   │   ├── octo_model.py          # ★ OctoModel：加载/保存/推理统一入口
│   │   ├── octo_module.py         # ★ OctoTransformer（架构核心）
│   │   └── components/
│   │       ├── action_heads.py    # Diffusion/Continuous/Discrete Head
│   │       ├── tokenizers.py      # Image/Language/Lowdim Tokenizer
│   │       ├── transformer.py     # BlockTransformer + 注意力规则
│   │       └── vit_encoders.py    # SmallStem16 等 ViT 编码器
│   ├── data/
│   │   ├── dataset.py             # ★ 单/多数据集构建 + 交错采样
│   │   ├── traj_transforms.py     # 轨迹级变换（窗口、action chunk）
│   │   └── oxe/                   # Open X-Embodiment 数据配置
│   └── utils/
│       └── gym_wrappers.py        # Gym 环境包装（HistoryWrapper）
├── scripts/
│   ├── train.py                   # 预训练主脚本
│   ├── finetune.py                # ★ 微调主脚本
│   └── configs/
│       ├── config.py              # 基础超参
│       └── finetune_config.py     # 微调配置（冻结策略）
├── examples/                      # ★ 6 个示例
│   ├── 01_inference_pretrained.ipynb
│   ├── 02_finetune_new_observation_action.py
│   ├── 03_eval_finetuned.py
│   └── 05_dataloading.ipynb
└── tests/                         # 调试用数据集 + 单测

核心模块

微调流程

三种冻结模式（命令行 --config=finetune_config.py:<mode>,<task>）：

命令行微调

python scripts/finetune.py \
    --config.pretrained_path=hf://rail-berkeley/octo-small-1.5 \
    --config=finetune_config.py:full,image_conditioned

Python API 微调（改观测/动作空间）

核心步骤（详见 examples/02_finetune_new_observation_action.py）：

1. OctoModel.load_pretrained() 加载预训练
2. 修改 config["model"]["observation_tokenizers"]（如删 wrist、加 proprio）
3. 修改 config["model"]["heads"]["action"]（如换成 L1ActionHead）
4. OctoModel.from_config() 重建模型 → merge_params() 合入预训练权重
5. freeze_weights(tx, ...) 冻结指定 key
6. 自定义 train_step（jax.jit + value_and_grad）跑训练

推理部署

from octo.model.octo_model import OctoModel

model = OctoModel.load_pretrained("hf://rail-berkeley/octo-small-1.5")

# 语言指令 → 动作
task = model.create_tasks(texts=["pick up the spoon"])
action = model.sample_actions(observation, task, rng=jax.random.PRNGKey(0),
    unnormalization_statistics=model.dataset_statistics["bridge_dataset"]["action"])

sample_actions 已用 @jax.jit 装饰，推理图编译后缓存（首次有编译开销）。部署到真机可用 HistoryWrapper 包装环境，自动管理历史窗口。

快速上手

安装

conda create -n octo python=3.10
conda activate octo
git clone https://github.com/octo-models/octo.git && cd octo
pip install -e .
pip install -r requirements.txt
# GPU 版（CUDA 11）
pip install --upgrade "jax[cuda11_pip]==0.4.20" \
    -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
# 验证安装
python scripts/finetune.py \
  --config.pretrained_path=hf://rail-berkeley/octo-small-1.5 --debug

加载模型 + 推理（最短路径）

import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
import jax, numpy as np
from octo.model.octo_model import OctoModel

# 1) 加载（自动从 HuggingFace 下载）
model = OctoModel.load_pretrained("hf://rail-berkeley/octo-small-1.5")
print(model.get_pretty_spec())

# 2) 准备观测（batch=1, history=1, 256×256 RGB）
img = np.random.randint(0, 255, (1, 1, 256, 256, 3), dtype=np.uint8)
observation = {"image_primary": img, "timestep_pad_mask": np.array([[True]])}

# 3) 构造任务（语言 或 目标图像二选一）
task = model.create_tasks(texts=["pick up the fork"])
# task = model.create_tasks(goals={"image_primary": goal_img[None]})

# 4) 采样动作（需传反归一化统计量以得到真实物理动作）
action = model.sample_actions(
    observation, task,
    unnormalization_statistics=model.dataset_statistics["bridge_dataset"]["action"],
    rng=jax.random.PRNGKey(0),
)
print(action.shape)  # (1, 4, 7) → [batch, action_horizon, action_dim]

硬件需求（资源友好是最大卖点）

关键数字