Learning Cognitive Maps from Transformer Representations for Efficient Planning in Partially Observed Environments
作者: Antoine Dedieu, Wolfgang Lehrach, Guangyao Zhou, Dileep George, Miguel Lázaro-Gredilla
分类: cs.LG, cs.AI
发布日期: 2024-01-11
💡 一句话要点
提出基于变换器的认知地图以解决部分可观测环境中的路径规划问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 变换器 认知地图 路径规划 部分可观测环境 机器学习 智能导航 压缩表示
📋 核心要点
- 现有的变换器模型无法学习可灵活查询的世界模型，限制了其在路径规划和导航中的应用。
- 本文提出的TDB通过多个离散瓶颈学习历史观察和动作的压缩表示，从而提取环境的认知地图。
- 实验结果显示，TDB在解决最短路径问题时速度显著提升，并在新环境中保持高准确性。
📝 摘要（中文）
尽管变换器在多种任务中表现优异，但其在路径规划和导航方面存在不足。本文针对部分可观测环境，提出了一种具有多个离散瓶颈的变换器（TDB），通过学习观察和动作的历史压缩表示，提取可解释的环境认知地图。实验表明，TDB在保持预测性能的同时，能够显著加快最短路径问题的求解速度，并在新环境中实现近乎完美的准确性。
🔬 方法详解
问题定义：本文旨在解决部分可观测环境中路径规划的挑战，现有变换器模型无法有效学习环境模型，导致规划困难。
核心思路：提出的TDB通过引入多个离散瓶颈，学习观察和动作的历史压缩表示，从而提取可解释的认知地图，支持路径规划。
技术框架：TDB的整体架构包括多个离散瓶颈模块，首先训练模型预测未来观察，然后从瓶颈索引中提取认知地图，最后与外部求解器结合解决路径规划问题。
关键创新：TDB的主要创新在于其能够在保持高预测性能的同时，显著加快路径规划的求解速度，这与传统的变换器模型有本质区别。
关键设计：在模型设计中，关键参数包括瓶颈的数量和结构，损失函数采用预测误差，网络结构结合了变换器的优势与瓶颈设计，确保了信息的有效压缩与提取。
🖼️ 关键图片
📊 实验亮点
实验结果显示，TDB在部分可观测环境中保持了接近完美的预测准确性，并在解决最短路径问题时速度提升了指数级，显著优于传统的变换器和LSTM模型。
🎯 应用场景
该研究的潜在应用领域包括机器人导航、自动驾驶、智能交通系统等，能够在复杂和动态环境中实现高效的路径规划，提升智能体的自主决策能力，具有重要的实际价值和未来影响。
📄 摘要（原文）
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.