Partial Policy Gradients for RL in LLMs

Puneet Mathur; Branislav Kveton; Subhojyoti Mukherjee; Viet Dac Lai

cs.LGMar 6, 2026

Partial Policy Gradients for RL in LLMs

Authors: Puneet Mathur, Branislav Kveton, Subhojyoti Mukherjee, Viet Dac Lai

Organizations: Adobe Research.

Abstract

Reinforcement learning is a framework for learning to act sequentially in an unknown environment. We propose a natural approach for modeling policy structure in policy gradients. The key idea is to optimize for a subset of future rewards: smaller subsets represent simpler policies, which can be learned more reliably because their empirical gradient estimates are more accurate. Our approach allows for modeling and comparison of different policy classes, including full planning, greedy, K-step lookahead, and segment policies. We evaluate the policies empirically on multiple persona-alignment conversational problems. Different policies excel in different problems, reflecting their different characteristics and highlighting the importance of our studied policy class.

Figures & tables

Figure 1: Examples of reward indices

\mathcal{R}_{t}

(yellow circles in columns) and action indices

\mathcal{S}_{t}

(blue circles in rows) in full, greedy, and

2

-step lookahead policies.

Policy	PC $\uparrow$	PC $\uparrow$	Policy	PC $\uparrow$	PC $\uparrow$
	Qwen			Llama
	Rollout	Step-level		Rollout	Step-level
$\color[rgb]{0,1,0}\tt Base$	0.799	0.886	$\color[rgb]{0,1,0}\tt Base$	0.425	0.425
PPO	0.875	0.895	PPO	0.613	0.624
$\color[rgb]{0,1,0}\tt GreedyPG$	0.623	0.930	$\color[rgb]{0,1,0}\tt GreedyPG$	0.739	0.833
$\color[rgb]{0,1,0}\tt 2\mathchar 45\relax Step\mathchar 45\relax PG$	0.930	0.936	$\color[rgb]{0,1,0}\tt 2\mathchar 45\relax Step\mathchar 45\relax PG$	0.721	0.759

Table 1: Evaluation on education domain in Consistent-LLMs dataset. Bold / underline is the best / second-best result.

Policy	PC $\uparrow$	PC $\uparrow$	Policy	PC $\uparrow$	PC $\uparrow$
	Qwen			Llama
	Rollout	Step-level		Rollout	Step-level
$\color[rgb]{0,1,0}\tt Base$	0.561	0.869	$\color[rgb]{0,1,0}\tt Base$	0.465	0.465
PPO	0.753	0.911	PPO	0.625	0.679
$\color[rgb]{0,1,0}\tt GreedyPG$	0.840	0.882	$\color[rgb]{0,1,0}\tt GreedyPG$	0.685	0.710
$\color[rgb]{0,1,0}\tt 2\mathchar 45\relax Step\mathchar 45\relax PG$	0.854	0.936	$\color[rgb]{0,1,0}\tt 2\mathchar 45\relax Step\mathchar 45\relax PG$	0.712	0.734

Table 2: Evaluation on therapy domain in Consistent-LLMs dataset. Bold / underline is the best / second-best result.

Policy	PC $\uparrow$	PC $\uparrow$	Policy	PC $\uparrow$	PC $\uparrow$
	Qwen			Llama
	Rollout	Step-level		Rollout	Step-level
$\color[rgb]{0,1,0}\tt Base$	0.760	0.775	$\color[rgb]{0,1,0}\tt Base$	0.478	0.478
PPO	0.818	0.845	PPO	0.556	0.558
$\color[rgb]{0,1,0}\tt GreedyPG$	0.762	0.881	$\color[rgb]{0,1,0}\tt GreedyPG$	0.799	0.850
$\color[rgb]{0,1,0}\tt 2\mathchar 45\relax Step\mathchar 45\relax PG$	0.782	0.813	$\color[rgb]{0,1,0}\tt 2\mathchar 45\relax Step\mathchar 45\relax PG$	0.929	0.939

Table 3: Evaluation on chatting domain in Consistent-LLMs dataset. Bold / underline is the best / second-best result.

Panel of Figure 2 — Figure 2: Residual of persona consistency scaled to $[-1,1]$ across trajectory steps for all domains with Qwen . $K=1$ represents $\color[rgb]{0,1,0}\tt GreedyPG$ and $K\in\{2,3,4,5\}$ represents $\color[rgb]{0,1,0}\tt K\mathchar 45\relax Step\mathchar 45\relax PG$ .

Panel of Figure 3 — Figure 3: Persona consistency as a function of sample size (number of training trajectories) for all domains with Llama . $K=1$ represents $\color[rgb]{0,1,0}\tt GreedyPG$ and $K\in\{2,3,4,5\}$ represents $\color[rgb]{0,1,0}\tt K\mathchar 45\relax Step\mathchar 45\relax PG$ .

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 4 — Figure 4: Scaling Laws for $\color[rgb]{0,1,0}\tt PPG$ : Across different LLMs (Llama-3.1-8B-Instruct, Qwen3-8B, and Gemma-7B-it), optimal value of lookahead $K$ in $\color[rgb]{0,1,0}\tt K\mathchar 45\relax Step\mathchar 45\relax PG$ scales with available training trajectories.

Panel of Figure 5 — Figure 5: Persona consistency of policy gradient methods vs trajectory length in education, therapy, and chatting domains for Llama3.1-8B-Instruct model. For each method, we report mean PC of all trajectories with $t$ steps, where $t=\{10,20,40,60\}$ .

Panel of Figure 6 — Figure 6: Persona consistency of policy gradient methods vs trajectory length in education, therapy, and chatting domains for Qwen3-8B model. For each method, we report mean PC of all trajectories with $t$ steps, where $t=\{10,20,40,60\}$ .

Panel of Figure 7 — Figure 7: Residual of Persona Consistency metric scaled to $[-1,1]$ across trajectory steps for Education, Therapy, Chatting domains using Llama3.1-8B-Instruct. K=1 represents $\color[rgb]{0,1,0}\tt GreedyPG$ ; K={2,3,4,5} represent $\color[rgb]{0,1,0}\tt K\mathchar 45\relax Step\mathchar 45\relax PG$ .

RL for Language Models

Reinforcement learning for language models uses feedback-driven policy updates to improve reasoning, alignment, and interactive behavior.

See all 433 papers

A First-Principles Derivation of LLM Policy Optimization: From Expected Reward to GRPO and Its Structural Extensions
Jun 15, 2026Jianghan Shen, Siqi Luo, Yue Li +9RL for Language Models Group Relative Policy Optimization
1Nanjing University · 2Shanghai Artificial Intelligence Laboratory · 3Shanghai Jiao Tong University +4
When are LLMs Sufficient Policy Optimizers for Sequential RL Tasks?
May 29, 2026Stephane Hatgis-Kessell, Emma BrunskillLanguage Model-Based Control Policy Optimization
Department of Computer Science, Stanford University
Modularized Reinforcement Learning on LLMs: From MDP Creation to Exploration and Learning
Jun 20, 2026Zhao Yang, Yuxuan Jiang, Ting-Chih Chen +18RL for Language Models Credit Assignment in RL
VU Amsterdam, The Netherlands · University of Maryland, Baltimore County, USA · Leiden University, The Netherlands +4

Jun 15, 2026J/K move · Enter open · S save

A First-Principles Derivation of LLM Policy Optimization: From Expected Reward to GRPO and Its Structural Extensions

Jianghan Shen, Siqi Luo, Yue Li, Jiyao Liu+8

1Nanjing University · 2Shanghai Artificial Intelligence Laboratory · 3Shanghai Jiao Tong University+4

RL for Language Models Group Relative Policy Optimization Policy Gradient Methods Proximal Policy Optimization Policy Learning

Open paper AskPDF arXiv

Policy gradient algorithms for language models optimize the same objective

J(θ) = \mathbb{E}*{τ\sim p*θ(τ)}[R(τ)]

, which has exactly two factors: the trajectory probability

p_θ(τ)

and the reward

R(τ)

. Every method from REINFORCE to PPO to GRPO and their descendants modifies one or both factors to address a specific failure in the preceding formulation. Existing surveys organize these methods by domain or chronology, which obscures the rationale behind each design choice and the precise location of its intervention within the gradient estimator. This survey revisits the landscape of LLM policy optimization from

J(θ)

on first principles and uses the trajectory side, induced by

p_θ(τ)

, and the reward side, induced by

R(τ)

, as the two axes along which methods are located. It covers the path from REINFORCE and PPO to GRPO, as well as post-GRPO variants, Agentic RL, and GRPO-OPD. The resulting framework is unified, diagnostic, and extensible: it analyzes methods from a shared objective, identifies which side each method modifies and why, and applies the same trajectory and reward axes across these settings. Across these settings, the framework also exposes compound failures that no single-side fix resolves and that therefore require joint design of the trajectory side and the reward side. The boundary cases and coupled failures identified by this map mark where existing solutions run out and provide a principled starting point for designing the next generation of LLM policy optimization algorithms.

Abstract

RL for Language Models

Explore similar work

A First-Principles Derivation of LLM Policy Optimization: From Expected Reward to GRPO and Its Structural Extensions

When are LLMs Sufficient Policy Optimizers for Sequential RL Tasks?

Modularized Reinforcement Learning on LLMs: From MDP Creation to Exploration and Learning