让视觉世界模型的物体在长时间推理中保持一致,避免消失或错乱。
Identifiable Token Correspondence for World Models

- 通过建模跨时间的令牌对应关系,解决物体在帧间不一致的问题。
- 在Craftax-classic上达到72.5%回报和35.6%得分,显著优于此前最佳结果。
- 无需修改原模型架构,可直接部署于现有基于变换器的世界模型。
基于令牌的变换器世界模型在视觉强化学习中表现优异,但在长时程推演中常出现物体重复、消失或错位等时间不一致性问题。主要原因是现有方法将下一帧预测视为纯令牌生成任务,忽视了令牌在时间上的持续性。本文提出可识别的令牌对应关系(ITC),作为基于令牌的变换器世界模型的解码步骤,将下一帧预测建模为包含隐式令牌对应变量的结构化分配问题:每个下一帧令牌要么来自前一帧的复制,要么是新生成。ITC不改变原有变换器架构与训练流程,可直接附加于现有主干网络。实验表明,该方法在4个挑战性基准上取得当前最优性能,在Craftax-classic基准上实现72.5%的回报和35.6%的得分,显著超越此前最佳结果(67.4%与27.9%)。代码已开源至https://github.com/snu-mllab/Identifiable-Token-Correspondence。
原文摘要 · Abstract (English)
Token-based transformer world models have shown strong performance in visual reinforcement learning, but often suffer from temporal inconsistency in long-horizon rollouts, including object duplication, disappearance, and transmutation. A key reason is that most existing approaches treat next-frame prediction purely as a token generation problem, without considering the persistence of tokens across time. We introduce Identifiable Token Correspondence (ITC), a decoding step for token-based transformer world models that formulates next-frame prediction as a structured assignment problem with latent token correspondence variables: each next-frame token is explained either by copying a token from the previous frame or by generating a new one. ITC leaves the transformer architecture and training procedure unchanged and can be added on top of existing backbones. Our experiments show state-of-the-art performance on 4 challenging benchmarks. The proposed method achieves a return of 72.5% and a score of 35.6% on the Craftax-classic benchmark, significantly surpassing the previous best of 67.4% and 27.9%. We release our source code on https://github.com/snu-mllab/Identifiable-Token-Correspondence.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。