Learning Cognitive Maps from Transformer Representations for Efficient Planning in Partially Observed Environments

要約

推論中にのみ明らかにされるコンテキスト内のタスクを含む、幅広いタスクで優れたパフォーマンスを発揮するにもかかわらず、ネクスト トークン予測用にトレーニングされたバニラ トランスフォーマーとバリアントは、(a) 柔軟にクエリや実行が可能な環境の明示的な世界モデルを学習しません。
(b) 計画やナビゲーションには使用できません。
この論文では、エージェントが移動中に知覚的にエイリアス化された観測を受け取る部分観測環境 (POE) について考察します。これにより、経路計画が困難になります。
我々は、(複数の) 離散ボトルネック TDB を備えた変換器を導入します。この変換器の潜在コードは、観測とアクションの履歴の圧縮表現を学習します。
履歴を考慮して将来の観測を予測するように TDB をトレーニングした後、アクティブなボトルネックのインデックスから環境の解釈可能な認知マップを抽出します。
これらのマップは外部ソルバーと組み合わせて、(制約された) パス計画の問題を解決します。
まず、POE でトレーニングされた TDB が、(a) バニラ変換器または LSTM のほぼ完璧な予測パフォーマンスを維持しながら、(b) 最短経路問題を指数関数的に高速に解決できることを示します。
第 2 に、TDB はテキスト データセットから解釈可能な表現を抽出し、同時にバニラ シーケンス モデルよりも高いコンテキスト内精度を達成します。
最後に、新しい POE では、TDB が (a) ほぼ完璧なインコンテキスト精度に達し、(b) 正確なインコンテキスト認知マップを学習し、(c) インコンテキスト パス計画の問題を解決します。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google