Identifiable Token Correspondence for World Models

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 at https://github.com/snu-mllab/Identifiable-Token-Correspondence.

When AI agents learn to play games or navigate virtual environments, they often rely on a world simulator — an internal model that predicts what will happen next so they can practice without needing real experience. But these world simulators have a problem: objects randomly duplicate, vanish, or morph into something else from one frame to the next. This happens because the world simulator predicts each piece of the next scene independently, with no sense that the object it sees now is the same object that was there a moment ago. We fixed this by adding a simple matching step to the world simulator's prediction process. Before finalizing the next frame, our method asks: "Which pieces of this new scene are just the old pieces in slightly different places?" It then carries those pieces forward directly, only generating truly new content where the scene has genuinely changed. This is done using a classical mathematical tool called optimal transport, and it requires no changes to how the underlying world simulator is trained. The result is a more coherent imagined world, which lets agents learn faster and achieve new state-of-the-art performance across a variety of challenging benchmarks.