In-Context Reinforcement Learning through Bayesian Fusion of Context and Value Prior

TL;DR

This paper tackles the challenge of in-context reinforcement learning under suboptimal pretraining by introducing SPICE, a Bayesian framework that learns a calibrated value prior from a deep ensemble and fuses it with test-time in-context evidence. The key innovation is a closed-form Bayesian fusion that yields per-action posteriors, enabling offline greedy decisions and online posterior-UCB exploration without gradient updates. The authors prove regret-optimal guarantees: logarithmic regret for stochastic bandits and minimax-optimal regret for finite-horizon MDPs, with any miscalibration from pretraining contributing only a constant warm-start term. Empirically, SPICE achieves near-optimal decisions on unseen tasks and substantially lowers regret across bandits and a sparse-reward Darkroom MDP under distribution shift, demonstrating practical adaptability and robustness. The work provides a scalable, gradient-free adaptation mechanism that leverages suboptimal historical data to improve fast, uncertainty-aware task transfer in RL.

Abstract

In-context reinforcement learning (ICRL) promises fast adaptation to unseen environments without parameter updates, but current methods either cannot improve beyond the training distribution or require near-optimal data, limiting practical adoption. We introduce SPICE, a Bayesian ICRL method that learns a prior over Q-values via deep ensemble and updates this prior at test-time using in-context information through Bayesian updates. To recover from poor priors resulting from training on sub-optimal data, our online inference follows an Upper-Confidence Bound rule that favours exploration and adaptation. We prove that SPICE achieves regret-optimal behaviour in both stochastic bandits and finite-horizon MDPs, even when pretrained only on suboptimal trajectories. We validate these findings empirically across bandit and control benchmarks. SPICE achieves near-optimal decisions on unseen tasks, substantially reduces regret compared to prior ICRL and meta-RL approaches while rapidly adapting to unseen tasks and remaining robust under distribution shift.

Figures (8)

• Figure 1: Training and Test-Time Overview. SPICE learns a causal-transformer backbone and a K-head value ensemble from offline trajectories, then performs test-time adaptation without gradients by combining the ensemble’s value prior with context-derived evidence via a closed-form Bayesian update. Circled numbers mark core contributions: ensemble prior with calibrated uncertainty, kernel-based evidence extraction from multi-episode context, and closed-form Bayesian fusion enabling greedy (offline) / posterior-UCB (online) action selection.
• Figure 2: Bandit performance evaluation.(a) Offline selection quality. (b) Online cumulative regret (zoomed view). (c) Online cumulative regret (full scale). Shaded regions are $\pm$ SEM over $N{=}200$ test environments.
• Figure 3: Robustness to reward noise. Final regret at $H{=}500$ for different noise levels ($\sigma\in\{0.0,0.3,0.5\}$). Bars are $\pm$ SEM over $N{=}200$ test environments.
• Figure 4: Darkroom (MDP) results. Models are pretrained on uniformly collected, weak‑last labeled trajectories and evaluated online on $N{=}100$ held‑out tasks for $H{=}100$ steps. Shaded regions denote $\pm$ SEM across tasks. This setup is intentionally worst-case: contexts come from a uniform random policy and labels are uninformative.
• Figure 5: Detailed test-time architecture diagram Given a query state $s_{\text{qry}}$ and a multi-episode context buffer of transitions $\{(s_t,a_t,r_t,s_{t+1})\}_{t=1}^T$, a frozen causal transformer encodes the query and context into latent features $h_{\text{qry}}$ and $\{h_t\}_{t=1}^T$. Test-time inference decomposes into three stages. (1) Prior: a $K$-head value ensemble provides a per-action value prior (ensemble mean and uncertainty), e.g., $(\bar{Q}(a), \sigma_Q^2(a))$. (2) Evidence: a kernel similarity module computes weights $\{w_t\}_{t=1}^T$ from latent similarity between $h_{\text{qry}}$ and $\{h_t\}$, and an evidence summariser aggregates the weighted context into action-wise sufficient statistics $(c_a,\tilde{y}_a)$ (pseudo-count and weighted target/return). (3) Fusion & decision: a closed-form Bayesian update via precision additivity fuses prior and evidence to produce per-action posterior parameters $(m_a^{\text{post}}, v_a^{\text{post}})$, used for offline greedy selection or online exploration via posterior-UCB to choose $a^{*}$. In the online setting, the newest transition is appended to the context buffer and the procedure repeats, enabling gradient-free adaptation driven purely by the evolving context.

Theorems & Definitions (25)

• Definition 1: SPICE posterior
• Definition 2: SPICE inference
• Theorem 1: SPICE's Regret-optimality with warm start in Bandits.
• Corollary 1: Bound of well-calibrated priors
• Corollary 2: Bound on weak priors
• Definition 3: MDP posterior and TD-based evidence
• Theorem 2: SPICE's Regret-optimality in Finite-Horizon MDPs
• Lemma 1: Bias-variance decomposition
• proof
• Lemma 2: Posterior concentration