OPTS-TTPO: Enhancing Finite-Sample Policy-Gradient Learning with Tree Search

Junyu Lu; Shichao Weng; Zhiqiang Wang; Haojie Luo; Jingfan Zhang; Yuhua Zhou; Cheng Du; Yuzhuo Zhang

cs.CLSep 30, 2026

OPTS-TTPO: Enhancing Finite-Sample Policy-Gradient Learning with Tree Search

Authors: Junyu Lu, Shichao Weng, Zhiqiang Wang, Haojie Luo, Jingfan Zhang, Yuhua Zhou, Cheng Du, Yuzhuo Zhang, +5 more

Organizations: Dobot Robotics · Independent Researcher · Zhejiang University · Fudan University · Osaka University

Abstract

The policy-gradient theorem gives the exact gradient under the current policy, but finite on-policy samples may miss rare high-return trajectories. We study whether tree search improves their coverage within a fixed budget while controlling gradient bias. We introduce On-Policy Parallel Tree Search (OPTS) and Tree Trajectory Policy Optimization (TTPO) using on-policy tree trajectories, which sample new suffixes from the current policy at visited states. This needs no action-distribution correction, although branching changes state visitation. Our Branch Aggregation Lemma shows that branch-weighted tree statistics recover chain expectations when branch choices and weights are fixed before outgoing transitions are sampled. OPTS selects expansion states using estimated performance differences. Under deterministic dynamics, exact values, and max-backup advantages, the induced search policy's expected return improves monotonically with the budget. We bound the gradient bias from adaptive expansion and show that max backup assigns prefix credit to actions leading to better discovered suffixes. Against a finite chain reference, TTPG's measured bias stays near its no-branching level, while NaivePG's bias grows from 0.1251 to 0.4884. At matched budgets, reward- and value-guided OPTS improve correct-answer coverage and majority-vote accuracy over independent sampling. At matched branch counts, OPTS + TTPG gains coverage with a modest bias increase relative to Fixed-branch + TTPG. Under matched interaction or rollout budgets, OPTS-TTPO improves MuJoCo tail returns over PPO by up to 28.6%, achieves a 34-22-1 win-loss-tie record against PPO on Atari-57 under the last-100-log mean-return metric, and improves micro-averaged avg@32 and pass@32 over PPO across all four Qwen3 models.

Figures & tables

Figure 1: OPTS-TTPO. OPTS selects states maximizing the length-penalized performance-difference estimate and samples on-policy suffixes from them in parallel; TTPO applies TreeGAE and branch weights to PPO updates.

Figure 2: NaivePG versus TTPG under global token-level normalization.

Figure 3: Search-budget improvements. Top: exact-value; bottom: learned critic.

Figure 4: Matched-budget coverage and accuracy. OPTS uses fixed

S_{\max}=3

; the IID curve is the matched baseline.

Figure 5: Mechanistic diagnostics. (a) Coverage gain versus relative bias difference from

s=0

. (b) Token-weighted extra prefix credit (max

-

mean) versus distance to the first downstream branch.

Figure 6: MuJoCo learning curves; tail

\Delta

is the final-100-point relative return improvement over PPO.

Variant	Hopper	Walker2d	HalfCheetah	Ant	Humanoid	Mean
Random + mean	+10.5	+5.6	-3.2	+15.0	-1.3	+5.3
Guided + mean	+8.2	+16.7	+9.5	-0.6	+6.3	+8.0
Guided + max, unweighted	+7.5	+6.8	+26.2	+8.7	+8.2	+11.5
Guided + max (full)	+11.2	+24.9	+28.6	+13.9	+15.0	+18.7

Table 1: MuJoCo component ablations. Tail-return change relative to PPO (%), averaged over ten seeds at

(\xi,S_{\max})=(0.6,1)

. Tail return averages the final 100 logged points; Mean weights tasks equally.

Metric	PPO	OPTS-TTPO	PPO	OPTS-TTPO	Tie
	Human-normalized IQM		Task wins
Full-training mean return	0.247	0.255	26	31	0
Last-100-log mean return	0.357	0.374	22	34	1

Table 2: Human-normalized IQM and task-level win counts on Atari-57. Each win compares the two methods after averaging the corresponding raw-return summary over three seeds for that game.

Methods		MATH500		MinervaMath		AMC23		AIME24		AIME25		AIME26		Macro Average		Micro Average
Methods		avg@32	pass@32	avg@32	pass@32	avg@32	pass@32	avg@32	pass@32	avg@32	pass@32	avg@32	pass@32	avg@32	pass@32	avg@32	pass@32
Qwen3-1.7B-Base
	PPO	0.6973	0.9080	0.2986	0.5515	0.4289	0.8500	0.0823	0.3333	0.0469	0.3333	0.0417	0.2667	0.2659	0.5405	0.5013	0.7384
	DAPO	0.6935	0.9100	0.2878	0.5257	0.4133	0.8750	0.0854	0.3667	0.0490	0.3000	0.0385	0.2667	0.2612	0.5407	0.4953	0.7328
	REINFORCE++	0.6877	0.9040	0.2911	0.5184	0.4016	0.8250	0.0823	0.3000	0.0521	0.3667	0.0354	0.2667	0.2584	0.5301	0.4924	0.7251
	OPTS-TTPO	0.7114	0.9220	0.2920	0.5404	0.4328	0.8000	0.1063	0.4333	0.0479	0.3000	0.0521	0.2667	0.2738	0.5437	0.5085	0.7428

Table 3: LLM performance at the synchronized step-400 checkpoint, from 32 independent responses per prompt. Macro Average weights the six benchmarks equally; Micro Average pools their 500, 272, 40, 30, 30, and 30 problems, respectively (902 in total). Column bests are bold.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Setting	MuJoCo	Atari-57
Optimizer	Adam, $\epsilon=10^{-5}$	Adam, $\epsilon=10^{-5}$
Learning rate	$3\times 10^{-4}$ , linearly annealed to 0	$2.5\times 10^{-4}$ , linearly annealed to 0
Rollout batch	$1\times 2048=2048$ transitions	$8\times 128=1024$ transitions
Minibatches / update epochs	32 (64 transitions each) / 10	4 (256 transitions each) / 4
Discount / GAE parameter	$\gamma=0.99$ , $\lambda=0.95$	$\gamma=0.99$ , $\lambda=0.95$
Policy clip / clipped value loss	0.2 / yes	0.1 / yes

Appendix

Table D.1: Optimization hyperparameters shared by PPO and OPTS-TTPO in the control experiments.

Method	Method-specific settings
PPO	4,096 prompts and rollouts per update; GAE with $\gamma=1$ and $\lambda=0.999$ ; learned critic with learning rate $10^{-5}$ and value clip 0.5; symmetric policy clip 0.2; token-level loss aggregation.
DAPO	512 prompts with 8 responses each; group-relative advantages and no critic; asymmetric policy clip $(0.2,0.28)$ with dual-clip constant 10; token-level loss aggregation; overlong-response buffer of 1,024 tokens with penalty factor 1.0.
REINFORCE++	512 prompts with 8 responses each; per-prompt mean reward baseline followed by global advantage whitening, with no critic; symmetric policy clip 0.2; token-level loss aggregation.
OPTS-TTPO	4,096 rollouts collected over the search rounds specified in Appendix D ; TreeGAE with a learned critic at learning rate $10^{-5}$ ; symmetric policy clip 0.2; branch-weighted token-level loss aggregation.

Appendix

Table D.2: Method-specific hyperparameters for LLM RLVR.

Figure E.1: Rollout scaling of OPTS across individual reasoning benchmarks. Top: reward-guided OPTS versus i.i.d. pass@k . Bottom: value-guided OPTS versus i.i.d. self-consistency. Each column reports one benchmark under matched rollout budgets

k\in\{8,16,32,64,128\}

at fixed

S_{\max}=3

Figure F.1: Learned-critic search-budget results by benchmark. Rows report reward-guided

\Delta J_{\lambda}

, reward-guided

\Delta J

, value-guided

\Delta J_{\lambda}

, and value-guided

\Delta J

, where each difference is relative to

S_{\max}=0

. Columns show the six reasoning benchmarks.

Figure G.1: Rebranching position matters. Avg@32 versus the maximum number of OPTS search rounds for performance-difference, uniformly random, and fixed-midpoint rebranching on the same trees and budgets.

Figure H.1: $\xi\times s$ grid on the MuJoCo development tasks. Top: full-training mean return; bottom: tail return. Black outline: the selected configuration $(0.6,1)$ ; white dashed outline: the best cell of each panel. The right-most column averages the min–max-normalized grids of Hopper-v4 and Humanoid-v4.

Figure I.1: Full Atari-57 learning curves for PPO and OPTS-TTPO. Each panel shows one game; curves are smoothed raw mean returns for different random seeds.

Metric	Comparison	Wins	Losses	Ties
Full training	Max vs. Mean	19	22	16
Full training	Max vs. PPO	27	30	0
Full training	Mean vs. PPO	31	26	0
Tail	Max vs. Mean	18	22	17
Tail	Max vs. PPO	34	22	1
Tail	Mean vs. PPO	34	22	1

Appendix

Table J.1: Backup-rule comparisons at fixed hyperparameters.

	Full			Tail
Game	PPO	Mean	Max	PPO	Mean	Max
ALE_Surround-v5	-2.83	-2.94	-2.94	-1.14	-0.91	-0.91
AlienNoFrameskip-v4	989.13	947.58	945.15	1186.07	1236.76	1372.65
AmidarNoFrameskip-v4	173.34	192.09	191.86	246.73	301.20	307.77
AssaultNoFrameskip-v4	922.34	967.38	968.90	1073.07	1278.37	1208.08
AsterixNoFrameskip-v4	1553.26	1671.41	1690.02	2051.14	2223.03	2194.67

Appendix

Table J.2: Per-game seed-mean returns (1), shown to two decimals. Full: all logged points; Tail: final 100 points. Win counts use unrounded values.

	Full			Tail
Game	PPO	Mean	Max	PPO	Mean	Max
KangarooNoFrameskip-v4	1229.65	1612.89	1587.39	1802.94	2458.78	2444.39
KrullNoFrameskip-v4	3762.75	3802.40	3776.63	4173.87	4482.65	4402.41
KungFuMasterNoFrameskip-v4	6102.91	5726.26	5726.26	5609.22	6604.17	6604.17
MontezumaRevengeNoFrameskip-v4	0.21	0.06	0.48	0.00	0.17	0.00
MsPacmanNoFrameskip-v4	1185.86	1131.68	1184.35	1471.88	1464.83	1551.24

Appendix

Table J.3: Per-game seed-mean returns (2), shown to two decimals. Full: all logged points; Tail: final 100 points. Win counts use unrounded values.

Figure K.1: Qwen3-1.7B per-checkpoint evaluation across six reasoning benchmarks. Each column shows one benchmark, with avg@32 above pass@32 ; all 20 checkpoints of every method are scored with the same verifier. Curves show a three-point moving average that retains the step-20 and step-400 endpoints.

Stage of one training step	Seconds	Share
Rollout generation (vLLM, all rounds)	348.1	43.7%
Critic update	214.4	26.9%
Actor update	131.4	16.5%
Critic value forwards (all rounds)	56.5	7.1%
Old log-probabilities	39.2	4.9%
Search module: $\widehat{\Delta}$ path refresh + selection	3.2	0.4%

Appendix

Table L.1: Wall-clock composition of an OPTS-TTPO step ( Qwen3-1.7B ,

S_{\max}=3

, eight H200 GPUs).

Explore similar work

EPIG-Tree: Compute-Optimal Branching for Gradient-Efficient Reinforcement Learning
Sep 17, 2026Nikita Khomich, Leopold Hermansson, Ido HakimiAutoregressive Rollout Credit Assignment
Beyond Outcome Rewards: Step-Level Self-Distilled Policy Optimization for Deep Search Agents
Aug 13, 2026Haoze Wu, Chuqiao Kuang, Tianyi Zhuang +1Search Agents Reward Signal
The Hong Kong University of Science and Technology · Huawei Technologies Ltd.
Maximizing Rollout Informativeness under a Fixed Budget: A Submodular View of Tree Search for Tool-Use Agentic Reinforcement Learning
May 6, 2026Yuelin Hu, Zhenbo Yu, Zhengxue Cheng +2Group-Based Reinforcement Learning Tree Search
Shanghai Jiao Tong University · Shanghai Maritime University

Sep 17, 2026, cs.LGJ/K move · Enter open · S save

EPIG-Tree: Compute-Optimal Branching for Gradient-Efficient Reinforcement Learning

Nikita Khomich, Leopold Hermansson, Ido Hakimi

Autoregressive Rollout Credit Assignment

Open paper AskPDF arXiv

Reward-based reinforcement learning for language models, exemplified by Group Relative Policy Optimization (GRPO), collapses an entire stochastic trajectory into a single scalar reward. This is clean and scalable, but it explores and allocates reward inefficiently: a trajectory may contain many causal decisions, recovery attempts, and environment-randomness events, yet every token or action inherits one trajectory-level advantage. We study tree-based rollout construction as a compute-allocation problem for policy-gradient estimation. Our central claim is that branches should be placed not where the policy is merely uncertain, but where an additional branch most reduces uncertainty about the policy gradient per unit of compute. From a law-of-total-variance decomposition of the local policy-gradient random variable, we derive two allocation laws: new branches reduce decision uncertainty, while repeated suffix rollouts reduce continuation uncertainty. The resulting EPIG-Tree score allocates branches using the already computed rollouts. It estimates occupancy- and score-weighted value uncertainty, along with a suffix law

n_e \propto w_e \|\nabla_θ\log π(a_e|h_e)\| σ_e / \sqrt{c_e}

. Empirically, EPIG reduces gradient MSE in cloned-state control, winning in all nine dense continuous-control environments of a 13-environment sweep and recovering the reference gradient direction near-perfectly, and it improves frozen-LLM gradient calibration relative to entropy branching. In online single-turn math, tree-local credit beats flat GRPO, while branch placement is secondary to token-level credit assignment. In online multi-turn Wordle, EPIG attains the highest final win rate (0.850), overtaking flat GRPO, which saturates early at 0.790, and entropy branching as training proceeds, confirming that the gradient-estimation advantage transfers to a stateful, large-action setting.

Abstract

Explore similar work

EPIG-Tree: Compute-Optimal Branching for Gradient-Efficient Reinforcement Learning

Beyond Outcome Rewards: Step-Level Self-Distilled Policy Optimization for Deep Search Agents

Maximizing Rollout Informativeness under a Fixed Budget: A Submodular View of Tree Search for Tool-Use Agentic Reinforcement Learning