041、Octo开源机器人基础模型：模块化Transformer的扩散策略实现
调试Octo的时候，我盯着终端里那行loss: nan看了整整一个下午。不是权重爆炸，不是学习率过高，是数据加载器里一个不起眼的dtype转换把图像像素值从uint8变成了float32之后忘了归一化——Transformer对输入尺度敏感得让人抓狂。这个坑让我意识到，Octo这类“基础模型”和之前玩过的单任务策略网络完全是两个物种，它的模块化设计让每个组件都能独立调试，但也意味着你得对每个接口的输入输出范围有近乎偏执的掌控。
Octo的核心思想其实很朴素：把机器人操作问题拆成“看什么”和“怎么动”两个子问题，前者交给预训练的视觉编码器，后者交给一个基于Transformer的动作解码器，中间用扩散模型把语言指令、机器人状态和视觉特征揉在一起。但朴素不等于简单，它的模块化设计让这个框架能同时支持多种机器人形态和任务描述方式，这才是它配得上“基础模型”这个称号的原因。
先看整体架构。Octo的输入有三个模态：视觉观测（通常是多个摄像头视角）、机器人本体状态（关节角度、夹爪开合度等）、以及任务描述（可以是自然语言，也可以是目标图像）。这三个模态各自走独立的编码器，视觉部分用的是预训练的ViT，但注意——它冻结了大部分层，只微调最后几层，这是为了保留视觉特征的同时让模型适应机器人数据的分布。本体状态是个简单的MLP，语言指令则用CLIP的文本编码器。这三个编码器的输出会被拼成一个token序列，送进一个标准的因果Transformer里。
但真正有意思的是动作头。Octo不用传统的回归或者分类来预测动作，而是用扩散模型。具体来说，Transformer输出的隐向量会作为条件，输入到一个轻量级的扩散解码器里，这个解码器负责从噪声中逐步去噪，最终生成一段动作序列。为什么用扩散？因为机器人动作分布往往是多模态的——同一个场景下，抓取一个杯子可以从左边抓也可以从右边抓，回归模型会取平均值导致动作模糊，扩散模型则能保留这种多峰性。
代码实现上，Octo的扩散动作头是个独立的nn.Module，核心逻辑在forward和loss两个函数里。forward负责训练时的加噪和去噪预测，loss计算预测噪声和真实噪声的MSE。这里有个关键细节：加噪过程不是一次性完成的，而是随机采样一个时间步t，然后根据t计算噪声水平。训练时，模型要预测的是“噪声”，而不是直接预测动作，推理时才从随机噪声开始，迭代去噪得到最终动作。
classDiffusionActionHead(nn.Module):def__init__(self,hidden_dim,action_dim,horizon,num_steps=100):super().__init__()self.horizon=horizon self.num_steps=num_steps# 这里用了一个轻量级的MLP作为去噪网络，输入是噪声动作+条件+时间步self.denoiser=nn.Sequential(nn.Linear(action_dim*horizon+hidden_dim+1,512),nn.GELU(),nn.Linear(512,512),nn.GELU(),nn.Linear(512,action_dim*horizon))# 预定义噪声调度，别用线性调度，效果差，用cosine调度self.noise_schedule=cosine_beta_schedule(num_steps)defforward(self,condition,actions):# condition: [B, hidden_dim], actions: [B, horizon, action_dim]B=actions.shape[0]# 随机采样时间步，这里踩过坑：t必须是float，不能是int，否则梯度传不过去t=torch.randint(0,self.num_steps,(B,),device=actions.device).float()# 获取对应的噪声水平alpha_bar=self.noise_schedule['alpha_bar'][t.long()].unsqueeze(1)# 加噪noise=torch.randn_like(actions)noisy_actions=torch.sqrt(alpha_bar).unsqueeze(-1)*actions+\ torch.sqrt(1-alpha_bar).unsqueeze(-1)*noise# 把噪声动作展平，和条件拼接noisy_flat=noisy_actions.view(B,-1)cond_expanded=condition.unsqueeze(1).expand(-1,self.horizon,-1).reshape(B,-1)t_expanded=t.unsqueeze(1)/self.num_steps# 归一化时间步，不然数值太大model_input=torch.cat([noisy_flat,cond_expanded,t_expanded],dim=1)# 预测噪声pred_noise=self.denoiser(model_input).view(B,self.horizon,-1)returnpred_noise,noise训练时，整个Octo的损失函数是扩散损失加上辅助的观测重建损失——后者是为了让Transformer的隐向量保留足够的视觉信息，不然模型会偷懒只依赖语言指令。这个辅助损失在早期训练阶段特别重要，我试过去掉它，结果模型学会了“不管看到什么，都输出一个平均动作”，损失降不下去。
推理阶段就完全不一样了。训练时是一次性预测噪声，推理时得循环去噪。从纯高斯噪声开始，逐步用模型预测的噪声去更新动作，每次更新步长由噪声调度决定。这里有个工程细节：推理时的去噪步数可以比训练时的num_steps少很多，比如训练用100步，推理用10步，效果几乎不变，但速度快了10倍。Octo官方代码里默认推理步数是10，我试过5步，动作会有点抖，但抓取成功率只掉了3个百分点。
@torch.no_grad()defsample(self,condition,num_steps=None):num_steps=num_stepsorself.num_steps B=condition.shape[0]# 从纯噪声开始，别用零初始化，扩散模型对初始噪声敏感actions=torch.randn(B,self.horizon,self.action_dim,device=condition.device)foriinreversed(range(num_steps)):t=torch.full((B,),i,device=condition.device).float()# 这里注意：推理时的时间步要除以训练时的总步数，保持尺度一致t_norm=t/self.num_steps noisy_flat=actions.view(B,-1)cond_expanded=condition.unsqueeze(1).expand(-1,self.horizon,-1).reshape(B,-1)t_expanded=t_norm.unsqueeze(1)model_input=torch.cat([noisy_flat,cond_expanded,t_expanded],dim=1)pred_noise=self.denoiser(model_input).view(B,self.horizon,-1)# 去噪更新，这里用简单的DDPM更新规则alpha_bar=self.noise_schedule['alpha_bar'][i]alpha_bar_prev=self.noise_schedule['alpha_bar'][i-1]ifi>0elsetorch.tensor(1.0)beta=1-alpha_bar/alpha_bar_prev# 计算后验均值mean=(1/torch.sqrt(alpha_bar))*(actions-(beta/torch.sqrt(1-alpha_bar))*pred_noise)ifi>0:noise=torch.randn_like(actions)actions=mean+torch.sqrt(beta)*noiseelse:actions=meanreturnactions模块化设计带来的好处是你可以单独替换任何一个组件。比如视觉编码器，Octo默认用ViT-S，但如果你处理的是高分辨率图像或者多视角输入，可以换成ViT-B甚至ViT-L，只需要保证输出维度一致。我试过把视觉编码器换成DINOv2，效果提升明显，尤其是在纹理丰富的场景里，抓取成功率从78%涨到了84%。但代价是训练时间翻倍，因为DINOv2的参数量是ViT-S的四倍。
另一个值得替换的组件是语言编码器。Octo默认用CLIP，但CLIP对动作指令的理解偏向于“物体描述”而非“动作描述”。比如“把红色杯子放到蓝色盘子里”，CLIP能理解红色杯子和蓝色盘子，但对“放”这个动作的语义编码很弱。我试过换成T5或者BERT，效果各有千秋。T5对长指令的理解更好，但推理速度慢；BERT快，但对复杂指令的理解不如T5。最终我留在了CLIP，因为Octo的预训练权重就是基于CLIP的，换编码器意味着从头训练语言对齐部分，成本太高。
训练Octo时最让我头疼的是数据混合。Octo支持多数据集联合训练，但不同数据集的机器人形态不同，动作空间维度也不同。Octo的解决方案是给每个数据集定义一个“动作适配器”——一个简单的线性层，把不同维度的动作映射到统一的隐空间。这个设计很巧妙，但实现时有个坑：不同数据集的采样权重需要精心调整，不然模型会偏向数据量大的那个数据集。我试过用均匀采样，结果模型在A数据集上表现很好，在B数据集上完全崩溃。后来改成按数据集大小平方根加权采样，平衡了很多。
还有一个容易被忽略的细节是数据增强。Octo对图像做了随机裁剪和颜色抖动，但对机器人状态数据没有做任何增强。我一开始觉得状态数据不需要增强，后来发现过拟合严重，训练集损失降到0.2，验证集还在0.8。后来给状态数据加了高斯噪声和随机缩放，验证损失降到了0.4。这个增强幅度要控制好，噪声标准差在0.01左右，缩放范围在0.95到1.05之间，太大会破坏物理约束。
部署到真实机器人上时，Octo的推理延迟是个问题。Transformer的序列长度虽然不长（通常几十个token），但扩散模型的迭代去噪过程很耗时。我在一个6自由度机械臂上测试，10步去噪需要约80毫秒，加上视觉编码和语言编码，总延迟在150毫秒左右。对于抓取任务来说这个延迟可以接受，但对于动态避障或者实时交互任务就太慢了。我试过把去噪步数降到5步，延迟降到90毫秒，但动作质量明显下降，抓取成功率从82%掉到71%。
最后说说我的经验性建议。第一，别一上来就训练完整Octo，先用官方预训练权重跑通推理流程，确认环境没问题再考虑微调。第二，微调时先冻结视觉编码器，只训练动作头和Transformer的后几层，等损失稳定后再解冻视觉编码器的最后两层。第三，扩散模型的噪声调度是玄学，cosine调度比线性调度稳定得多，但如果你发现训练不稳定，试试把num_steps从100降到50，有时候减少步数反而更稳定。第四，多数据集联合训练时，一定要监控每个数据集单独的损失，不要只看总损失，不然某个数据集可能被“淹没”。第五，部署时用TensorRT或者ONNX加速Transformer部分，扩散去噪部分可以保留PyTorch，因为它的计算图相对简单，优化空间不大。
Octo不是银弹，它更像一个精心设计的积木盒，每个模块都可以替换和调整。但正是这种模块化，让你能在不重写整个框架的情况下，针对自己的机器人形态和任务场景做定制。如果你正在做多任务机器人操作，或者想探索视觉-语言-动作模型的泛化能力，Octo是个值得投入时间的起点。但记住，基础模型只是起点，真正的价值在于你如何调整它去适配你的机器人、你的场景、你的数据。
