In-Context Curiosity: Distilling Exploration for Decision-Pretrained Transformers on Bandit Tasks

📄 arXiv: 2510.00347v1 📥 PDF

作者: Huitao Yang, Guanting Chen

分类: cs.LG, cs.AI, cs.MA

发布日期: 2025-09-30

💡 一句话要点

提出基于预测误差的好奇心驱动的离线预训练方法，提升决策Transformer在Bandit任务中的泛化性。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 决策预训练Transformer 上下文好奇心 离线强化学习 分布外泛化 多臂老虎机

📋 核心要点

  1. 现有决策预训练Transformer(DPTs)难以泛化到预训练数据分布之外，限制了其应用范围。
  2. 论文提出上下文好奇心机制，通过预测误差作为内在奖励，鼓励模型在离线预训练阶段进行更广泛的探索。
  3. 实验表明，该方法能有效提升DPT在奖励方差较高的测试环境中的鲁棒性，尤其是在预训练数据有限的情况下。

📝 摘要（中文）

随着大型语言模型(LLMs)能力的不断增长，将其应用于决策任务的兴趣日益浓厚。一种常见的流程是决策预训练Transformer(DPTs)。然而，现有的DPTs训练方法通常难以泛化到其预训练数据分布之外。为了缓解这一限制，我们提出了上下文好奇心——一种轻量级的、受探索启发的离线预训练正则化方法——并引入了预测驱动Transformer(PPT)框架。PPT通过一个辅助奖励预测器来增强DPT，使用预测误差作为内在的好奇心信号，以鼓励训练期间更广泛的探索。在高斯多臂老虎机上的概念验证实验表明，PPT表现出更好的鲁棒性：它缓和了在测试环境中奖励方差较高时DPT中观察到的性能下降，尤其是在预训练数据多样性有限时。虽然离线数据的质量仍然是根本，但我们的初步结果表明，好奇心驱动的预训练为增强上下文强化学习代理中的分布外泛化提供了一个有希望的方向。

🔬 方法详解

问题定义：论文旨在解决决策预训练Transformer(DPTs)在面对与预训练数据分布不同的新环境时，泛化能力不足的问题。现有方法依赖于大量高质量的预训练数据，但在数据多样性有限的情况下，性能会显著下降。尤其是在奖励分布发生变化时，DPTs难以适应。

核心思路：论文的核心思路是引入“上下文好奇心”的概念，通过预测奖励的误差来驱动模型进行更广泛的探索。具体来说，模型会尝试预测每个动作的奖励，预测误差越大，说明模型对当前状态的理解越不足，因此会给予更高的内在奖励，鼓励模型探索这些未知的状态。

技术框架：论文提出了预测驱动Transformer(PPT)框架，该框架在DPT的基础上增加了一个辅助的奖励预测器。整体流程如下：1) 使用离线数据集预训练DPT；2) 训练奖励预测器，使其能够根据状态和动作预测奖励；3) 在训练过程中，计算奖励预测器的预测误差，并将其作为内在奖励添加到原始奖励中；4) 使用增强后的奖励重新训练DPT。

关键创新：论文的关键创新在于将预测误差作为一种内在的好奇心信号，并将其融入到离线预训练过程中。这种方法不需要额外的环境交互，就可以有效地提升模型的探索能力和泛化性能。与传统的探索方法相比，上下文好奇心更加轻量级，易于集成到现有的DPT框架中。

关键设计：奖励预测器可以使用任何回归模型，论文中使用了简单的神经网络。内在奖励的权重是一个重要的超参数，需要根据具体任务进行调整。损失函数包括原始奖励的损失和预测误差的损失。网络结构方面，PPT沿用了DPT的Transformer结构，并在其基础上添加了奖励预测分支。

📊 实验亮点

在高斯多臂老虎机实验中，PPT在测试环境奖励方差较高时，显著缓解了DPT的性能下降。尤其是在预训练数据多样性有限的情况下，PPT的性能提升更为明显。实验结果表明，PPT能够有效地提升DPT在分布外环境中的泛化能力，验证了上下文好奇心驱动预训练的有效性。

🎯 应用场景

该研究成果可应用于各种需要决策能力的场景，例如推荐系统、自动驾驶、机器人控制等。通过提升模型在未知环境中的泛化能力，可以减少对大量特定领域数据的依赖，降低部署成本，并提高系统的鲁棒性和适应性。未来，该方法有望扩展到更复杂的强化学习任务中。

📄 摘要（原文）

As large language models (LLMs) continue to grow in capability, there is increasing interest in incorporating them into decision-making tasks. A common pipeline for this is Decision-Pretrained Transformers (DPTs). However, existing training methods for DPTs often struggle to generalize beyond their pretraining data distribution. To explore mitigation of this limitation, we propose in-context curiosity -- a lightweight, exploration-inspired regularizer for offline pretraining -- and introduce the Prediction-Powered Transformer (PPT) framework. PPT augments DPT with an auxiliary reward predictor, using prediction error as an intrinsic curiosity signal to encourage broader exploration during training. In proof-of-concept experiments on Gaussian multi-armed bandits, PPT shows improved robustness: it moderates the performance degradation observed in DPT when test environments exhibit higher variance in reward, particularly when pretraining data has limited diversity. While the quality of offline data remain fundamental, our preliminary results suggest that curiosity-driven pretraining offers a promising direction for enhancing out-of-distribution generalization in in-context RL agents.