Context Value Informed In-Context Reinforcement Learning

arXiv:2509.23209 · cs.LG, cs.AI · Submitted 2025-09-27 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: Today's paper: "Context Value Informed In-Context Reinforcement Learning".

Jane: The gist In-Context Reinforcement Learning methods often suffer from performance degradation at testing time due to Contextual Ambiguity,

Tom: First, who's behind it and why it matters.

Paper summary: Tom: So we're looking at this paper called Context Value Informed In-Context Reinforcement Learning. It’s about how these in-context reinforcement learning methods can keep getting better when they are actually being tested on new tasks without any parameter updates.

Jane: Right, so the main thing here is that existing ICRL methods have a problem where they start to perform worse at testing time because of something called Contextual Ambiguity. This paper claims it solves that by introducing something called Context Value into the training process so the performance keeps improving steadily.

Tom: Exactly. The core idea they are pushing is that when you train these models on data that is getting better and better, they can’t actually maintain that improvement when testing because their own random actions can create a history that looks like it came from a worse policy, which causes them to make bad decisions.

Lu: That's the context of the Contextual Ambiguity problem they're looking at. They call it when a single stochastically poor action generates a context that misleadingly resembles a history from a weaker, sub-optimal policy, which leads to performance degradation during testing because of sampling randomness, which causes decision uncertainty <ref:2509.23209#pg1>.

Meng: If I'm shipping models to production and things start breaking at scale like that, it’s a huge deployment risk. So what's the fix they are proposing here?

Jane: They introduce this theoretical construct they call Context Value, which they define as the ideal performance theoretically achievable by a policy given the information in context C, meaning VC equals J(pi star C) <ref:2509.23209#pg2>.

Tom: So basically, it’s giving the model an unambiguous quality label for whatever context it's currently in, which lets the policy bypass guessing based on noisy historical interactions and instead adjust its behavior using this value signal.

Lu: They prove that this Context Value signal tightens the lower bound on how much worse the performance gap can get relative to a policy that is actually improving monotonically <ref:2509.23209#pg2>. Theorem one shows that incorporating this Context Value yields a tighter performance bound for the value-informed policy than for the baseline policy <ref:2509.23209#pg1>.

Meng: That sounds like it provides a stronger mathematical guarantee of stability, but how do you actually get this Context Value signal into the system during training? That's where I get skeptical about implementation costs.

Tom: They propose two ways to estimate this Context Value practically. First, they have CV-ICRL-phi(C), which parameterizes the context value as a function of context using an auxiliary output head in the Transformer-based ICRL model <ref:2509.23209#pg3>.

Jane: So this head gets supervised during training by the source policy return, which is J(pi t) <ref:2509.23209#pg1>. And there's a second method, CV-ICRL-phi(t), which ties the estimated Context Value to the monotonically increasing variable, timestep <ref:2509.23209#pg3>.

Paper summary: Tom: For example, in their Dark Room testbed, they use a specific function for this estimation where VC equals min t / one thousand two hundred one <ref:2509.23209#pg3>. It’s an explicit formula tied to the training progression itself.

Lu: The implication here is that they are using the structure of the training data, specifically how far along in time you are, to guide what the current context is actually worth <ref:2509.23209#pg3>. This suggests a way to maintain monotonic improvement even when testing occurs outside of the initial training distribution.

Jane: So if we look at the experimental validation, they tested this on Dark Room and Minigrid testbeds, and they found that CV-ICRL substantially improves stability in test time adaptation across various tasks Abstract.

Meng: And what are the concrete numbers showing how much better it is? I need to know what breaks when you deploy this.

Tom: They showed that CV-ICRL consistently lowers the Degradation Frequency across all tested tasks, which is a key metric for stability five point two. Specifically in Minigrid, they compared it against baseline methods like AD and IDT and found CV-ICRL achieved superior overall performance in Table one five point two.

Lu: They also provided empirical evidence on Minigrid that suggests the generalization ability of AD-like ICRL algorithms across significantly different task types is improved by this method five point two. It shows it’s not just working for one specific environment.

Jane: So, to summarize, the paper Context Value Informed In-Context Reinforcement Learning proposes using an explicit Context Value signal to give the in-context policy a quality label for its context, which prevents performance degradation during testing by maintaining a tighter bound on improvement <ref:2509.23209#pg2>.

Tom: The authors, Wenhao Zhang and Shao Zhang, are arguing that relying only on trajectory-level return feedback can lead to high-variance updates when the test time rollout is limited Abstract.

Jane: It really boils down to them addressing that Contextual Ambiguity problem by using this Value signal so the policy doesn't get misled by poor historical interactions <ref:2509.23209#pg1>.

Tom: And if you only listen to the main idea, it’s about adding a quality metric for context to stabilize performance and keep it improving during deployment Abstract.

Lu: The implication is that future work could explore how this Context Value evolves dynamically, maybe something more complex than just a timestep variable or an auxiliary head.

Meng: For production purposes, the cost of running this with the auxiliary head seems manageable if we can keep the computational overhead low during inference on real traffic <ref:2509.23209#pg3>.

Jane: So for someone just listening, it means they're looking at a new way to train AI agents so that when you take them out into the wild on new tasks, their performance doesn't suddenly drop because of how they interpreted past data Abstract.

Conclusion: Tom: So we're wrapping up on Context Value Informed In-Context Reinforcement Learning, or CV-ICRL, and what these authors actually did here is introduce a specific way to give that in-context policy a quality label for its current context.

Jane: Right, so we started by talking about how in-context reinforcement learning methods can get unstable when testing on new tasks because they suffer from what they call contextual ambiguity. This paper shows how adding this Context Value helps keep the model making better decisions during that test time rollout.

Lu: What's interesting here is that the authors introduce this Context Value as a theoretical construct, saying it represents the ideal performance achievable given all the information in context C.

Meng: And what makes me look at it is that they actually provide two ways to estimate this value practically, one using an output head in the model and another tying it to time itself. That’s a concrete mechanism we can actually build something on.

Tom: Exactly, and the experiments on things like the Dark Room and Minigrid show that this approach substantially improves stability across all those environments. We're seeing lower degradation frequency in their tests for every task they ran.

Jane: So, what does this mean for someone just listening who isn't diving into the math? It means we’re moving away from just hoping the model is right and instead giving it a better idea of what success actually looks like in its current situation.

Lu: The theoretical guarantee they provide is that incorporating this Context Value yields a tighter performance bound for the value-informed policy than for a baseline policy. It shows mathematically how much better this signal is at constraining the potential error.

Tom: I'm really interested in how these authors connect it back to the actor-critic principle, using trajectory return feedback to create these lower-variance updates, which is a clean mechanical reason why their training should be more stable than what we usually see.

Jane: And that leads us to the main implication: this method provides an explicit signal for context quality, proving it can effectively alleviate the performance degradation issues we've been seeing in testing.

Wenhao Zhang, Shao Zhang, Xihuai Wang, Yang Li, Ying Wen

Shanghai Jiao Tong University

cs.LG, cs.AI

Submitted: 2025-09-27

Updated: 2026-10-14

Comments: 35 pages, 10 figures. Camera-ready version for NeurIPS 2026. Updated title, substantially revised presentation, and expanded experiments

Code: https://github.com/Bluixe/CV-ICRL

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 83/100

The gist: The gist In-Context Reinforcement Learning methods often suffer from performance degradation at testing time due to Contextual Ambiguity, which this paper addresses by introducing Context Value

Key concepts

Contextual Ambiguity: This occurs when the model's random actions create an interaction history that looks like it belongs to a poor policy from training data. This leads the model to make bad decisions early on, causing a cycle of poor action selection and performance degradation during testing.
Context Value (VC): Context Value is defined as the ideal performance achievable by a policy given the current context. It acts as an unambiguous quality label for the current interaction history, helping the policy bypass noisy historical data and provide a clear target for improvement.
CV-ICRL Framework: This framework incorporates Context Value into testing time inference. By using this value, the system can establish a tighter performance bound. This ensures that policy updates are more stable because they are guided by an explicit measure of how good the current context is relative to an ideal policy.
CV-ICRL-ϕ(C) Estimation: This practical method estimates Context Value by using an auxiliary output head in the model, trained to predict VC based on the context. This prediction is supervised during training using the source policy's return, providing a data-driven way to quantify the quality of a specific context.

Terminology

In-Context Reinforcement Learning — 2× in this paper
Reinforcement Learning — 2× in this paper · explained in 4 episodes

Summary

The gist In-Context Reinforcement Learning methods often suffer from performance degradation at testing time due to Contextual Ambiguity, which this paper addresses by introducing Context Value Informed ICRL (CV-ICRL) to maintain monotonic improvement.

Problem Identification

In-Context Reinforcement Learning (ICRL) methods trained on monotonically improving data fail to show a continue improved testing time performance because they suffer from Contextual Ambiguity, where the model’s own stochastic actions can generate an interaction history that misleadingly resembles that of a sub-optimal policy from the training data, initiating a vicious cycle of poor action selection (1). This phenomenon is observed in case studies and is attributed to sampling randomness which causes "short-term contexts may contain low-reward samples that mislead the model and induce suboptimal decisions; consequently, the model misidentifies its stage and transitions prematurely to a more advanced policy, exacerbating performance degradation" (3.3).

Context Value as a Solution

To fundamentally address Contextual Ambiguity, the paper introduces Context Value (VC) as the ideal performance theoretically achievable by a policy given the information in context C, i.e., VC = J(π∗C) (4). This construct is designed to provide an unambiguous quality label for the current context C which allows the policy to bypass inference from noisy historical interactions (4). The paper proves that this signal tightens the lower bound on the performance gap relative to an ideal, monotonically improving policy (4.1).

CV-ICRL Framework and Guarantees

Context Value Informed ICRL (CV-ICRL) is proposed to resolve ambiguity by incorporating this value into testing time inference (4.2). The theoretical framework establishes that if the policy perfectly learns the context-optimal policy, its expected return will be optimal: J(πC) = VC = J(π∗C) (4.1). Theorem 1 proves that incorporating Context Value yields a tighter performance bound, showing that the upper bound for the value-informed policy πC (·C, VC) is strictly tighter than the bound for the baseline policy πC (·C) (4.1).

Practical Estimation Methods

The paper proposes two practical algorithms to estimate Context Value:

CV-ICRL-ϕ(C): This method parameterizes VC as a function of context, implemented as an auxiliary output head in the Transformer-based ICRL model (D.3). During training, this head is supervised by the source policy return J(πt) (Algorithm 2).
CV-ICRL-ϕ(t): Motivated by the idea that Context Value should ideally increase with information gathering, this method ties estimated VC to the monotonically increasing variable, timestep (D.3). In Dark Room, it uses the function VC = min t / 1200, 1 (D.3).

Experimental Validation

Experiments conducted on the Dark Room and Minigrid testbeds demonstrate that CV-ICRL effectively mitigates performance degradation and improves overall ICRL abilities across various tasks and environments (abstract). The results show that CV-ICRL consistently lowers the Degradation Frequency across all tested tasks (5.2), leading to superior overall performance compared to baseline methods like AD, IDT, and AD-ϵ in the Minigrid environment (Table 1). Furthermore, experiments on Minigrid provide the first empirical evidence of the generalization ability of AD-like ICRL algorithms across significantly different tasks types (5.2).

Conclusion

CV-ICRL successfully addresses performance degradation by providing an explicit signal for context quality, proving that it can effectively alleviate this issue, leading to significant improvements in overall performance (5.2). The study confirms that the method improves stability and achieves gains in metrics such as average episode return (5.2).

--- Page 1 ---

CV-ICRL Framework and Guarantees

Context Value Informed ICRL (CV-ICRL) is proposed to resolve ambiguity by incorporating this value into testing time inference (4.2).

Improvements for AI systems

Bold header: Contextual Ambiguity Resolution via Context Value (CV-ICRL)

This improvement allows agents to bypass the perilous inference from noisy historical interactions by using Context Value as an explicit signal representing the ideal performance theoretically achievable by a policy given the current context. This enables the agent to maintain performance monotonicity, addressing the flaw where models exhibit severe performance regression during inference due to stochastic actions generating misleading history.

Bold header: Monotonic Performance Guarantee

The proposed Context Value tightens the lower bound on the performance gap relative to an ideal, monotonically improving policy, as proven by Theorem 1. This provides stronger guarantees of performance monotonicity during inference, ensuring that J(πCj) = VCj ≥ VCi for successive contexts.

Bold header: Task Generalization and Robustness

The CV-ICRL method demonstrates superior generalization capability across different task types in the Minigrid environment, showing the first empirical evidence of the generalization ability of AD-like ICRL algorithms across significantly different tasks types. This confirms the system's ability to perform learn-to-learn by adapting effectively to novel tasks beyond those seen in training.

Bold header: Adaptive Context Value Estimation

The system incorporates practical methods for estimating Context Value at both training and testing time, specifically proposing two algorithms: CV-ICRL-ϕ(C) and CV-ICRL-ϕ(t). This allows the policy to be conditioned not just on the current state, but on the history of recent interactions by dynamically calculating a context quality label.

Bold header: Enhanced Stability via Context Value Informed Policy

The CV-ICRL methods effectively mitigate performance degradation and improve stability, as shown by consistently lowering the Degradation Frequency across diverse tasks. This leads to higher Average Episode Return (AER) and more reliable test-time performance compared to baseline AD methods.

Abstract

In-context reinforcement learning (ICRL) has emerged as an effective paradigm for test time adaptation to unseen tasks without parameter updates. However, existing ICRL methods can exhibit brittle and unstable adaptions, and the mechanisms underlying such adaptions remain poorly understood. We provide a policy-gradient view of ICRL and argue that relying on trajectory-level return feedback can lead to high-variance updates when test time rollout is limited. Motivated by the actor--critic principle, we propose Context Value Informed ICRL (CV-ICRL), which equips the in-context policy with an explicit context value: the expected discounted return conditioned on the current state and accumulated interaction context. CV-ICRL trains a value head and writes its prediction back into the context as a value token, enabling TD-style return targets for lower-variance policy updates. Experiments on the Dark Room, Minigrid, and Procgen testbeds show that CV-ICRL substantially improves the stability of test time adaptation and achieves higher returns across tasks and environments. Code and data are available at https://github.com/Bluixe/CV-ICRL.