Agentic Transformers Provably Learn to Search via Reinforcement Learning
- URL: http://arxiv.org/abs/2606.00183v1
- Date: Fri, 29 May 2026 14:58:03 GMT
- Title: Agentic Transformers Provably Learn to Search via Reinforcement Learning
- Abstract summary: Tree search is a central abstraction behind many language-agent reasoning and decision-making tasks.<n>We study this question in a $k$-ary tree environment, where an agentic transformer observes only its trajectory history through interaction.<n>We first construct a two-head transformer that implements randomized depth-first search (DFS)<n>We then analyze the training dynamics of policy gradient under a depth-wise curriculum, showing that this same DFS mechanism emerges in stages from sparse reinforcement feedback.
- Score: 80.1412817870784
- License: http://creativecommons.org/licenses/by/4.0/
- Abstract: Tree search is a central abstraction behind many language-agent reasoning and decision-making tasks: agents must explore actions, remember failures, and backtrack toward promising alternatives. Yet, we lack a theoretical understanding of how transformer-based policies acquire such search capabilities from the training dynamics of reinforcement learning (RL). We study this question in a stochastic $k$-ary tree environment, where an agentic transformer observes only its trajectory history through interaction and receives a terminal reward for reaching a hidden leaf goal node. We first construct a two-head transformer that implements randomized depth-first search (DFS): one head tracks previous actions, while the other detects failure outcomes and triggers backtracking. We then analyze the training dynamics of policy gradient under a depth-wise curriculum, showing that this same DFS mechanism emerges in stages from sparse reinforcement feedback without expert demonstrations. The resulting policy exhibits depth generalization: after training only on depth-$1$ and depth-$2$ trees, it succeeds on deeper full trees. We further show that, under imbalanced goal distributions, discounting the return leads to a ranked DFS policy that prioritizes higher-probability branches. Overall, our results identify a mechanistic normal form for transformer-based search, in which attention heads specialize and cooperate to extract decision-relevant traces from context and convert them into agentic action selection via RL training.
Related papers
- Beyond Outcome Rewards: Step-Level Self-Distilled Policy Optimization for Deep Search Agents [15.62688304085298]
Deep search agents operate over trajectories spanning dozens of steps.<n>On-policy self-distillation (OPSD) addresses this by using the model's own logits as dense token-level teachers.<n>We propose Step-Level Self-Distilled Policy Optimization (SSPO), which converts teacher-student disagreement into step-level advantage weights.
arXiv Detail & Related papers (2026-08-13T03:13:18Z) - SEED: Self-Evolving On-Policy Distillation for Agentic Reinforcement Learning [38.54413500261524]
Outcome-based reinforcement learning (RL) provides a practical optimization paradigm, but its sparse trajectory-level rewards offer limited guidance on intermediate decisions.<n>We propose SEED (SElf-Evolving On-Policy Distillation), a self-evolving framework that converts completed on-policy trajectories into training-time hindsight skills.
arXiv Detail & Related papers (2026-07-16T09:57:18Z) - RubricEM: Meta-RL with Rubric-guided Policy Decomposition beyond Verifiable Rewards [76.17893114021757]
Deep research systems that plan, search, evaluate evidence, and synthesize long-form reports lack ground-truth answers and span many tool-augmented decisions.<n>In this work, we argue that rubrics should serve not merely as final-answer evaluators, but as the shared interface that structures policy execution, judge feedback, and agent memory.<n>We introduceEM, a rubric-guided reinforcement learning framework that combines stagewise policy decomposition with reflection-based meta-policy evolution.
arXiv Detail & Related papers (2026-05-11T17:40:38Z) - Learning to Retrieve from Agent Trajectories [72.8923565916533]
We argue that retrieval models for agentic search should be trained directly from agent interaction data.<n>We introduce learning to retrieve from agent trajectories as a new training paradigm, where supervision is derived from multi-step agent interactions.<n>Our results highlight agent trajectories as a practical and scalable supervision source, pointing to a promising direction for retrieval in the era of agentic search.
arXiv Detail & Related papers (2026-03-30T17:59:02Z) - Search-R2: Enhancing Search-Integrated Reasoning via Actor-Refiner Collaboration [49.9937230730202]
We propose Search-R2, a novel Actor-Refiner collaboration framework that enhances reasoning through targeted intervention.<n>Our approach decomposes the generation process into an Actor, which produces initial reasoning trajectories.<n>We show that Search-R2 consistently outperforms strong RAG and RL-based baselines across model scales.
arXiv Detail & Related papers (2026-02-03T15:32:09Z) - AT$^2$PO: Agentic Turn-based Policy Optimization via Tree Search [19.443576967819684]
LLM agents have emerged as powerful systems for tackling multi-turn tasks by interleaving internal reasoning and external tool interactions.<n>We present AT$2$PO, a unified framework for multi-turn agentic RL that addresses three core challenges: limited exploration diversity, sparse credit assignment, and misaligned policy optimization.
arXiv Detail & Related papers (2026-01-08T09:35:49Z) - Multi-head Transformers Provably Learn Symbolic Multi-step Reasoning via Gradient Descent [66.78052387054593]
This work investigates how transformers learn to solve symbolic multi-step reasoning problems through chain-of-thought processes.<n>We analyze two intertwined tasks: a backward reasoning task, where the model outputs a path from a goal node to the root, and a more complex forward reasoning task.<n>We show that trained one-layer transformers can provably solve both tasks with generalization guarantees to unseen trees.
arXiv Detail & Related papers (2025-08-11T17:40:47Z) - RAGEN: Understanding Self-Evolution in LLM Agents via Multi-Turn Reinforcement Learning [125.96848846966087]
Training large language models (LLMs) as interactive agents presents unique challenges.<n>While reinforcement learning has enabled progress in static tasks, multi-turn agent RL training remains underexplored.<n>We propose StarPO, a general framework for trajectory-level agent RL, and introduce RAGEN, a modular system for training and evaluating LLM agents.
arXiv Detail & Related papers (2025-04-24T17:57:08Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.