TreeFlex: Agents Navigating Their Own Search Trees
Dongzhuo Li, Fangzhou Cheng, Yuandong Tian and colleagues (Meta AI) introduce TreeFlex, in which a single policy model stores its intermediate states as nodes in a harness-maintained search tree and decides itself which node to resume, branch from, verify or stop at.
Ask this paper
Two roles, one model. In the Reflection role the model sees the whole tree through short node summaries and picks a node plus an instruction. In the Execution role the harness restores that node's path and the model works with normal task tools until it saves a node or proposes a solution.
No fixed traversal. There is no node scoring function and no fixed search order. The policy chooses whether to go deeper, branch, revisit, verify or stop.
Training. Navigation, task actions, checkpointing and proposal review are trained jointly with a tree-aware GRPO variant using outcome rewards plus process rewards.
Deep research. A 4B model reaches 48.8% on a deep-research benchmark versus 41.5% for flat ReAct and 42.8% for FoldAgent, under the same active-context limit, with Pass@8 about ten points higher.
SWE-bench Verified. From the same 9B base, TreeFlex resolves 54.9% versus 43.3% for FoldAgent and 45.2% for its step-wise variant.
Abstract
A ReAct-style agent accumulates observations, intermediate results, and failed attempts in a single context. As this context grows, stale and off-path information consumes its limited capacity and obscures relevant evidence, making focused reasoning increasingly difficult. We introduce TreeFlex, a framework in which the agent's own policy decides where to search. The agent records each intermediate state, with a short summary, as a node in a search tree maintained by the harness. The policy is a single language model that acts in two roles. In the Reflection role, it has a global view of the tree through these summaries and picks a node to resume, together with an instruction. In the Execution role, the harness restores the path that led to that node, and the model carries out the instruction with the usual task tools until it saves a new node or proposes a solution, which Reflection then approves or sends back to another node. Through these choices, the policy decides whether to deepen the current line, branch from or revisit an earlier state, verify a proposal, or stop, without a fixed traversal order or a score for each node. We train navigation, task actions, checkpointing, and proposal review jointly with a tree-aware adaptation of GRPO, using task-level outcome rewards augmented by process rewards. After small-scale RL of a 4B model, with the same active-context limit for all methods, TreeFlex reaches 48.8% accuracy on a challenging deep-research benchmark, versus 41.5% for flat ReAct and 42.8% for FoldAgent, and its Pass@8 is about ten points higher than either. On SWE-bench Verified, with all three systems trained with RL from the same 9B model, TreeFlex resolves 54.9% of tasks, versus 43.3% for FoldAgent and 45.2% for its step-wise variant.