One-shot World Models Using a Transformer Trained on a Synthetic Prior

📄 arXiv: 2409.14084v2 📥 PDF

作者: Fabio Ferreira, Moreno Schlageter, Raghu Rajan, Andre Biedenkapp, Frank Hutter

分类: cs.LG, cs.AI

发布日期: 2024-09-21 (更新: 2024-10-24)

💡 一句话要点

提出基于Transformer的One-Shot World Model，利用合成先验数据进行环境建模。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 世界模型 Transformer 上下文学习 合成数据 环境适应 强化学习 Prior-Fitted Networks

📋 核心要点

  1. 现有世界模型通常在真实环境数据上训练，泛化能力有限，难以适应新的真实环境。
  2. OSWM利用Transformer架构，通过在合成数据上进行上下文学习，实现对环境动态的快速适应。
  3. 实验表明，OSWM在简单环境中能够快速适应并训练出有效的智能体策略，但在复杂环境迁移上仍有挑战。

📝 摘要（中文）

本文提出了一种名为One-Shot World Model (OSWM) 的Transformer世界模型，该模型以上下文学习的方式，完全从先验分布中采样的合成数据中学习。先验由多个随机初始化的神经网络组成，每个网络模拟目标环境中每个状态和奖励维度的动态。通过随机掩蔽上下文位置的下一个状态和奖励，并查询OSWM以基于剩余的转换上下文进行概率预测，从而采用Prior-Fitted Networks的监督学习程序。在推理时，OSWM能够通过提供1k个转换步骤作为上下文，快速适应简单的网格世界、CartPole gym和自定义控制环境的动态，并成功训练解决环境问题的智能体策略。然而，迁移到更复杂的环境仍然是一个挑战。尽管存在这些局限性，但这项工作是完全从合成数据中学习世界模型的重要一步。

🔬 方法详解

问题定义：现有世界模型依赖于真实环境的观测数据进行训练，这限制了它们在新环境中的泛化能力。当环境发生变化时，需要重新训练模型，成本较高。因此，如何利用少量样本甚至零样本快速适应新环境是一个关键问题。

核心思路：本文的核心思路是利用合成数据训练一个通用的世界模型，使其能够通过上下文学习快速适应新的真实环境。通过构建一个包含多个随机初始化神经网络的先验分布，模拟各种可能的环境动态，从而使模型能够学习到环境动态的通用表示。

技术框架：OSWM的整体框架包括以下几个主要步骤：1) 构建合成数据先验：使用多个随机初始化的神经网络模拟环境动态；2) 数据生成：从先验分布中采样生成大量的合成数据；3) 模型训练：使用Transformer架构，在合成数据上进行上下文学习训练；4) 环境适应：在新的真实环境中，提供少量样本作为上下文，使模型快速适应环境动态；5) 策略训练：在适应后的世界模型上训练智能体策略。

关键创新：最重要的技术创新点在于利用合成数据先验进行世界模型的训练，并结合Transformer的上下文学习能力，实现了对新环境的快速适应。与传统的在真实数据上训练世界模型的方法相比，该方法具有更好的泛化能力和适应性。

关键设计：OSWM的关键设计包括：1) 使用Transformer作为世界模型的架构，利用其强大的序列建模能力；2) 构建包含多个随机初始化神经网络的先验分布，模拟各种可能的环境动态；3) 采用Prior-Fitted Networks的监督学习程序，通过随机掩蔽上下文位置的下一个状态和奖励，训练模型进行概率预测；4) 使用1k个转换步骤作为上下文，进行环境适应。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，OSWM能够通过提供1k个转换步骤作为上下文，快速适应简单的网格世界、CartPole gym和自定义控制环境的动态，并成功训练解决环境问题的智能体策略。这表明OSWM具有一定的环境适应能力，能够利用少量样本快速学习新环境的动态。

🎯 应用场景

该研究成果可应用于机器人控制、游戏AI、自动驾驶等领域。通过利用合成数据训练世界模型，可以降低对真实环境数据的依赖，加速智能体在新环境中的学习和适应过程。未来，该方法有望应用于更复杂的真实世界环境，实现更智能、更鲁棒的智能体。

📄 摘要（原文）

A World Model is a compressed spatial and temporal representation of a real world environment that allows one to train an agent or execute planning methods. However, world models are typically trained on observations from the real world environment, and they usually do not enable learning policies for other real environments. We propose One-Shot World Model (OSWM), a transformer world model that is learned in an in-context learning fashion from purely synthetic data sampled from a prior distribution. Our prior is composed of multiple randomly initialized neural networks, where each network models the dynamics of each state and reward dimension of a desired target environment. We adopt the supervised learning procedure of Prior-Fitted Networks by masking next-state and reward at random context positions and query OSWM to make probabilistic predictions based on the remaining transition context. During inference time, OSWM is able to quickly adapt to the dynamics of a simple grid world, as well as the CartPole gym and a custom control environment by providing 1k transition steps as context and is then able to successfully train environment-solving agent policies. However, transferring to more complex environments remains a challenge, currently. Despite these limitations, we see this work as an important stepping-stone in the pursuit of learning world models purely from synthetic data.