In-Context Curiosity: Distilling Exploration for Decision-Pretrained Transformers on Bandit Tasks
作者: Huitao Yang, Guanting Chen
分类: cs.LG, cs.AI, cs.MA
发布日期: 2025-09-30
💡 一句话要点
提出上下文好奇心机制,增强决策预训练Transformer在Bandit任务中的泛化能力
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 决策预训练Transformer 上下文学习 好奇心驱动 离线强化学习 泛化能力
📋 核心要点
- 决策预训练Transformer(DPTs)在决策任务中表现出色,但泛化能力受限于预训练数据分布。
- 论文提出上下文好奇心机制,通过奖励预测误差作为内在好奇心信号,鼓励更广泛的探索。
- 实验表明,该方法能提升DPT在奖励方差较高环境下的鲁棒性,尤其是在预训练数据有限时。
📝 摘要(中文)
随着大型语言模型(LLMs)能力的不断增长,将其应用于决策任务的兴趣也日益浓厚。决策预训练Transformer(DPTs)是常用的方法。然而,现有的DPTs训练方法通常难以泛化到预训练数据分布之外。为了缓解这一局限性,我们提出了上下文好奇心——一种轻量级的、受探索启发的离线预训练正则化方法,并引入了预测驱动Transformer(PPT)框架。PPT通过一个辅助奖励预测器来增强DPT,使用预测误差作为内在的好奇心信号,以鼓励训练期间更广泛的探索。在高斯多臂老虎机上的概念验证实验表明,PPT表现出更好的鲁棒性:它缓和了在测试环境中奖励方差较高时DPT的性能下降,尤其是在预训练数据多样性有限的情况下。虽然离线数据的质量仍然是根本,但我们的初步结果表明,好奇心驱动的预训练为增强上下文强化学习代理的分布外泛化提供了一个有希望的方向。
🔬 方法详解
问题定义:现有决策预训练Transformer (DPTs) 在决策任务中表现良好,但其泛化能力严重依赖于预训练数据的分布。当测试环境与预训练环境存在差异,特别是奖励分布的方差增大时,DPTs 的性能会显著下降。这限制了 DPTs 在实际应用中的适用性。
核心思路:论文的核心思路是引入一种内在的好奇心机制,鼓励 DPTs 在离线预训练阶段进行更广泛的探索。通过预测环境的奖励,并利用预测误差作为好奇心信号,引导模型关注那些不确定性较高的状态和动作,从而提升其对未知环境的适应能力。
技术框架:论文提出了预测驱动Transformer (PPT) 框架。该框架在标准的 DPT 基础上增加了一个辅助的奖励预测器。PPT 的训练过程包括两个主要部分:一是利用离线数据训练 DPT,使其学习决策策略;二是训练奖励预测器,使其能够预测给定状态和动作下的奖励。预测误差被用作内在奖励,与环境提供的外在奖励结合,共同驱动 DPT 的训练。
关键创新:论文的关键创新在于将好奇心驱动的探索机制引入到离线决策预训练中。与传统的 DPT 训练方法不同,PPT 不仅仅依赖于离线数据中的经验,而是主动地探索那些模型预测不确定性较高的区域,从而提升了模型的泛化能力。
关键设计:PPT 的关键设计包括:1) 奖励预测器的网络结构,可以选择简单的线性模型或更复杂的神经网络;2) 好奇心信号的计算方式,可以使用预测误差的绝对值或平方误差;3) 内在奖励和外在奖励的权重比例,需要根据具体任务进行调整,以平衡探索和利用。
🖼️ 关键图片
📊 实验亮点
在高斯多臂老虎机实验中,PPT 在测试环境奖励方差高于预训练数据时,显著缓解了 DPT 的性能下降。具体来说,当测试环境的奖励方差是预训练数据的两倍时,PPT 的性能优于 DPT,表明其具有更强的鲁棒性和泛化能力。该结果验证了好奇心驱动的预训练能够有效提升 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.