做决策智能和机器人学习这两年我越来越确信一件事——世界模型和策略网络真正拉开差距的地方不在模型本身有多大而在对场景的组织方式。STICA是我们在复现多物体交互控制方案时遇到的一个很有意思的架构核心思想就写在名字里先用世界模型把场景“拆成一个一个物体”再用注意力机制让策略网络“只看该看的那几个”。第一眼看到这个设计时觉得它平平无奇不就是先分组再聚焦吗可真正跑起来才发现这一步拆解把物体中心表示和注意力机制的优势全部激活了训练稳定性、泛化性和可解释性都上了一个台阶。如果你也在做机器人操作、游戏AI或者多物体交互场景下的强化学习STICA这套思路值得仔细琢磨一遍。它不依赖超大模型也不依赖海量算力靠的是结构上的合理性让模型知道场景里有哪些东西、每样东西在哪、正在发生什么变化再让策略网络有选择地消费这些信息。这篇分享没有太多花哨的理论我会把拆场景的方法、策略网络的注意力实现、复现时的参数设置和我踩过的坑都尽量写清楚。1. 为什么不直接把整幅画面交给模型1.1 整景编码的算力浪费与泛化短板最早做视觉强化学习的时候大家的惯用操作很简单把游戏画面或者机器人相机的RGB图像塞进CNN压成一个高维向量或者切成patch变成一串token再交给策略网络输出动作。这种方式在单物体、背景简单、光照稳定的环境里是能跑的一旦进入多物体交互场景问题就开始集中暴露。第一个问题是算力浪费。画面上有成百上千个patch但真正跟当前决策有关的往往只有两三个物体。策略网络要把所有patch都做一次自注意力等于把大量计算花在了背景墙、阴影、桌面上最后再费力地学会“忽略它们”。这不是能力问题是路径浪费。第二个问题是泛化能力差。CNN和patch-based token不具备天然的物体边界模型看到的是像素块和频率特征。当物体位置移动、数量变化、颜色换掉模型之前学到的特征组合往往失效。我最直观的体感是用固定画面训练一个推箱子的策略训练集里箱子永远在左边真到测试环境把箱子摆到右边策略立刻懵掉。它学的根本不是“找箱子”而是“在画面的某个固定位置输出动作”。这些问题的根源在于我们让模型在没有任何物体语义的情况下直接从像素回归到决策。中间缺了一层“场景的组织和分组”。人类看到桌子上的杂货堆会天然把杯子、笔、键盘分成不同实体再根据任务决定处理哪个。这种能力不是靠堆算力硬学出来的而是靠对场景结构的先验归纳。1.2 物体中心表示带来了什么根本性改变物体中心表示Object-centric Representation要解决的问题就是给模型一个显式的“实体感”。它把整个场景编码成一组槽slot每个槽负责描述一个物体包含这个物体的外观特征、位置、存在概率等。这一组槽就是策略网络和世界模型的统一状态接口。拿Slot Attention这类方法举例它的逻辑是输入仍然是一张特征图但经过若干次迭代绑定后输出是K个固定维度的slot向量每个slot会通过注意力权重绑定到画面里的某一个物体或概念。这个过程有点像聚类但比聚类更强的地方在于它自带位置编码和迭代修正能把“同一个物体”在不同光照、遮挡、位姿下的像素聚到同一个slot里。物体中心表示带来的最根本改变是模型终于可以“指着一个东西”说话了。决策逻辑可以从“看到像素颜色分布产生动作”升级成“看到slot 3是蓝色盒子、slot 5是抓手所以控制抓手靠近slot 3”。这种抽象粒度对后续的世界模型预测和策略学习都更友好因为物体之间的交互规律比像素之间的相关性稳定得多也更容易跨场景迁移。我在实际使用中的体会特别深一旦把场景显式拆成了物体再去加关系推理、给策略加注意力筛选都变成了顺理成章的事。因为每个slot本身语义干净注意力机制聚焦到某个slot就意味着“聚焦到某个物体”而不是聚焦到一堆意义不明的patch。1.3 STICA两段式方案的整体分工STICA聪明的地方在于它把两件性质完全不同的事拆给了两个模块。世界模型负责“拆物体”和“预测未来”。它先把当前画面编码成物体级别的slot然后基于这些slot预测下一时刻的slot状态。也就是说世界模型不是对像素做预测而是对物体做预测。这个过程既给了模型一个紧凑的场景状态又让模型在想象中推演物体如何运动、如何交互相当于在决策之前多了一层“脑内模拟”。策略网络负责“选择并决策”。它拿到世界模型输出的slot集合先根据当前任务目标算一个查询向量再用注意力机制去和所有slot做匹配挑出最该关注的几个物体最后基于这些选中的物体信息输出动作。这样策略网络不需要一开始就处理全场景而是先做一轮稀疏聚焦。这种分工背后有一个非常重要的理由世界模型和策略网络的优化目标不同。世界模型希望保留足够信息来预测物体未来的状态所以它需要尽量完整地组织场景策略网络希望过滤掉与任务无关的干扰所以它需要尽量稀疏地选择信息。把这两个目标揉在一个网络里往往互相打架但拆成两段各自的训练都变得更干净。2. 世界模型先“拆物体”Object-Centric编码的实现细节2.1 从像素到SlotSlot Attention为什么适合干这个活要让世界模型具备物体中心表示最核心的模块就是Slot Attention。它本质上是一个基于注意力机制的迭代绑定过程输入是一张视觉特征图输出是一组slot。具体拆开来看流程大概是这样的先用CNN把图像编码成特征图每个空间位置对应一个特征向量并加上位置编码。初始化K个slot向量可以直接用可学习的embedding也可以从输入特征中采样初始化。重复执行三轮左右的迭代。每一轮里每个slot都通过一个线性层生成Query输入特征通过另一个线性层生成Key和Value。计算所有slot与所有输入位置的点积注意力在slot维度上做Softmax得到每个slot对图片哪个区域的绑定权重。用绑定权重去对Value做加权求和得到每个slot的更新量。把更新量通过GRU或残差连接合并进原slot最后做一次LayerNorm。这里的关键是Softmax是在slot维度上做的而不是在空间维度上。这意味着所有slot会互相竞争图像区域一个区域只能被少数几个slot抢到从而逼着每一个slot去负责一个语义上独立的实体。如果K设得比实际物体数大一些多余的slot通常会被推到背景或者“不存在”的角落去。我复现时发现迭代次数一般设置在3次最稳妥。次数太少slot还没有完全稳定下来会出现一个物体被两个slot切碎的情况次数太多训练开销增加而且容易把所有slot都逼到同一个高激活物体上造成slot坍塌。2.2 拆完之后的世界模型怎么建模物体动态拿到一组slot之后世界模型的预测就变得非常直接。STICA的典型做法是把这K个slot当作一个序列用Transformer或者共享权重的MLP去建模它们在时间上的演化。每一步预测的不是下一帧像素而是对每个slot做状态更新。如果要考虑物体之间的交互可以在slot集合内部先做一次自注意力。这一步很像机器人桌面上“两个物体会碰撞”这种关系推理自注意力让每个slot能够看到其他slot的状态然后把信息融合进自己的下一时刻预测。这么做的好处是模型学到的交互规律与物体具体外观无关blue cube还是red cube都能套用同一套物理关系。之后世界模型可以输出两种东西。一种是直接预测下一时刻的slot向量另一种是预测一个状态差分再加到当前slot上。后者在实际中更稳因为差分预测迫使模型把精力放在“变化量”上而不是死记硬背当前状态。对于不会动的背景slot差分预测会自动趋近于零这就大大减少了无用计算。预测得到的slot序列可以用来做多步想象和规划。比如在机器人操作任务里世界模型在潜在空间里推演“抓手靠近盒子之后盒子会被推动多远”策略网络不需要每次都回到真实环境做试错而是在想象中评估多个动作序列。2.3 世界模型的训练损失与观测重建世界模型不能只靠slot本身训练因为网络很容易“偷懒”让slot编码一些没意义的信息。STICA的做法通常分两头训练。一头是重构损失。把slot向量通过一个解码器重建出原始图像或者原始图像的特征图。这逼着slot保存足够多的外观信息和空间信息否则没法还原画面。第一次看到重构效果的时候我印象很深刻模型真的能按物体边界把画面切分成几块每一块对应着一个slot的重建结果边缘虽然有点糊但物体位置和颜色都还原得基本准确。另一头是动态损失。用当前时刻的slot和动作预测下一时刻的slot然后和下一时刻的真实slot特征做MSE或者余弦相似度损失。这两股损失交替约束让slot不仅要描述得清楚还要预测得准。实际训练中可以分阶段先固定视觉编码器只训动态模块等动态误差降到一定程度再放开视觉编码器端到端微调。如果一上来就端到端联合训练经常会出现Slot Attention无法稳定收敛的问题。3. 策略网络“只看该看的那几个”注意力机制怎么落地3.1 让任务描述作为Query世界模型已经给出了K个物体slot接下来问题变成策略网络怎么知道该看哪几个一个很自然的方案是用“任务描述”作为注意力里的Query。这个任务描述可以是自然语言比如“把蓝色盒子推到黄色盒子旁边”也可以是更结构化的目标向量比如目标位置坐标加上目标物体特征。在STICA的参考实现里我们用一个小的文本编码器把任务文本转成任务向量再通过一个线性层转成Query。这样做有一个好处同一套slot集合面对不同任务时关注点完全不同。同样一个桌面场景“捡起红色杯子”和“把蓝色盒子推开”会激活完全不同的注意力权重。这种查询驱动的信息筛选机制本质上就是Transformer交叉注意力里的标准用法但把它放在物体中心slot上特别出效果因为候选集里的每个token都有明确含义。3.2 用交叉注意力从Slot里挑出关键物体具体来看策略网络把任务Query输入一个交叉注意力层让它和所有slot做匹配。过程是任务Query生成Q向量所有slot生成K和V向量计算点积注意力分数然后用Softmax归一化。在实际任务里比如桌面上有五个物体当前任务只要求拿走红色的杯子。理想情况下注意力分数会把大部分质量压在“红色杯子”这个slot上其他slot的分数趋近于零。但纯Softmax有一个毛病它不会产生真正的稀疏所有物体都分到了一点注意力结果策略网络还是会被噪声干扰。所以我更推荐在STICA的场景里用硬Top-K选择。也就是先算出每个slot的注意力分数按分数从高到低排只保留前M个slot作为策略网络的实际输入。这个Top-K的选择过程如果是离散的可以用Gumbel-Softmax或者Straight-Through Estimator来保持梯度可回传如果发现离散选择不好训可以先退化成加权求和等模型稳定了再切开练。我在实验里最喜欢把K设成2或3。场景里5个物体加上一个抓手其实真正决定动作的就是目标物体和抓手本身再加一个可能被遮挡的近邻物体。这个数量既保留了关键信息又把输入维度降得非常小策略网络基本不会过拟合。3.3 从Top-K到动作输出策略头怎么接选出的M个slot向量加上机器人自身的关节状态或抓手位置会拼成一个紧凑向量送入一个两到三层的MLP策略头输出动作或者动作分布。这里的动作空间要根据任务设定。连续控制任务通常输出均值和对数方差做成高斯策略离散动作任务直接输出分类概率。训练框架可以兼容PPO或者SAC我自己的偏好是SAC加一个稍微保守的熵系数因为STICA的潜在状态空间通常比较平滑SAC的探索效率会比PPO更稳一点。这个策略头非常轻量参数量很少。它看到的输入永远是M个语义清晰的物体slot而不是一整张满是噪声的特征图。这带来一个实际收益训练数据量可以明显少一轮。用几万步交互数据简单桌面操作任务就能看到策略接近收敛放到整景编码的CNN策略上这个数据量可能才刚刚出现一点触地动作。4. 完整复现流程与参数参考4.1 实验环境与数据准备建议先别冲真实机器人先用仿真环境把流程跑通。我用的第一个环境是模仿MetaWorld的桌面抓取任务场景里有3到5个不同颜色形状的物体一个机械臂抓手任务随机指定其中一个物体作为目标。视觉输入是128x128的RGB图像机械臂状态单独做一个低维向量。数据收集分成两个阶段第一阶段用随机策略和简单的脚本策略混合探索跑大约5万步存成replay buffer用来预训练世界模型第二阶段开启在线强化学习世界模型和策略网络交替更新。仿真环境的一个好处是物体数量、颜色、位置可以脚本化生成可以很方便地构造“训练时3个物体、测试时5个物体”这样的泛化实验。这种设置最能体现STICA的优势也很适合用来观察不同模块的问题。4.2 核心代码骨架下面贴一段STICA核心流程的简化参考实现重点在于展示物体中心编码和策略注意力选择的组织方式。代码只做结构示意训练细节需要按环境调。import torch import torch.nn as nn class SlotAttention(nn.Module): def __init__(self, dim, num_slots, iters3): super().__init__() self.dim dim self.num_slots num_slots self.iters iters self.scale dim ** 0.5 self.slot_embed nn.Parameter(torch.randn(num_slots, dim)) self.q_proj nn.Linear(dim, dim) self.k_proj nn.Linear(dim, dim) self.v_proj nn.Linear(dim, dim) self.gru nn.GRUCell(dim, dim) self.norm nn.LayerNorm(dim) def forward(self, feat): # feat: B, N, D b, n, d feat.shape k self.k_proj(feat) v self.v_proj(feat) slots self.slot_embed.unsqueeze(0).expand(b, -1, -1) for _ in range(self.iters): q self.q_proj(slots) dots torch.einsum(bid,bjd-bij, q, k) * (self.scale ** -0.5) attn torch.softmax(dots, dim1) updates torch.einsum(bij,bjd-bid, attn, v) slots self.gru(updates.reshape(-1, d), slots.reshape(-1, d)).reshape(b, -1, d) slots self.norm(slots) return slots class CrossAttentionSelect(nn.Module): def __init__(self, dim, topk2): super().__init__() self.topk topk self.q_proj nn.Linear(dim, dim) self.k_proj nn.Linear(dim, dim) self.v_proj nn.Linear(dim, dim) def forward(self, task_query, slots): # task_query: B, 1, D ; slots: B, K, D q self.q_proj(task_query) k self.k_proj(slots) v self.v_proj(slots) scores torch.einsum(bqd,bkd-bqk, q, k) / (q.shape[-1] ** 0.5) attn torch.softmax(scores, dim-1) # 取 attention 最高的 topk 个 slot 做加权求和 topk_idx attn.topk(self.topk, dim-1).indices gathered torch.gather(slots.unsqueeze(1).expand(-1, 1, -1, -1), 2, topk_idx.unsqueeze(-1).expand(-1, 1, -1, slots.shape[-1])) gathered gathered.squeeze(1) # B, topk, D return gathered, attn整个流程很直观视觉特征先进Slot Attention得到K个物体slot世界模型在slot上做预测策略网络用任务Query做交叉注意力再从slot里选出Top-K用于动作MLP。这里没有花哨的封装但跑实验的效果出乎意料地稳。4.3 关键参数与训练节奏我复现时使用的一组比较稳定的参数如下不同环境可能需要微调。模块参数参考值说明视觉编码器backboneResNet-18输出特征图8x8后接1x1卷积降维Slot Attentionslot数K6比场景最大物体数多1到2Slot Attentionslot维度128过大会增加训练成本过小则丢失外观信息Slot Attention迭代次数3太少绑定不稳太多易collapse世界模型交互建模单层Transformer自注意力在slot之间交换信息策略选择Top-K2到3根据任务复杂度选择策略头MLP[256, 256]输入M个slot加机器人状态优化器统一AdamWlr3e-4世界模型和策略分开调度批量大小混合批次32注意平衡随机探索数据比例训练节奏上我的建议是前4万步只训练世界模型让Slot Attention先把场景拆干净。等重构损失下降到一定程度再开始训练策略网络。策略网络训练初期可以把世界模型冻结等策略有明显进步后再放开世界模型做联合微调。这种分阶段训练能显著降低两个模块互相带偏的概率。5. 复现中的高频坑位与排查方法5.1 Slot坍塌所有槽绑定到同一个物体上复现期间最折磨我的问题就是Slot Attention坍塌。现象是K个slot最终全都描述同一个物体其他物体被完全忽略重构图里只看得见一块区域。排查下来主要有三个原因第一个是学习率太高Slot Attention的迭代绑定过程非常敏感大学习率会让注意力权重快速收敛到少数槽上建议把世界模型学习率降到1e-4再试。第二个是迭代次数过多超过四次以后slot之间会过度竞争也容易坍塌我后来统一固定为三次。第三个是输入特征图分辨率太低物体在特征图上只剩一两个像素点注意力没法区分区域把特征图分辨率提到8x8以上会明显改善。另外可以通过一个简单trick缓解坍塌初始化slot的时候从输入特征中随机采样几个位置的特征做初始化而不是完全用固定可学习slot embedding。这样一开始每个slot就有了不同的“初始视野”竞争压力会小很多。5.2 策略注意力“雨露均沾”或聚焦错误另外一个频繁出现的问题是策略网络的注意力没有区分度。任务已经写成“拿红色杯子”但注意力分数给所有物体都差不多Top-K选出来的slot里全是无关物体策略自然一塌糊涂。我分析下来根因往往是任务Query编码和物体slot编码没有在同一语义空间里对齐。文本编码器输出的向量和图像slot的向量差别太大交叉注意力根本匹配不起来。解决办法是在训练初期加一个对齐损失让对应正确物体的slot向量和任务Query向量在表征空间里尽量靠近不相关物体推远。这个损失权重不用太大0.1左右就能起到引导作用。另外可以把任务Query的维度调成和slot一致甚至共享几层投影让交互更容易学到。5.3 动态预测误差滚动放大世界模型在想象中做多步预测时误差会逐步累积。前几步还算准到了第五步以后物体位置开始漂移甚至出现slot之间“穿模”这种不合理状态。经验做法是两条路并走。一是在训练动态模型时加入噪声扰动让模型学会纠正偏差而不是死记硬背确定性轨迹二是训练结束后做一个“预测-修正”循环每预测两步就从当前观测重新编码一次slot用真实观测校准残差。这样牺牲一点想象力步长但换来策略稳定性的提升在真实机器人部署时非常值得。5.4 新场景物体数量变化时怎么办STICA依赖固定slot数量如果训练时K6测试场景突然出现8个物体多出来的物体没有slot可用会被模型忽略。比较好的处理方式是训练时就把K设成略大于可能出现的最多物体数并且额外训练一个“存在性预测头”。这个头判断每个slot是否真的对应一个实体物体。策略选择物体前先过滤掉那些存在性分数很低的slot避免把背景空槽当成真物体。如果物体数量浮动区间特别大也可以考虑动态slot分配机制那就要引入一些边界检测和slot新增删除的逻辑复杂度会上一个台阶。6. 实测心得与扩展方向6.1 和常规Transformer注意力方案对比我做过一组对比实验同一套仿真环境一边用STICA一边用带标准Transformer注意力的整景编码方案。整景方案把画面切成64个patch全局做自注意力再输出动作STICA把画面拆成6个物体slot策略只用Top-2。结果最明显的差异在数据效率和泛化上。整景方案在训练集场景里也能跑赢随机策略但一旦把物体位置挪到没见过的区域策略退化非常快STICA在物体重新排列之后依旧能保持较高的成功率。原因很简单STICA的策略输入里根本不含“绝对坐标对应的背景特征”它看到的是“目标物体slot5、抓手slot1”位置只是slot内部的一个属性决策逻辑天然具备了平移不变性。另外在可解释性上STICA要好太多。整景方案你只能看最后的attention map但那个map在整个patch网格上人很难理解STICA的注意力分数是“当前策略关注的是第几个物体”可视化出来就是原图上画个框。调试时遇到策略失败我可以很快判断出是选错了物体还是动作输出错了排查效率完全不在一个量级。6.2 还能往哪些方向扩展STICA这套结构不止能在强化学习里用。我现在正在尝试的方向是把它接到多模态大模型的任务规划里让世界模型负责把物理场景拆成物体策略部分换成大模型的动作生成模块这样大模型就能基于“物体级状态”而不是像素来推理。实验下来会少很多幻觉式的规划。另外一个有意思的扩展是长期预测。物体中心表示天然适合做拼接把连续时间步的slot序列按物体ID关联起来就得到了每个物体自己的轨迹。这比直接预测整幅视频帧要轻量得多后续甚至可以加上物体粒度和物理属性的预测比如质量、摩擦力这些信息对抓取规划特别有用。在实际项目里踩过几次坑之后我的体会是STICA真正的价值不是某个模块的单独炫技而是把“场景理解”和“决策聚焦”用一条清晰的链路串了起来。世界模型不贪心专注于把场景拆对策略网络不贪心只看和任务强相关的几个物体。这种化整为零、擦亮焦点的设计在很多复杂控制任务里都值得借鉴。