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

2025-09-30

AI Summary

To address the poor out-of-distribution (OOD) generalization of Decision Pre-trained Transformers (DPTs), this paper proposes a lightweight intrinsic curiosity-based regularization method. Specifically, it leverages prediction error of state transitions as an intrinsic curiosity signal to explicitly encourage exploration of novel state spaces during offline pre-training. This work is the first to integrate an exploration mechanism into the DPT framework, yielding the Prediction-Powered Transformer (PPT)—a novel architecture that jointly incorporates an auxiliary reward predictor and an intrinsic curiosity regularization term, coupled with offline reinforcement learning for uncertainty-aware adaptive decision-making. Empirical evaluation on a Gaussian multi-armed bandit benchmark demonstrates that our method significantly mitigates performance degradation under OOD test conditions induced by increased reward variance; the improvement is especially pronounced when pre-training data diversity is limited.

Technology Category

• Machine Learning: Online Learning & Bandits
• Humans and AI: Human-Aware Planning and Behavior Prediction
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

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.

Problem

• Improving generalization of Decision-Pretrained Transformers beyond training data
• Addressing performance degradation in test environments with higher reward variance
• Enhancing exploration during pretraining when offline data has limited diversity

Innovation

• Auxiliary reward predictor for Decision-Pretrained Transformers
• Prediction error as intrinsic curiosity signal
• Lightweight exploration regularizer for offline pretraining