Agentic Transformers Provably Learn to Search via Reinforcement Learning
Summary
This paper theoretically studies how transformer-based policies acquire search capabilities from reinforcement learning training dynamics in a stochastic tree environment. It shows that a two-head transformer can implement depth-first search and that this mechanism emerges naturally from sparse reward signals under a depth-wise curriculum.
View Cached Full Text
Cached at: 06/02/26, 03:39 PM
# Agentic Transformers Provably Learn to Search via Reinforcement Learning Source: [https://arxiv.org/abs/2606.00183](https://arxiv.org/abs/2606.00183) [View PDF](https://arxiv.org/pdf/2606.00183) > 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\. ## Submission history From: Yuejie Chi \[[view email](https://arxiv.org/show-email/403b6d8d/2606.00183)\] **\[v1\]**Fri, 29 May 2026 14:58:03 UTC \(1,178 KB\)
Similar Articles
Neuro-Symbolic Injection of LTLf Constraints in Autoregressive Reinforcement Learning Policies
Introduces a neurosymbolic framework that injects LTLf constraints into transformer-based reinforcement learning policies via differentiable automaton representations and a logic-based loss, improving constraint satisfaction while maintaining competitive returns.
Syntax vs. Semantics: How Transformers Learn Deep Dependencies
This paper introduces a mechanistic framework analyzing transformer learning dynamics, identifying gradient starvation as a barrier to deep semantic dependencies and validating chain-of-thought strategies for effective learning.
MechRL: Reinforcement Learning Agents Perform Circuit Discovery for Mechanistic Interpretability
Proposes MechRL, a reinforcement learning approach to automate circuit discovery in transformer language models. A PPO agent trained on multiple tasks discovers attention head circuits that match known canonical circuits and generalizes to a held-out task.
@swyx: co-sign. a very handy mental framework for what kinds of learning transformers do well today, and why it runs into limi…
The article discusses a mental framework for understanding what transformers learn well and their limitations, arguing that scaling current paradigms may be inefficient compared to approaches that hypothesize and seek truth, referencing the need for adversarial world models and reinforcement learning.
Revealing Safety-Critical Scenarios for UTM via Transformer
This research paper proposes a transformer-based reinforcement learning framework to automatically generate safety-critical test scenarios for Unmanned Traffic Management (UTM) systems, achieving an 8× improvement in vulnerability discovery efficiency over expert-guided testing.