Reflective Context Learning: Studying the Optimization Primitives of Context Space

Nikita Vassilyev; William Berrios; Ruowang Zhang; Bo Han; Douwe Kiela; Shikib Mehri

cs.LGApr 3, 2026

Reflective Context Learning: Studying the Optimization Primitives of Context Space

Authors: Nikita Vassilyev, William Berrios, Ruowang Zhang, Bo Han, Douwe Kiela, Shikib Mehri

Organizations: Contextual AI

Abstract

Generally capable agents must learn from experience in ways that generalize across tasks and environments. The fundamental problems of learning, including credit assignment, overfitting, forgetting, local optima, and high-variance learning signals, persist whether the learned object lies in parameter space or context space. While these challenges are well understood in classical machine learning optimization, they remain underexplored in context space, leading current methods to be fragmented and ad hoc. We present Reflective Context Learning (RCL), a unified framework for agents that learn through repeated interaction, reflection on behavior and failure modes, and iterative updates to context. In RCL, reflection converts trajectories and current context into a directional update signal analogous to gradients, while mutation applies that signal to improve future behavior in context space. We recast recent context-optimization approaches as instances of this shared learning problem and systematically extend them with classical optimization primitives, including batching, improved credit-assignment signal, auxiliary losses, failure replay, and grouped rollouts for variance reduction. On AppWorld, BrowseComp+, and RewardBench2, these primitives improve over strong baselines, with their relative importance shifting across task regimes. We further analyze robustness to initialization, the effects of batch size, sampling and curriculum strategy, optimizer-state variants, and the impact of allocating stronger or weaker models to different optimization components. Our results suggest that learning through context updates should be treated not as a set of isolated algorithms, but as an optimization problem whose mechanisms can be studied systematically and improved through transferable principles.

Figures & tables

Classical Concept	RCL Analogue	Prior Work
Parameters $\theta$	Context artifact $\mathcal{C}$ (playbook, memory, tools, guidelines)	All methods below
Forward pass $\hat{y}=f_{\theta}(x)$	Trajectory $\tau=\mathcal{A}(x,\mathcal{C})$	ReAct ( Yao et al., 2023 ) ; Voyager ( Wang et al., 2023 )
Loss $\mathcal{L}(\hat{y},y^{*})$	Outcome signal $r=\mathcal{R}(\tau,y^{*})$	—
Gradient $\nabla_{\theta}\mathcal{L}$	Reflective diagnostic $\Delta=g(\tau,r,\mathcal{C})$	ProTeGi ( Pryzant et al., 2023 ) ; TextGrad ( Yuksekgonul et al., 2024 )
Optimizer step	Context update $\mathcal{C}\leftarrow f(\mathcal{C},\Delta)$	ACE ( Zhang et al., 2026 )
Minibatch $\mathcal{B}\subset\mathcal{D}$	Trajectory batch per reflection step	ProTeGi ( Pryzant et al., 2023 ) ; TF-GRPO ( Cai et al., 2025 )

Table 1: Functional correspondence between classical gradient-based learning and reflective context learning, with representative prior work that explored each context-space analogue.

Figure 1: The RCL optimization loop with primitives mapped to stages. A batch of $B$ tasks is sampled from $\mathcal{D}$ (with replay, § 3.4 ), each executed $G$ times (§ 3.1 ). Failed traces are passed — alongside dual-trace annotations (§ 3.2 ) — to a multi-head reflector (§ 3.3 ) that produces per-trace diagnostics $\Delta_{i}$ . The mutator $f$ aggregates these into structured edits to individual playbook entries $e_{i}$ , conditioned on a rolling optimization state document (§ 3.5 ).

Primitive	Stage	Pathology	Prior Work
Batching	Execution	High variance from single samples	ProTeGi; TF-GRPO
Grouped Rollouts	Execution	Confounded attribution	TF-GRPO
Credit Assignment	Reflection	Sparse terminal reward	TextGrad
Auxiliary Losses	Reflection	Surface-level diagnostics	—
Failure Replay	Sampling	Forgetting learned tactics	Dyn. Cheatsheet; ExpeL
Optimizer State	Mutation	Oscillation from stateless updates	ERM; OPRO

Table 2: Optimization primitives, the classical pathologies they address, and the RCL loop stage each targets.

Baselines
	AppWorld				BrowseComp+		RewardBench2
	Normal		Challenge		Test Set		Test Subset
	Lite	Nano	Lite	Nano	Lite	Nano	Lite	Nano
Seed	78.0	81.5	64.3	75.8	28.9	40.7	75.9	67.2
GEPA	82.7 +4.7	76.2 -5.3	62.6 -1.7	66.7 -9.1	35.3 +6.4	48.7 +8.0	79.2 +3.3	76.5 +9.3
ACE	83.3 +5.3	86.9 +5.4	69.1 +4.8	80.7 +4.9	37.3 +8.4	50.0 +9.3	79.2 +3.3	62.7 -4.5

Table 3: Main results. Each primitive is added individually to the ACE baseline. AppWorld: TGC; BrowseComp+/RewardBench2: accuracy. Lite = Gemini 3.1 Flash Lite, Nano = GPT-5.4 Nano. All runs use Claude Opus; 30 training iterations. Deltas vs. Seed.

	AppWorld				BrowseComp+		RewardBench2
	Normal		Challenge		Test Set		Test Subset
	Lite	Nano	Lite	Nano	Lite	Nano	Lite	Nano
RCL (all primitives)	89.3	89.1	71.9	83.7	37.3	51.3	82.4	69.8
$-$ Failure Replay	85.7 -3.6	88.5 -0.6	71.5 -0.4	74.9 -8.8	38.7 +1.4	33.3 -18.0	83.1 +0.7	59.4 -10.4
$-$ Optimizer State	89.9 +0.6	85.5 -3.6	71.2 -0.7	82.5 -1.2	34.0 -3.3	46.7 -4.6	74.6 -7.8	64.1 -5.7
$-$ Credit Assignment	89.9 +0.6	85.0 -4.1	73.4 +1.5	82.1 -1.6	32.7 -4.6	41.3 -10.0	74.3 -8.1	75.1 +5.3

Table 4: Leave-one-out ablation of the composed RCL optimizer. Each row removes one primitive from the full RCL composition. Deltas are relative to full RCL . Positive deltas indicate settings where removing a primitive improves over the full composition, while negative deltas indicate settings where that primitive contributes positively in composition. Largest degradation per column is underlined .

Figure 2: Learning dynamics on AppWorld dev (Gemini 3.1 Flash-Lite, 57 tasks). Solid lines: current TGC at each checkpoint. Dashed lines: recently solved rate (fraction of tasks solved $\geq$ 1 $\times$ in the trailing 5 iterations). Colored shading: active instability — the gap between the recently solved rate and current TGC, measuring tasks solved within the window but not currently retained. Gray shading: stale regressions — the gap between the all-time per-example best-so-far and the recently solved rate, measuring tasks solved historically but not within the window. Stars: peak TGC achieved during training. Green verticals: first iteration at which every dev task has been solved at least once (full coverage).

Figure 3: Design choice analysis (Gemini 3.1 Flash-Lite). (a) Seed robustness on AppWorld Challenge: RCL converges to 72–76 TGC from all three seeds; ACE without primitives diverges from weaker initializations. (b) Reflector $\times$ mutator allocation across benchmarks. Performance depends on the interaction between both roles and the task regime; no single configuration dominates uniformly. (c) Per-trace vs. batched reflection with $B{=}3$ . Batched reflection helps on harder tasks (AW Challenge, BC+) but hurts when failures are diverse.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Reflector ( $g$ )	Mutator ( $f$ )	State / Memory	Regime
Loop development
Reflexion ( Shinn et al., 2023 )	Verbal self-critique	Append to memory	Episodic history	Single-ep.
ExpeL ( Zhao et al., 2024 )	Experience extraction	Insight reuse	Extracted knowledge	Cross-task
Agent-Pro ( Zhang et al., 2024 )	Belief + policy critique	DFS-style search	World model beliefs	Policy
Dyn. Cheatsheet ( Suzgun et al., 2026 )	Failure summarization	Curated append	Persistent cheatsheet	Task learn.
ACE ( Zhang et al., 2026 )	Trajectory critique	Structured delta edits	Playbook history	Policy

Appendix

Table 5: Prior context-optimization methods as instances of the RCL loop, grouped by developmental thread.

Figure 4: Training dynamics for all primitives with a 5-iteration sliding window. Format follows Figure 2 : solid = current TGC; dashed = recently solved rate (per-example union over 5 iterations); colored shading = active instability (solved within window but not now); gray shading = stale regressions (solved historically but not within window). This is the strictest recency condition and produces the widest active instability gaps.

Figure 5: Training dynamics with a 10-iteration sliding window. Tasks solved anywhere in the last 10 iterations count as recently solvable, shifting instability from the active to the stale category. Active gaps narrow relative to the 5-iteration view (Figure 4 ), but the relative ordering of primitives is preserved.

Figure 6: Training dynamics with the all-time best-so-far envelope (no window). The dashed line is now monotonically non-decreasing and converges to 100% for all runs. All instability is active by definition (the stale category is empty). The remaining gap measures capabilities ever demonstrated but not currently retained. This is the most permissive view and produces the narrowest gaps.

Panel of Figure 7 — Figure 7: Base system prompt for AppWorld benchmark. The playbook section header is followed by dynamically generated playbook entries during evaluation.

Panel of Figure 8 — Figure 8: Base system prompt for BrowseComp benchmark. The playbook section header is followed by dynamically generated playbook entries during evaluation.

Panel of Figure 9 — Figure 9: Base system prompt for RewardBench2 benchmark. The playbook section header is followed by dynamically generated playbook entries during evaluation.

Panel of Figure 10 — Figure 10: Seed Playbook – AppWorld. Contains 9 entries across 5 sections.

Panel of Figure 11 — Figure 11: Seed Playbook – RewardBench2. Contains 6 entries across 6 sections.

Panel of Figure 12 — Figure 12: Seed Playbook (Decent) – AppWorld. Contains 7 entries across 4 sections.

Panel of Figure 13 — Figure 13: Seed Playbook (High Quality) – AppWorld. Contains 9 entries across 5 sections.

LLM Agent Self-Improvement

Large language model agent self-improvement studies how agents use experience and feedback to improve their behavior, tools, and workflows.

See all 261 papers

Feb 17, 2026

Contextual Intelligence The Next Leap for Reinforcement Learning

Reinforcement learning (RL) has produced spectacular results in games, robotics, and continuous control. Yet, despite these successes, learned policies often fail to generalize beyond their training distribution, limiting real-world impact. Recent work on contextual RL (cRL) shows that exposing agents to environment characteristics -- contexts -- can improve zero-shot transfer. So far, the community has treated context as a monolithic, static observable, an approach that constrains the generalization capabilities of RL agents. To achieve contextual intelligence we first propose a novel taxonomy of contexts that separates allogenic (environment-imposed) from autogenic (agent-driven) factors. We identify three fundamental research directions that must be addressed to promote truly contextual intelligence: (1) Learning with heterogeneous contexts to explicitly exploit the taxonomy levels so agents can reason about their influence on the world and vice versa; (2) Multi-time-scale modeling to recognize that allogenic variables evolve slowly or remain static, whereas autogenic variables may change within an episode, potentially requiring different learning mechanisms; (3) Integration of abstract, high-level contexts to incorporate roles, resource & regulatory regimes, uncertainties, and other non-physical descriptors that crucially influence behavior. We envision context as a first-class modeling primitive, empowering agents to reason about who they are, what the world permits, and how both evolve over time. By doing so, we aim to catalyze a new generation of context-aware agents that can be deployed safely and efficiently in the real world.

André Biedenkapp

Albert-Ludwigs-Universität Freiburg, Germany

Contextual RL Zero-Shot RL Reinforcement Learning Generalization in RL Transfer Learning

Ask PDF

Apr 8, 2026

CLEAR: Context Augmentation from Contrastive Learning of Experience via Agentic Reflection

Large language model agents rely on effective model context to obtain task-relevant information for decision-making. Many existing context engineering approaches primarily rely on the context generated from the past experience and retrieval mechanisms that reuse these context. However, retrieved context from past tasks must be adapted by the execution agent to fit new situations, placing additional reasoning burden on the underlying LLM. To address this limitation, we propose a generative context augmentation framework using Contrastive Learning of Experience via Agentic Reflection (CLEAR). CLEAR first employs a reflection agent to perform contrastive analysis over past execution trajectories and summarize useful context for each observed task. These summaries are then used as supervised fine-tuning data to train a context augmentation model (CAM). Then we further optimize CAM using reinforcement learning, where the reward signal is obtained by running the task execution agent. By learning to generate task-specific knowledge rather than retrieve knowledge from the past, CAM produces context that is better tailored to the current task. We conduct comprehensive evaluations on the AppWorld and WebShop benchmarks. Experimental results show that CLEAR consistently outperforms strong baselines. It improves task completion rate from 72.62% to 81.15% on AppWorld test set and averaged reward from 0.68 to 0.74 on a subset of WebShop, compared with baseline agent. Our code is publicly available at https://github.com/awslabs/CLEAR.

Linbo Liu, Guande Wu, Han Ding +7

AWS AI

LLM Context Management LLM Agents LLM Agent Training

Ask PDF

Oct 13, 2026

Contextual Credit Policy Optimization

Group-based reinforcement learning enables large language model agents to learn from sparse outcome rewards in long-horizon tasks. Recent methods derive step-level credit by comparing observation-matched experiences across trajectories. However, identical observations can occur in different contexts, and matched peers are not equally informative for a given decision. Moreover, favorable terminal outcomes can obscure redundant actions or detours within successful trajectories. To address these two challenges, we propose Contextual Credit Policy Optimization (CCPO). For the first, CCPO weights observation-matched peers from other trajectories by contextual similarity, computed from frozen LLM representations and accumulated trajectory statistics, allowing peers from similar contexts to contribute more to a step's baseline. For the second, CCPO pairs historical credit, which compares the realized return against the weighted peer baseline, with future credit, which measures the change in contextual outcome value over a short realized continuation. A step is thus credited for the progress it makes, not only for its trajectory's eventual outcome. Both credits are computed in hindsight from rollouts already collected for training, enabling CCPO to address these challenges without an auxiliary value critic. Across two Qwen2.5-Instruct backbones (1.5B and 7B), CCPO achieves up to 97.4% success on ALFWorld and 79.4% on WebShop and outperforms GRPO by up to 21.6 and 18.5 points. Ablations show that contextual weighting outperforms uniform weighting, and that historical and future credit serve as complementary learning signals.

Huiwen Han, Ziyao Zeng, Zihao Zhu +6

Stanford University · Texas A&M University · Yale University +3

RL for LLM Agents Step-Level Credit Assignment in RL RL Post-Training Sparse-Reward RL

Ask PDF

Abstract

LLM Agent Self-Improvement

Explore similar work

Contextual Intelligence The Next Leap for Reinforcement Learning

CLEAR: Context Augmentation from Contrastive Learning of Experience via Agentic Reflection

Contextual Credit Policy Optimization