Despite their stellar performance on a wide range of tasks, including in-context tasks only revealed during inference, vanilla transformers and variants trained for next-token predictions (a) do not learn an explicit world model of their environment which can be flexibly queried and (b) cannot be used for planning or navigation. In this paper, we consider partially observed environments (POEs), where an agent receives perceptually aliased observations as it navigates, which makes path planning hard. We introduce a transformer with (multiple) discrete bottleneck(s), TDB, whose latent codes learn a compressed representation of the history of observations and actions. After training a TDB to predict the future observation(s) given the history, we extract interpretable cognitive maps of the environment from its active bottleneck(s) indices. These maps are then paired with an external solver to solve (constrained) path planning problems. First, we show that a TDB trained on POEs (a) retains the near perfect predictive performance of a vanilla transformer or an LSTM while (b) solving shortest path problems exponentially faster. Second, a TDB extracts interpretable representations from text datasets, while reaching higher in-context accuracy than vanilla sequence models. Finally, in new POEs, a TDB (a) reaches near-perfect in-context accuracy, (b) learns accurate in-context cognitive maps (c) solves in-context path planning problems.

翻译：尽管Transformer及其变体在广泛任务（包括仅在推理时揭示的上下文任务）中表现出色，但用于下一词元预测训练的普通Transformer（a）并未学习其环境的显式世界模型，该模型可灵活查询，（b）也无法用于规划或导航。本文考虑部分可观测环境（POE），其中智能体在导航时会接收到感知混叠的观测，这使得路径规划变得困难。我们提出一种带有（多个）离散瓶颈的Transformer（TDB），其潜在编码学习观测与行动历史的压缩表征。在训练TDB基于历史预测未来观测后，我们从其活跃的瓶颈索引中提取环境可解释的认知地图。这些地图随后与外部求解器配对，以解决（带约束的）路径规划问题。首先，我们证明在POE上训练的TDB（a）保持了普通Transformer或LSTM近乎完美的预测性能，同时（b）以指数级速度解决最短路径问题。其次，TDB从文本数据集中提取可解释表征，同时达到比普通序列模型更高的上下文准确率。最后，在新的POE中，TDB（a）达到近乎完美的上下文准确率，（b）学习准确的上下文认知地图，（c）解决上下文路径规划问题。