STICA:对象中心世界模型如何提升强化学习决策与泛化 我看一个自动驾驶决策日志的时候发现一个特别有意思的现象模型在仿真里已经能稳稳跑完绕障任务但测试时路边多了一个气球广告牌车就开始左右摇摆。排查到底层才发现我把整帧画面直接压成一个特征向量交给了策略网络世界模型把“所有东西的混合平均”当成了场景状态策略网络根本不知道该把注意力放在哪里。这个痛点在做机器人控制、自动驾驶决策、具身智能的人身上太常见了。后来我在复现一个叫 STICA 的方法时思路被彻底打开——它让世界模型先把场景拆成一个个物体再让策略网络只挑和当前任务相关的几个物体去看整条链路一下子就干净了。这不是什么玄学而是一套可以用公式写清楚的建模思路而且它带给我的收益是实打实的样本效率提升、跨场景泛化变强、决策过程还能被拆出来解释。这篇文章我会把 STICA 的整体设计、两个阶段各自怎么实现、训练过程中的坑和消融实验结论都讲透适合正在做决策模型、世界模型、仿真环境强化学习的朋友拿去对照自己的项目。1. 从“整图打平”到“按物体记账”STICA针对的是哪个具体瓶颈1.1 一个全局特征向量装不下“谁在哪、谁在动”大多数世界模型的做法是把一帧图像通过编码器压成一个潜变量 z_t再用这个 z_t 去预测下一帧的 z_{t1}最后策略网络基于 z_t 输出动作。问题在于这个 z_t 是全局的它把画面里所有信息做了一次均值化处理。想象你在马路上开车画面里有一辆即将变道的前车、一个注意不到的广告牌、一片随风吹动的树叶。对决策来说前车是信号后两者是噪声可它们在全局潜变量里的梯度贡献是混在一起的策略网络根本分不清谁重要。我实测过一个很典型的例子在同一个仿真环境里模型在没有干扰物的场景训练测试时什么干扰都不加成功率 92%但只要在画面角落放一个闪烁的霓虹灯牌成功率直接掉到 61%。这个掉点不是因为模型不会开车而是全局特征把“霓虹灯闪烁”当成了一种状态变化策略网络开始为了应付闪烁而改变决策。这是全局表示的天然缺陷它混淆了“场景里的所有东西”和“当前任务需要的东西”。1.2 对象中心表示把每个潜在物体单独“记账”STICA 的解法很朴素先在特征层面把场景拆成一组物体槽object slots每个槽负责记录一个潜在物体。这就像记账以前你把整个月的所有流水写在一张纸上现在你按“房租”“餐饮”“交通”分科目记每一笔变化发生在哪个科目下一目了然。具体来说对象中心表示要求编码器输出的不是单个 z_t而是一组槽位集合 S_t {s_t^1, s_t^2, ..., s_t^K}每个槽位 s_t^i 里面既包含这个物体的外观属性颜色、纹理、类别信息也包含几何位置中心点、边界框或分割掩码。槽位之间相互独立各自做状态转移。这样做的第一个好处是当场景里某个物体发生变化时只有对应槽位的特征会变不会污染其他物体第二个好处是后续策略网络可以用一个选择器直接对槽位做筛选而不是在混合特征里“盲猜”。1.3 STICA 的总体设计拆解在前聚焦在后STICA 把整个系统明确切成两段。第一段是对象中心世界模型它只负责一件事——把原始观测拆成一个个物体并对每个物体的动力学做预测。第二段才是策略网络但它不直接看原始观测也不看全部物体槽而是通过一个选择器根据当前任务目标挑出“该看的那几个”物体再把这几路特征拼起来做决策。这个设计的第一性原理是场景理解与动作决策应该解耦。让世界模型学全部物体的物理规律让策略网络只做筛选和动作映射。比如在追捕任务里世界模型需要同时建模追捕者和障碍物但策略网络只需要重点看追捕者和离自己最近的两个障碍物剩下的可以不看。STICA 的名字就是这套思路的浓缩——先拆scene decomposition再选task-conditioned selection。它不是什么复杂的网络结构创新而是在系统架构层面做了一个非常关键的责任划分。2. 第一段拆解对象中心世界模型怎么把场景拆成“一个个物体”2.1 像素分组到物体槽用注意力完成“找物体”这件事如何在没有任何监督标注的情况下把一帧图像拆成 K 个物体我复现时用的机制是 slot attention思路本质上是一种“可学习的聚类”。图像先被编码成一组像素特征比如来自 CNN 的特征图然后初始化 K 个槽位向量通过多轮交叉注意力在像素特征与槽位之间迭代。每一轮里每个像素会计算自己与每个槽位的相似度然后按照 softmax 的权重把自己“投给”几个槽位槽位再根据收到的像素信息更新自身特征。这里我贴一段我项目里实际用过的伪代码语言层面就是 PyTorch 风格# slots: (B, K, D)初始化为可学习或从观测采样 # features: (B, H*W, D) 像素特征 for _ in range(num_iter): attn torch.einsum(bnd,bmd-bnm, features, slots) # 像素-槽位相似度 attn attn / sqrt(D) attn attn.softmax(dim-1) # 每个像素对K个槽做归一化 updates torch.einsum(bnd,bnm-bmd, features, attn) # 加权聚合像素特征 slots transformer_block(slots updates) # 更新槽位核心点在于归一化是在像素维度上做的也就是说每个像素会尽量只归属到一个槽位这让槽位之间天然具备竞争关系。这个机制在最开始的时候槽位可能是随机分配到某个颜色块但经过重构损失和预测损失的约束槽位会逐渐稳定到“一个槽对应一个可分离的物体”。在环境物体边界清晰时这个机制效果很好在物体彼此遮挡严重时需要额外引入深度信息或运动信息这点后面踩坑部分再细说。2.2 世界模型推理公式从 S_t 到 S_{t1} 的完整链路有了物体槽世界模型的推理就可以写成一套具体的公式这也是我理解“世界模型推理公式”最直接的方式。给定当前观测 o_t完整的前向过程是分解S_t Encoder_obj(o_t)把观测映射为 K 个物体槽选择A_t Selector(S_t, g)根据任务意图 g 选出一个索引子集预测对每个选中的槽位做一步转移 s_{t1}^i Dyn(s_t^i, a_t, z_t^i)其中 z_t^i 是随机扰动用来建模该物体的随机性决策a_t π_θ( Agg({s_{t1}^i | i ∈ A_t}) )策略网络只看选出来的少数槽位重构o_{t1} Decoder(S_{t1})用来提供自监督信号。这里最关键的是第三步里的 Dyn 设计。我在刚开始实现时踩过的一个思维误区是把所有槽位拼成一个序列丢进 Transformer 做统一的动力学预测。后来发现这样会让槽位之间过度耦合——一个槽的输出会被另一个槽的状态干扰物体边界就被“磨”模糊了。更好的做法是让每个槽位用一个轻量 GRU 或带 Mask 的注意力模块做自回归预测只有必要时才加一个稀疏的“关系图网络”来建模物体间的交互。这个取舍决定了拆解是否真正干净。2.3 槽位数量怎么定少了混叠多了碎裂槽位 K 是对象中心模型里最难调的参数。K 设得太小两个物体会被塞进同一个槽。我见过一个场景画面里有两辆速度方向完全不同的车K 设成 3 个槽其中一个槽的重构画面里同时出现了两辆车的叠影这个槽位的动力学预测完全没法做因为两辆车的转移规律互相冲突。K 设得太大物体碎片化一辆车可能会被拆成“车头”“车尾”“轮子”三个槽每个槽之间的边界不稳定会导致后续选择器选出无关碎片。我从实际调试中得到的经验是先用“场景中同时出现的动态物体峰值数量 2”作为初始 K然后观察两步预测误差。如果误差在某个槽位上持续偏高就增加 K如果多个槽位的特征向量两两相似度长期超过阈值比如余弦相似度大于 0.85就减少 K。大部分情况下K 落在 5 到 10 之间比较合适。当然也有一些环境里物体数量本身动态变化这时可以在固定 K 的基础上增加一个“空槽”机制让某些槽位学会表示“没有物体”再配一个 mask 头输出掩码。但代价是训练难度明显上升建议先从固定 K 开始。3. 第二段聚焦策略网络“只看该看的那几个”是怎么训练出来的3.1 选择器到底学到了什么STICA 里的策略网络前面加了一个轻量选择器 Selector它的输入是所有物体槽 S_t 和任务向量 g输出是每个槽位的显著性打分。这个打分表示“在当前任务目标下这个物体对我的决策有多少影响”。为什么不能直接用全部槽位因为决策系统中真正关键的常常只有两三个物体其余的可信度不高或没必要。把所有槽都塞给策略网络会让策略网络把参数学到“如何忽略无关物体”上而且这种忽略是隐式的完全不可控。STICA 的做法把“忽略哪些物体”这件事显式化了选择器先做一轮粗筛策略网络只处理少量相关的物体。这样做的好处是策略网络参数量可以大幅缩小而且我们随时能检查选择器到底选了哪几个物体决策过程的可解释性直接拉满。我自己的一个直观感受是加了选择器之后策略网络输入维度可以砍掉一半以上。一个 12 个槽位的场景只选 3 个槽策略网络输入特征是从 12D 降到 3D但成功率反而上升了。原因很简单模型不需要再“硬扛”大量无关特征的干扰了。3.2 top-k 选择的梯度问题用 soft top-k 保住学习信号如果直接用 index 做 top-k那选择器就没法训练了因为索引选择对打分是不可导的。我的做法是采用 Gumbel-Softmax straight-through estimator 的变体前向传播阶段用 hard top-k 选出要保留的槽位索引反向传播阶段把梯度直接复制回所有槽位的打分分数。这里有段核心实现可以参考scores selector(slots, task_embedding) # (B, K) mask torch.zeros_like(scores) idx scores.topk(k, dim-1).indices mask.scatter_(-1, idx, 1.0) # 硬 mask # 训练时 soft_mask F.softmax(scores / temperature, dim-1) mask mask soft_mask - soft_mask.detach() # 前向用硬mask反向梯度从soft_mask流出temperature 的退火策略很关键。初始温度设高soft_mask 趋近均匀分布这样选择器可以探索所有槽位训练中期逐渐降低温度让选择变得更尖锐。我遇到过不降温的场景——选择器最后变成一个“软平均”跟直接看全部槽位没什么区别这就失去了 STICA 的聚焦意义。3.3 选择器偷懒的两种姿势以及对应的正则手段选择器有一个非常常见的学习陷阱整个人会“偷懒”。一种偷懒姿势是它发现只看某个固定槽位就能在训练集上拿到不错奖励于是不管任务内容是什么它都选同一个物体这种行为我称它为“选择器锁定”。另一种偷懒姿势是选择器根本无视物体槽的内容单纯把任务向量映射成一个固定选择分布这在多任务训练时尤其严重。对付第一种锁定我会在训练前中期给选择器打分加一个熵奖励项。让选择分数的分布熵不低于某个下限避免它收敛到 one-hot。注意这个正则要逐渐退火否则后期选择器永远在“乱看”干扰策略收敛。对付第二种“只依赖任务向量”的问题我采用了信息自由度正则L_reg -λ * I(global_state; selection_mask)也就是约束选择器产生的 mask 必须携带至少一定量的场景状态信息如果 mask 与所有槽位特征完全不相关就会被这个正则惩罚。这样能逼着选择器真正去看每个槽位的状态而不是翻个“任务查询表”就完事。4. 端到端还是分阶段训练方案与消融实验的答案4.1 我最终采用的“两阶段课程式训练”关于 STICA 这类模型的训练方式社区里一直有两种声音一种是所有模块一起端到端训练另一种是先把世界模型训练好再冻结再训练策略部分。我两边都试过最终项目里采用的是两阶段课程训练。第一阶段只训练对象中心世界模型。用重构损失 未来帧预测损失一起约束优化目标是让 Decoder 能从物体槽里重建出当前帧并从当前槽预测下一帧、再重建下一帧。这个阶段我不要策略网络有任何输出因为如果选择器和策略网络还没学好它们的梯度会污染世界模型的表示空间。第二阶段冻结世界模型的编码器、动力学和译码器只训练选择器 策略网络。奖励信号直接通过策略网络传回选择器。这样做的优势是世界模型先形成了稳定的物体概念选择器在“已经知道什么是物体”的基础上学习筛选两者不会互相带偏。代价是第一阶段如果拆解质量不够第二阶段再怎么调选择器也没用。所以第一阶段我会花大量时间和计算量去把重构误差和预测误差压下来。4.2 消融实验四个对照组逼出了问题的关键为了验证 STICA 的收益到底来自“拆”还是“选”我做过一组消融实验环境是一个带干扰物的连续控制的追捕场景一个智能体要去追移动靶标场景里同时有 3 个完全不相关的干扰物体颜色不同、运动模式不同。对照组设计如下方法世界模型策略输入成功率无干扰成功率有干扰A. 全局潜变量基线单 z_t 预测全局 z_t89.2%64.7%B. 对象中心看全部槽对象中心融合所有槽93.5%79.1%C. 对象中心随机选槽对象中心随机 3 个槽82.4%73.6%D. 完整 STICA对象中心选择器 Top-394.3%91.8%从这个结果能看出两件事。第一对象中心表示确实带来了泛化收益B 组比 A 组在干扰场景下高出 14 个百分点第二无脑随机选 3 个槽是严重负收益即使有对象中心支撑C 组成功率还是低于 A 组说明聚焦这件事必须“选得对”盲选不如全看。真正的跳变发生在 D 组同样的输入规模只保留 3 个槽因为选择器选得准干扰场景下成功率从 79.1% 提升到了 91.8%。4.3 结果分析里最值得关注的三个数如果只看上面的表格还不够我还发现另外三个更值得关注的数据。一是训练样本效率D 组只需要 A 组大约 40% 的交互步数就达到 85% 成功率。二是策略网络参数量D 组策略网络只需要约 50 万参数A 组用了 150 万但最终性能仍然低于 D 组这说明此前很大一部分参数学到的其实是如何忽略干扰而不是如何决策。三是可解释性我可以直接打印出选择器选出的槽位然后发现模型在接近目标时总是优先选择“追捕目标”而在躲避障碍时优先选择“最近障碍物”这种行为模式几乎与直觉判断一致。不过这里要泼一盆冷水STICA 在物体很少、背景极简的环境里不一定比全局特征模型好。比如一个桌面上只有一个目标球和一个机械臂的抓取任务总共就两个物体拆不拆都无所谓全局特征直接预测反而更快更稳。STICA 的收益有边界它更适合物体数量中等4 到 20 个、且物体间存在明显运动差异的场景。5. 复现 STICA 翻过的车四条最容易踩的坑5.1 槽位注意力分配翻转前一帧还是车后一帧变成树这是我在复现时遇到最恶心的问题。物体槽的编号在帧与帧之间没有对齐约束前一帧 3 号槽编码的是车后一帧 3 号槽编码的可能变成了树因为重构损失只关心“物体重建得像不像”不关心“同一个物体是否一直在同一个槽里”。这个不变量一旦被打破槽位的动力学预测 Dyn 就是学废的——它上一帧在预测车的轨迹下一帧要去预测树的飘动。我的解决办法有两个组合使用效果最好。一是给槽位加一个显式的位置编码让每个槽在空间语义上有偏置减少随机交换概率二是加一个时序一致性损失强制槽位特征在相邻两帧之间的匹配代价矩阵尽量接近单位阵。如果训练时间紧第二种方式更关键因为它是在损失层面直接约束。5.2 选择器过早锁死训练刚开始就只看一个物体选择器在训练早期非常脆弱它一旦发现“只看目标物体”能拿到奖励就会一直锁死在这个策略上从此不再关注障碍物和干扰物。最气人的是这种情况下单看成功率可能还挺高模型似乎“学会了”任务但场景里只要多换一个障碍物位置成功率立刻崩掉。根因是反馈信号在早期不足选择器还没见过“因为没看障碍物所以撞上了”这种反例。我的修复策略是给选择器加“探索期”训练前 2 万步选择器的打分加入均值为 0、方差从 1.0 衰减到 0.1 的高斯噪声并提高选择分数熵正则的权重逼着它多试。探索期结束后再进入精调阶段。这个技巧看起来傻实际效果极其显著尤其是任务里有多个候选物体时探索期能把选择器的“视野”撑开。5.3 编码器“偷看未来”遮挡场景下的因果泄漏在训练时我们要预测 t1 时刻的槽位但如果槽位编码器在解码重构时不小心用到了下一帧的全局特征作条件那模型就会走捷径它从下一帧里直接读取物体信息来重构当前帧槽位表示里根本没有真正的“当前状态”。这个问题在静态图像数据集里很隐蔽因为重构损失照样很低但一旦进入交互式环境模型就完全失控。检测方式是做一步遮挡测试把 t 时刻某个物体遮住如果 t1 时刻的预测仍然很准说明模型很可能在“偷看”。实现上要严格控制信息流槽位编码器只能接收当前帧及之前帧的信息不能使用全局 ResNet 的“未来视角”。我用的是因果卷积和逐层 Mask 的注意力结构确保每一组槽位特征的感受野严格截止到当前时间步。这个坑在复现时最容易被忽略因为训练时损失很漂亮到真机部署就现原形。5.4 长程预测误差累积世界模型当“预报器”还是“短时预报器”最后一个坑是对象中心世界模型做多步自回归预测时误差会越来越明显。尤其是槽位动力学里的随机扰动项 z_t^i 如果没有做重参数化或方差估计过大预测 5 步以上画面就会开始抖动。我最初想把它当长程规划器用结果效果一塌糊涂。后来我做了两处修正。一是训练阶段加入计划式采样scheduled sampling在训练时按一定概率把上一帧的真实槽位替换预测槽位减少误差累积的影响。二是推理阶段把 STICA 的使用场景限制在“短时预报器”而不是“长程世界推演器”让它预测未来 2~3 步然后配合模型预测控制MPC滚动窗口做动作规划每步重新用真实观测去更新槽位。这样既发挥了“只看该看物体”的优势又避开了自回归误差放大的问题。我自己的体会是STICA 带来的最大变化不是单点指标而是建模思路的转变世界模型不再是一个黑盒压缩器而是一个能明确区分“场景里有什么”和“任务需要什么”的系统。把选择器打印出来看看它这步选了哪几个物体为什么选它们你一眼就能看明白——这种可调试性远比提升几个百分点的成功率更有吸引力。后续如果再有人问我怎么做可解释的决策模型我大概率会先推荐从这个“先拆解、再聚焦”的框架入手。