Identifiable Token Correspondence for World Models
For researchers in visual reinforcement learning, this work improves long-horizon rollout consistency with a novel probabilistic formulation.
The paper addresses temporal inconsistency in transformer-based world models for visual RL by modeling token correspondence across time, achieving state-of-the-art performance with 72.5% return and 35.6% score on Craftax-classic, surpassing prior best of 67.4% and 27.9%.
Transformer-based 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 explicitly modeling correspondence between tokens across time. We formulate next-frame prediction as a structured probabilistic inference problem with latent token correspondence variables, deriving a model in which each next-frame token is explained either by copying a token from the previous frame or by generating a new token. 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.