TreeFlex: Tree-Aware Policy Optimization

Why People Wanted It

In ordinary agentic reinforcement learning, a rollout is one growing token sequence. A tree search does not fit that picture: Reflection can resume a saved node, and each Execution segment starts from the context saved at its selected node. The same search therefore contains several sequences with different histories, yet only its complete outcome determines whether the task succeeded.

So the real question was never just how to score an answer. It was how to treat the whole search as one learning sample without pretending that every generated token came from one uninterrupted context.

The Core Move

Treat each complete search as one group-relative rollout, but compute each token’s learning signal in the causal block and context where that token was generated.

The source gives the reward and loss structure in prose and named quantities, but does not provide a complete displayed objective formula in the cited pages. The essential distinction is: a task reward belongs to the search; a process penalty belongs only to the turn it concerns; and restored context is input, not newly generated output to train.

Quiz

In TreeFlex, what counts as one member of a group-relative policy-optimization group?

Each generated token

Each complete search for the task

Each individual Execution segment

Each checkpoint card

$ check

Why One Complete Search Is the Rollout

Now the version for readers who want more of the machinery: group-relative comparison happens between complete searches for the same task, not between isolated segments. The old policy samples $K$ complete searches, each with a terminal score $R_i$; the group mean and standard deviation are computed from those terminal scores alone. In words, the baseline compares task outcomes without letting a local process penalty shift the baseline for other turns.

A canonical example is a research search that issues several instructions, performs multiple retrieval segments, then submits one answer. That answer receives the terminal score for the whole search, and the score applies to every generated output in it. A malformed tool call, by contrast, is a turn-level event: its penalty applies only to tokens in that turn. If all searches in the group have the same score, the normalization divisor is one, preserving penalties at their original scale. Without penalties, every generated token in a search has the same advantage.

The task score is binary in the described settings: a deep-research answer scores one if the precision-first judge accepts it, and a SWE search scores one if the official tests mark the task resolved; other outcomes score zero. Process signals are non-positive penalties. A malformed or rejected tool call receives $Q=-1$; in deep research, judged out-of-scope work can receive $Q=-0.2$. In words, these signals discourage specific turns without inventing a separate bonus for choosing a tree shape or checkpointing. What should count as a fair comparison when searches produce very different amounts of text?

Why Restored Context Creates Separate Causal Blocks

Here is the part that gets technical: a causal block is a maximal sequence in which context grows only by appending tokens, preserving the prefix from which every token was sampled. Reflection is one block because its context grows throughout the search; each Execution segment is its own block because the next segment may begin from restored context.

For a token at rollout $i$, block $b$, and position $t$, the notation $z_{i,b,t}$ names the generated token, $h_{i,b,t}$ its generating context, and $m_{i,b,t}\in\{0,1\}$ its loss mask. In words, the mask is one for tokens generated by the model in that block and zero for the prompt, restored context, harness messages, and tool observations. Instructions and summary cards copied into the other role’s context are masked there too: the instruction is trained when Reflection writes it, and the card when Execution writes it, not again when the other role reads it.

Take an Execution segment that writes a checkpoint card, followed by a later segment resumed from that checkpoint. The later segment can use the restored state as context, but the earlier segment’s tokens are not newly generated there and do not receive a second training loss. This is why each token’s log-probability is evaluated under the context from which it was generated. What happens when a search has many blocks, or Execution contributes far more tokens than Reflection?

Why Loss Aggregation Changes Which Decisions Count

Now the deeper accounting choice: PPO’s clipped surrogate is applied to generated tokens, with each token’s importance ratio evaluated under its own generating context. The weights determine whether equal influence belongs to tokens, blocks, complete rollouts, or roles; they do not change which tokens are trainable.

A token mean gives every token equal weight, so a long search contributes more. A block mean averages within each block and then across blocks, so a short Execution segment can count as much as a long one, and searches with more blocks contribute more. A rollout mean averages within each complete search and then across searches, giving each search equal total weight regardless of length or block count. A role mean instead first averages within Reflection and Execution, then across roles; mixing it with a base average uses a weight $\lambda\in[0,1]$. In words, increasing $\lambda$ gives the comparatively sparse Reflection outputs more influence against Execution’s larger token volume.

The described training setup uses rollout mean with $\lambda=0$. The cited text also describes dual clipping for negative advantages with $c=3$, but the pages do not give a complete readable displayed objective to reproduce here. The tradeoff is clear: choosing the averaging unit changes what the optimizer treats as equally important, even though the generated tokens and their causal masks stay the same. Which unit best matches what you want a search policy to learn?

Minimal Causal Code

This sketch makes the mask and context boundary visible. It is not a reproduction of the full PPO objective; the cited pages describe the clipped objective but do not provide a complete legible formula here.

# Each block stores only newly generated tokens and the context used to generate them.
for rollout in group:
    for block in rollout.blocks:
        for token, context, generated_by_model in block.steps:
            if not generated_by_model:
                continue  # prompt, restored context, harness output, tool observation

            log_prob = policy.log_prob(token, context)
            advantage = rollout.task_advantage
            if block.turn_has_process_penalty:
                advantage += block.process_penalty

            loss = clipped_surrogate(log_prob, block.old_log_prob(token), advantage)
            train(loss)

What You Gain, and What You Lose

You can assign task-level credit to the complete search while keeping token training faithful to the context that produced each output. Turn-specific penalties can target local failures without moving the group baseline, and the loss average makes explicit whether tokens, blocks, rollouts, or roles receive equal weight.

The cost is bookkeeping: you must preserve each block’s generating context, mask restored or externally supplied text, associate penalties with the right turn, and choose an aggregation rule. Different aggregation rules emphasize different behavior; they are not interchangeable shortcuts. TreeFlex’s described setup uses rollout averaging without role mixing, so it does not explicitly equalize Reflection and Execution contributions.

Quiz

A later Execution segment receives a restored context containing an earlier segment’s card. Which tokens receive training loss in the later segment?

The restored card and all later generated tokens

Only Reflection’s instructions

Every token in the search, recomputed under the latest context

Only tokens newly generated in the later segment’s causal block

$ check