Paper deep dive

Partial Policy Gradients for RL in LLMs

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

Year: 2026Venue: arXiv preprintArea: cs.LGType: PreprintEmbeddings: 129

Intelligence

Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 93%

Last extracted: 3/13/2026, 12:16:43 AM

Summary

The paper introduces 'Partial Policy Gradients' (PPG), a reinforcement learning framework for LLMs that optimizes for subsets of future rewards to improve statistical efficiency. By decomposing rewards and focusing on K-step lookahead policies, the authors demonstrate improved persona consistency in conversational agents across domains like education and therapy, mitigating persona drift compared to standard full-planning policy gradients.

Entities (6)

Partial Policy Gradients · methodology · 98%K-step lookahead · policy-class · 95%Consistent-LLMs · dataset · 92%Gemma · llm · 90%Llama · llm · 90%Qwen · llm · 90%

Relation Signals (3)

K-step lookahead → instanceof → Partial Policy Gradients

confidence 98% · Full planning, greedy, and K-step lookahead policies are all instances of our framework.

Partial Policy Gradients → evaluatedon → Consistent-LLMs

confidence 95% · We empirically evaluate our policies on four persona-alignment conversational problems.

Partial Policy Gradients → optimizes → subset of future rewards

confidence 95% · The key idea is to optimize for a subset of future rewards: smaller subsets represent simpler policies.

Cypher Suggestions (0)

No Cypher suggestions yet.

Abstract

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.

PDF

Open source PDF →Open local PDF →

Trouble viewing inline? Open PDF directly →

Full Text

128,532 characters extracted from source content.

Expand or collapse full text

Partial Policy Gradients for RL in LLMs Puneet Mathur 1 Branislav Kveton 1 Subhojyoti Mukherjee 1 Viet Dac Lai 1 Abstract Reinforcement learning is a framework for learn- ing to act sequentially in an unknown environ- ment. We propose a natural approach for model- ing 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 em- pirical 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. Dif- ferent policies excel in different problems, reflect- ing their different characteristics and highlighting the importance of our studied policy class. 1. Introduction Reinforcement learning (RL) (Sutton and Barto, 1998) is a framework for learning to act sequentially in an unknown environment with the goal of maximizing future rewards. Due to its generality and broad applicability, RL has been studied extensively and many RL algorithms have been proposed, including temporal-difference learning (Sutton, 1988), Q-learning (Watkins and Dayan, 1992), policy gra- dients (Williams, 1992), and actor-critic methods (Sutton et al., 2000). RL is more data demanding than supervised learning because it plans for the future (Sutton and Barto, 1998). Therefore, many early works on planning and RL tried to increase the statistical efficiency of learning algo- rithms by assuming additional structure, such as a known model of the environment (Bellman, 1957; Boutilier et al., 1995; Koller and Parr, 1999), linear value or Q functions (Tsitsiklis and Van Roy, 1997; de Farias and Van Roy, 2003; Hauskrecht and Kveton, 2004), and baseline subtraction in policy gradients (Sutton et al., 2000; Baxter and Bartlett, 2001; Munos, 2006). Proximal policy optimization (PPO) (Schulman et al., 2017) 1 Adobe Research.Correspondence to: Puneet Mathur <puneetm@adobe.com>. Preprint. March 9, 2026. and group-relative policy optimization (GRPO) (Shao et al., 2024) are arguably the most popular RL algorithms for large language models (LLMs) (Ouyang et al., 2022). PPO is a policy gradient algorithm where the advantage is estimated using a per-token reward model. When the policy is learned from off-policy data, the advantage is corrected by clipped propensity scores (Horvitz and Thompson, 1952; Ionides, 2008). The challenge with applying PPO in practice is that a per-token reward model has to be learned. GRPO is a policy gradient algorithm where the advantage is estimated by simulation, which eliminates the need for a per-token reward model. The challenge with applying GRPO is that the advantage is attributed equally to all tokens, potentially reducing statistical efficiency. We propose a natural approach for introducing policy struc- ture in policy gradients. The first step is decomposing the to- tal reward additively over time. This decomposition may be a part of problem definition, for instance when the problem is a Markov decision process (MDP) with per-step rewards (Bellman, 1957; Puterman, 1994), or obtained by a simple algebraic transformation. The second step is optimization for subsets of future rewards. The subsets lead to simpler policies than full planning, which can be learned more re- liably because their empirical gradient estimates are more accurate. We propose both online and offline algorithms for learning our policies, and analyze them. Full planning, greedy, andK-step lookahead policies are all instances of our framework. LLMs are increasingly used to simulate human personas for social role-play (Tseng et al., 2024; Jandaghi et al., 2024). Despite that, off-the-shelf LLMs frequently drift from their assigned personas, contradict to their earlier statements, or abandon role-appropriate behavior (Ji et al., 2025; Abdulhai et al., 2025). We apply our policies to LLMs to maintain consistent and faithful human personas in role-playing di- alogues. This problem is challenging because many dia- logues have long horizons, on average20steps and up to 60. Our results show thatK-step lookahead policies lead to consistent personas in extended dialogues, mitigating the persona drift that plagues base models. We make the following contributions: 1.We propose a framework for modeling policy structure 1 arXiv:2603.06138v1 [cs.LG] 6 Mar 2026 Partial Policy Gradients for RL in LLMs in policy gradients. The key idea is to optimize for a subset of future rewards, which can be either observed or obtained by a simple algebraic transformation. Our formulation is more general than those in the closest related works, which focused only on specific policies (Kveton et al., 2025; Guo et al., 2025), sometimes even without observations. 2.Our framework trades off the complexity of learned policies for the statistical efficiency of learning them. Specifically, optimization for smaller subsets of future rewards leads to learning simpler policies, which can be learned more reliably because their empirical gra- dient estimates are more accurate. We propose both online and offline algorithms for learning the policies. 3. We instantiate the framework: online policy gradient (Williams, 1992), policy gradient for adaptive submod- ular policies (Kveton et al., 2025) and segment-level credit assignment Guo et al. (2025), can be all viewed as instances of our framework. This is the first work that proposes and empirically evaluates K-step looka- head policies in LLMs. 4. We empirically evaluate our policies on four persona- alignment conversational problems across multiple do- mains: education, therapy, chatting, and generic. We experiment with three LLMs: Qwen (Yang et al., 2025), Llama (Grattafiori et al., 2024), and Gemma (Team et al., 2025). We show thatK-step lookahead policies lead to consistent persona across extended dialogues (Section 5.4) by preventing persona drift (Section 5.5). The optimal value ofKdepends on the domain cam- plexity and the amount of training data (Section 5.6). The paper is organized as follows. We present our setting in Section 2. In Section 3, we present our framework for opti- mizing a subset of future rewards, interpret it, and analyze it. We instantiate the framework in Section 4 and evaluate it empirically in Section 5. Related works are reviewed in Section 6, and we conclude in Section 7. 2. Setting We start with introducing our notation. The marginal and conditional probabilities under the probability measurep arep(X = x)andp(X = x| Y = y), respectively; and we writep(x)andp(x| y)when the random variables are clear from context. The indicator function is1·. For a positive integern, we define[n] = 1,...,n. We denote thei-th entry of vectorvbyv i . If the vector is already indexed, such as v j , we write v j,i . We consider a general reinforcement learning problem (Sut- ton and Barto, 1998) where an agent interacts with an en- vironment in an episode consisting of multiple steps. The agent takes actions conditioned on the interaction history and the environment responds. When the episode ends, it is assigned a reward. The reward measures the quality of the interaction and the agent maximizes it. We formalize the problem as follows. The agent first ob- serves contextx ∈ S, whereSis the space of all token sequences. The context defines the task. The interaction be- tween the agent and environment consists of steps indexed byt∈ N, whereNis a set of positive integers. At stept, the agent takes an actiona t ∈Sand the environment responds with an observationy t ∈ S. The interaction history up to stepnis a trajectoryτ n = (a 1 ,y 1 ,...,a n ,y n )ofnactions and observations. We assume that the number of stepsnis fixed. The reward is a non-negative function ofxandτ n , denoted byr(x,τ n ) ≥ 0, and measures the quality of the interaction. We do not make any additional assumptions on the reward. This is to maintain generality and because our algorithms (Section 3) do not require it. The agent follows a policy conditioned on its history. Specif- ically, the probability that actionais taken in contextxand historyτ t−1 isπ(a | x,τ t−1 ;θ), and is parameterized by θ ∈ Θ. We callθa policy andΘthe space of policy pa- rameters. The probability of observingy t conditioned on historyτ t−1 and actiona t is denoted byp(y t | x,τ t−1 ,a t ). We slightly abuse our notation and denote the probability of trajectory τ n in context x under policy θ by π(τ n | x;θ) = n Y t=1 p(y t | x,τ t−1 ,a t )π(a t | x,τ t−1 ;θ). The factorization follows from the chain rule of probability. The expected value of policyθ, whereqis a distribution over contexts x, is defined as V (θ) = E x∼q,τ n ∼π(·|x;θ) [r(x,τ n )] .(1) Our goal is to learn a policyθ ∗ = arg max θ∈Θ V (θ)with the highest expected value. Our formulation is more gen- eral than those in the closest related works, which focused only on specific policies and did not consider observations (Kveton et al., 2025; Guo et al., 2025). Our framework can model multiple use cases. For instance, suppose that we want to maximize the pedagogical value of a conversation overnsteps (Scarlatos et al., 2025). Then τ n is ann-step conversation andr(x,τ n )is its pedagogical value. We experiment with such problems in Section 5. In math reasoning problems (Shao et al., 2024),τ n would be agent’s reasoning andr(x,τ n )would indicate whether its response is correct. 3. Algorithms This section has four main parts. In Section 3.1, we intro- duce policy gradients. In Section 3.2, we present partial 2 Partial Policy Gradients for RL in LLMs Algorithm 1 Meta policy gradient algorithm. 1:Input: Learning rate schedule(α j ) j∈N , batch sizeB, gradient estimator g 2: Initialize θ 3: j ← 1 4: while (not convergence) do 5:for i = 1,...,B do 6:x i ∼ q 7:τ i,n ∼ π(·| x i ;θ) 8:r i ← r(x i ,τ i,n ) 9:BatchB ←(x i ,τ i,n ,r i ) B i=1 10: θ ← θ + α j g(B) 11: j ← j + 1 12: Output: Learned policy θ policy gradients. We interpret them in Section 3.3 and intro- duce offline partial policy gradients in Section 3.4. 3.1. Policy Gradient We want to maximize the expected policy valueV (θ)in(1). While this can be done many ways (Sutton and Barto, 1998), the most straightforward method for complex policies, like those represented by LLMs, are policy gradients (Williams, 1992). The key idea in an online policy gradient is to update the policyθiteratively by gradient ascent. The gradient of V (θ) at θ is ∇V (θ) = E [r(x,τ n )∇ logπ(τ n | x;θ)](2) = E " r(x,τ n ) n X t=1 ∇ logπ(a t | x,τ t−1 ;θ) # , whereE [·] = E x∼q,τ n ∼π(·|x;θ) [·]is an expectation over the same random variables as in(1). This identity can be derived by a direct application of the score identity (Aleksandrov et al., 1968). All policy gradient algorithms in this work are variants of the meta algorithm in Algorithm 1. The gradient in Line 10 is estimated from a batchBof trajectories (Line 9) of sizeB, which are sampled in Line 7 and assigned rewards in Line 8. We denote thei-th trajectory inBbyτ i,n ; and its context and reward byx i andr i , respectively. The gradient estimator for the classic online policy gradient is g(B) = 1 B B X i=1 r i n X t=1 ∇ logπ(a i,t | x i ,τ i,t−1 ;θ).(3) We refer to this instance of Algorithm 1 as PG. 3.2. Partial Policy Gradient While the gradient in(2)is valid, the reward is attributed equally to each step, since all∇ logπ(a t | x,τ t−1 ;θ)are multiplied by the same rewardr(x,τ n ). If the attribution could be improved, we could reduce the variance in esti- mator(3)and improve statistical efficiency. We propose a general attribution framework that allows for such tradeoffs. The framework has two components. Factored reward. We assume that the reward factors as r(x,τ n ) = n X t=1 r t , r 1 ,...,r n ≥ 0,(4) wherer t = r(x,τ t )denotes the reward at stept, and is a function ofxandτ t . The rewardr t with such properties can be either observed, for instance if the environment is an MDP with per-step rewards, or designed. We present a general design that guarantees (4) next. Lemma 1. Letf (x,τ t )be any function ofxandτ t such that 0 = f (x,τ 0 )≤·≤ f (x,τ n ) = r(x,τ n ). Let r t = f (x,τ t )− f (x,τ t−1 ) hold for allt ∈ [n]. Then P n t=1 r t = r(x,τ n )andr t ≥ 0 for all t∈ [n]. The proof of the claim is in Section A.1. Sincef (x,τ t ) can be viewed as a trajectory reward up to stept,r t is the per-step gain in the total trajectory reward at stept. Iffis not monotone, a monotone ̄ f can be obtained by a transfor- mation ̄ f (x,τ t ) = max ℓ∈[t] f (x,τ ℓ ). Because of this, and without loss of generality, we assume thatr t ≥ 0in the rest of our paper, unless stated otherwise. One natural choice of fis an LLM judge (Zheng et al., 2023) that can judge partial trajectories. Similar decompositions were proposed before in more specific problems, without observations (Guo et al., 2025) and in optimization of submodular policies (Kveton et al., 2025). Reward attribution. We consider a general partial policy gradient of the form E " n X t=1 X ℓ∈R t r ℓ ! ∇ logπ(a t | x,τ t−1 ;θ) # ,(5) whereR t ⊆ [n]\[t−1]are reward indices affected by action at steptthrough∇ logπ(a t | x,τ t−1 ;θ). Specifically,R t is the subset of future steps whose rewards are affected by a t conditioned on the history. The gradient can be estimated empirically as g(B) = 1 B B X i=1 n X t=1 X ℓ∈R t r i,ℓ ! ∇ logπ(a i,t | x i ,τ i,t−1 ;θ), (6) whereτ i,n is thei-th trajectory in batchB; anda i,ℓ ,r i,ℓ , andτ i,ℓ−1 are the action, reward, and trajectory history at steptof this trajectory, respectively. When(6)is used in 3 Partial Policy Gradients for RL in LLMs Algorithm 1, we call itpartialpolicygradient (PPG) since only a subset of future rewards is used in policy optimiza- tion. Note thatPPGreduces toPGwhenR t = [n]\ [t− 1] for allt∈ [n], as we show in Section 4.1. Our framework al- lows us to reason about otherR t and understand associated statistical tradeoffs. PPGis more statistically efficient thanPG, and we prove this formally in Section A.2. Specifically, Theorem 5 says that a partial policy gradient estimatorg(B)that optimizes for a subset of rewards of another estimatorg ′ (B)concentrates at E [g(B)]faster thatg ′ (B)concentrates atE [g ′ (B)]. The key idea in the proof is to treat(6)as a sum of random vectors, independent acrossibut dependent acrosst. Then we use Hoeffding’s inequality, which is a standard tool in concentra- tion analyses (Raginsky and Sason, 2012; Boucheron et al., 2013), to argue that the sum of a subset of random vectors concentrates faster than the sum of the original set, when the vectors are sampled from the same distribution. 3.3. Interpretation We interpret gradient ascent with respect to(5)next. That requires additional notation. LetS t =ℓ∈ [t] : t∈R ℓ ⊆ [t]be action indices that affect the reward at stept,r t . Note that by definition, n X t=1 X ℓ∈R t r ℓ ! ∇ logπ(a t | x,τ t−1 ;θ)(7) = n X t=1 r t X ℓ∈S t ∇ logπ(a ℓ | x,τ ℓ−1 ;θ). This is because we can first sum over log-probabilities and then over the associated rewards, or the other way around. Figure 1 illustratesR t andS t in full, greedy, andK-step lookahead policies in Section 4. We also define ̄ S t = [t] t , the complement ofS t with respect to the firsttsteps. For anyS,Z ⊆ [n], we define E S [·] = E x∼q,τ S ∼π(·|x;θ) [·] , E S [·|Z] = E τ S ∼π(·|x;θ) [·|x,τ Z ] . HereE S [·]is the expectation over actions and observations at stepsS, and contextx; andE S [·|Z]is the expectation over actions and observations at stepsS, conditioned on actions and observations at stepsZ, and contextx. With the new notation in hand, we get the following claim. Lemma 2. Letx ∼ qandτ n ∼ π(· | x;θ). Then for any R t ⊆ [n]\ [t− 1], the gradient in (5) is equal to E " n X t=1 ∇E S t r t ̄ S t # . The proof is in Section A.3. The lemma says that the gra- dient in(5)increases the sum of expected per-step rewards r t through actions at stepsS t that affect the corresponding expected rewards. The histories ̄ S t are sampled fromπ. We discuss different forms of these policies in Section 4. 3.4. Offline Partial Policy Gradient So far, we focused on online policy gradients (Section 3.2), where the trajectories are sampled from learned policyπ. Another common and practical setting is offline RL (Jaques et al., 2020; Levine et al., 2020), where the trajectories are sampled from a logged dataset. We represent it by another distributionρand propose an offline partial policy gradient of the form E ρ " n X t=1 X ℓ∈R t r ℓ ! ∇ logπ(a t | x,τ t−1 ;θ) # ,(8) whereE ρ [·] = E x∼q,τ n ∼ρ(·|x;θ) [·]. Note that the only dif- ference from(5)is that the trajectories are sampled from another policyρ. Therefore, the empirical estimator of(8) from batchB has the same algebraic form as (6). Based on our discussion, an offline variant ofPPGcan be obtained by changing the sampling distribution in Line 7 of Algorithm 1 toτ i,n ∼ ρ(·| x;θ). Our gradient estimator is motivated by Mukherjee et al. (2025) and can be related to maximizing the original online objective as follows. Lemma 3. Let all quantities be defined as in Lemma 2. Let E ρ be an expectation with respect to another policyρ. Then for anyR t ⊆ [n]\ [t− 1], the gradient in (8) is equal to E ρ " n X t=1 ∇ E ρ S t " r t X ℓ∈S t logπ(a ℓ ,y ℓ | x,τ ℓ−1 ;θ) ̄ S t # | z f t (θ) # . Whenr t ≥ 0for allt∈ [n],f t (θ)≤ E S t r t ̄ S t holds. In general,|f t (θ)− E S t r t ̄ S t | ≤ C(π,ρ) , whereC(π,ρ) is a constant depending on the propensity ratios ofπandρ, and their logarithms. The proof is in Section A.4. The first claim says that the gradient increases the sum of expected per-step rewardsr t weighted by log-probabilities under logging policyρ, de- noted byf t (θ), through actions at stepsS t that affect them. The second claim says that the true objectiveE S t r t ̄ S t that we want to optimize is bounded from below byf t (θ), which we optimize. The lower bound holds under the as- sumption thatr t ≥ 0. Without it,f t (θ)is stillC(π,ρ)close to E S t r t ̄ S t , as stated in the last claim. Since Lemma 3 provides guarantees for negative rewards, offlinePPGcan be implemented with advantages and we do that in Section 5. The advantage for steptin trajectoryiis 4 Partial Policy Gradients for RL in LLMs a 2 a 3 a 1 y 1 y 2 y 3 r 1 r 2 r 3 R 1 Trajectory R 2 R 3 1 2 3 1 2 3 S 1 S 2 S 3 Full 2 3 3 1 2 1 R 1 R 2 R 3 1 Greedy 2 3 R 1 R 2 R 3 1 2 2 3 2-step 2 3 3 1 2 1 3 2 1 Figure 1. Examples of reward indicesR t (yellow circles in columns) and action indicesS t (blue circles in rows) in full, greedy, and 2-step lookahead policies. computed as ̃r i,t = P ℓ∈R t r i,ℓ − μ G σ G ,(9) whereμ G andσ G are the mean and standard deviation of P ℓ∈R t r j,ℓ computed jointly overj ∈Gandt∈ [n], where Gis the group of trajectories that trajectoryibelongs to. In our experiments in Section 5, we group the trajectories by the persona and topic. 4. Partial Policy Gradient Instances To show the generality of our framework, we derive and compare several instances ofPPG. Our concentration argu- ments rely on Theorem 5, which is stated and proved in Section A.2. Segment policies are derived in Section A.5. Offline policy gradients can be derived analogously. 4.1. Full Policy Gradient Suppose that all actionsa t in(5)affect all future rewards. ThenR t = [n]\ [t− 1]for all stepst∈ [n]. By Lemma 2, whereS t = [t] and ̄ S t =∅, the gradient is E " n X t=1 ∇E [t] [r t ] # = E " ∇ n X t=1 E [t] [r t ] # (10) = E " ∇E " n X t=1 r t ## =∇E [r] . In plain English, all actions up to steptaffect the reward at stept; and the gradient is a sum of gradients of the corre- sponding expected rewards. Note that it is identical to(2), because future actions do not affect past rewards. 4.2. Greedy Policy Gradient Suppose that all actionsa t in(5)affect only the immediate reward. ThenR t =tfor allt∈ [n]. By Lemma 2, where S t =t and ̄ S t = [t− 1], the gradient is E " n X t=1 ∇E t [r t | [t− 1]] # .(11) In plain English, only the action at steptaffects the reward at stept; and the gradient is a sum of gradients of the corre- sponding expected rewards. We refer toPPGwithR t =t as GreedyPG. This policy is simpler than that in Section 4.1, where the reward is affected by each previous action, and therefore its gradient can be estimated more efficiently. Specifically, by Theorem 5, the empirical gradient inGreedyPGconcen- trates faster than that inPGsinceS t = t ⊆ [t]. While this policy class was proposed, analyzed, and evaluated in LLMs by Kveton et al. (2025), it was done only for adaptive submodular maximization (Golovin and Krause, 2011). The connection to statistical efficiency was not established. 4.3. K-Step Lookahead Policy Gradient Suppose that all actionsa t in(5)affect the nextKrewards. ThenR t = [mint + K− 1,n]\ [t− 1]for allt ∈ [n]. By Lemma 2, whereS t = [t]\ [maxt− K, 0]and ̄ S t = [maxt− K, 0], the gradient is E " n X t=1 ∇E [t]\[maxt−K,0] [r t | [maxt− K, 0]] # . (12) In plain English, only the lastKactions up to steptinclud- ing affect the reward at stept; and the gradient is a sum of gradients of the corresponding expected rewards. We call PPGwith thisR t asK-Step-PG. Our paper is the first work that proposes this policy in LLMs. For a given lookahead, such asK = 3, we write3-Step-PG. Note that1-Step-PG is identical to GreedyPG in Section 4.2. This policy is simpler than that in Section 4.1, where the reward is affected by each previous action, but also more complex than that in Section 4.2, where the reward is af- fected only by the immediate action. The policy naturally 5 Partial Policy Gradients for RL in LLMs interpolates between these two extremes. Specifically, since t⊆S t = [t]\ [maxt− K, 0]⊆ [t] holds for anyK, we have by Theorem 5 that the empirical gradient inK-Step-PGconcentrates faster than that inPG but slower than that inGreedyPG. The complexity of the policy increases and its concentration rate decreases asK increases. 5. Experiments We describe our benchmarks in Section 5.1, evaluation pro- tocol in Section 5.2, and baselines in Section 5.3. The poli- cies are evaluated in Section 5.4. We analyze their stability in Section 5.5 and statistical efficiency in Section 5.6. In addition, we study the policies across conversation lengths in Section C, present scaling laws for Qwen, Llama, and Gemma models in generic domain in Section B, and give qualitative examples in Section G. 5.1. Data Settings Problem Formulation. Our learning problem is maintain- ing consistent and faithful human personas over time in role-playing dialogues like therapy, education, and chatting. We instantiate the framework in Section 2 as follows. The contextxis a persona definition given at the start of the interaction in a system prompt. A trajectoryτ n is a di- alogue history between an LLM-simulated agent and the environment. The agent’s responses are actionsa t and the environment’s replies are observations y t . Dataset. We experiment with Consistent-LLMs benchmark (Abdulhai et al., 2025), which provides human-persona dia- logues across 3 distinct domains: (i) Education: Tutoring sessions where a student agent maintains specific learning preferences and the tutor acts as the environment; (i) Ther- apy: Counseling sessions where a patient agent maintains consistent mental health backgrounds while the therapist serves as the environment; and (i) Chatting: Casual di- alogue where a simulated human is the agent maintaining biographical consistency (family, career, and so on) against a conversational environment. Each domain comprises of 6 500trajectories (5 200training and1 300test) of vary- ing lengths (2to60steps, with23steps on average). The trajectories are generated using a Llama-3.1-8B model. 5.2. Evaluation Abdulhai et al. (2025) proposed two metrics to evaluate role-playing consistency: (i) Prompt consistency measures if an agent’s response aligns with their assigned persona or background. (i) Index consistency evaluates if an agent’s response contradicts to their own previous responses in the conversation history based on their past prompt consistency scores. Since index consistency is a secondary metric de- rived from prompt consistency, we use prompt consistency in our work; and call it persona consistency (PC) because it measures the consistency of a persona. Using this terminol- ogy, the reward at step t, r t = r(x,τ t ), is its PC. All policies are learned offline from the training trajectories of Abdulhai et al. (2025). At test time, we measure the PC using an LLM judge with the same prompts as in Abdulhai et al. (2025). Specifically, for each agent’s responsea t , we ask the LLM judge to evaluate if it contradicts the persona description provided as contextxand score it1(consistent) if the judge outputs “NO” (no contradiction) or0(inconsis- tent) if the judge outputs “YES” (contradiction detected). The judge is implemented using GPT-4o mini with tempera- ture0, and we report the mean scores across all steps of the evaluated trajectories. See prompts in Section E. We employ two complementary evaluation protocols to as- sess persona consistency under different distribution condi- tions. The first follows the rollout methodology from Ab- dulhai et al. (2025), while the second is designed to isolate consistency improvements from error propagation effects. Rollout evaluation. Following Abdulhai et al. (2025), we generate complete conversation trajectories starting with only the initial persona description as context, alternating between responses of the agent and a fixed simulated envi- ronment implemented using Llama-3.1-70B model. This exposes the agent to distribution shifts, as inconsistencies compound over time, making it a challenging test of main- taining long-horizon consistency. Step-level evaluation. To disentangle consistency improve- ments from error compounding, we introduce a controlled evaluation setting where each step is assessed independently. Specifically, we take the test trajectories of Abdulhai et al. (2025) and for each step in them, we generate the response of the agent given the conversation history of that particular trajectory. This design isolates the model’s inherent consis- tency capability from the cascading effects of its own errors without model-induced distribution shift. 5.3. Baselines and Compared Methods We have two baselines.Baseis an unmodified base model used zero-shot to show what can be attained without learn- ing. PPO (Schulman et al., 2017) was used as a policy optimization algorithm in Abdulhai et al. (2025), and hence is a state-of-the-art baseline in our domain. We experiment with two agent models: Llama-3.1-8B-Instruct and Qwen3- 8B, which we abbreviate asQwenandLlama, respectively. All our policies are learned offline by an offline policy gra- dient (Section 3.4), whereρsamples trajectories uniformly at random from the training set. The reward at steptis the persona consistency at stept(Section 5.2). The rewards are 6 Partial Policy Gradients for RL in LLMs Qwen Llama RolloutStep-levelRolloutStep-level PolicyPC↑PC↑PolicyPC↑PC↑ Base0.7990.886 Base0.4250.425 PPO0.8750.895PPO0.6130.624 GreedyPG0.6230.930 GreedyPG0.7390.833 2-Step-PG0.9300.936 2-Step-PG0.7210.759 3-Step-PG0.9480.960 3-Step-PG0.5730.645 4-Step-PG0.971 0.9544-Step-PG0.8370.858 5-Step-PG0.8810.937 5-Step-PG0.7750.896 PG0.9760.990 PG0.9130.974 Table 1. Evaluation on education domain in Consistent-LLMs dataset. Bold / underlineis the best / second-best result. Qwen Llama RolloutStep-levelRolloutStep-level PolicyPC↑PC↑PolicyPC↑PC↑ Base0.5610.869 Base0.4650.465 PPO0.7530.911PPO0.6250.679 GreedyPG0.8400.882 GreedyPG0.6850.710 2-Step-PG0.8540.936 2-Step-PG0.7120.734 3-Step-PG0.9370.940 3-Step-PG0.7800.784 4-Step-PG0.8350.830 4-Step-PG0.6730.687 5-Step-PG0.869 0.890 5-Step-PG0.6900.740 PG0.7850.815 PG0.5860.595 Table 2. Evaluation on therapy domain in Consistent-LLMs dataset. Bold / underlineis the best / second-best result. Qwen Llama RolloutStep-levelRolloutStep-level PolicyPC↑PC↑PolicyPC↑PC↑ Base0.7600.775 Base0.4780.478 PPO0.8180.845PPO0.5560.558 GreedyPG0.7620.881 GreedyPG0.7990.850 2-Step-PG0.7820.813 2-Step-PG0.9290.939 3-Step-PG0.9850.995 3-Step-PG0.821 0.835 4-Step-PG0.8310.925 4-Step-PG0.6920.715 5-Step-PG0.9580.9645-Step-PG0.8020.824 PG0.8670.813 PG0.7270.771 Table 3. Evaluation on chatting domain in Consistent-LLMs dataset. Bold / underlineis the best / second-best result. standardized as in(9). We learn a full planning policy (PGin Section 4.1), greedy policy (GreedyPGin Section 4.2), and K-step lookahead policies (K-Step-PGin Section 4.3) for K =2, 3, 4, 5. Additional training details in Section F. 5.4. Main Results We start by evaluating all policies in all domains: education, therapy, and chatting; and report their average persona con- sistency in Tables 1 to 3. While the best performing policy is domain-dependent, policy gradients universally outper- formBaseand PPO, demonstrating the robustness of offline RL by policy gradients. In chatting and therapy domains, K-step lookahead policies (2-Step-PGand3-Step-PG) per- form extremely well across bothQwenandLlamamodels. PGis the best policy in the education domain, because a faithful simulation of a tutored student requires long-term planning. We observe consistent results across rollout and step-level evaluation methodologies, suggesting that our gains are robust despite the cascading errors commonly ob- served in long dialogues. Our results show similar trends for bothQwenandLlamamodels, suggesting that our ap- proach generalizes across multiple LLM architectures. 5.5. On Effectiveness of Simpler Policies Now we try to understand the best performing policies. For all policies, we compute the mean persona consistency at each conversation step, subtract the mean over all steps, and renormalize to[−1, 1]jointly over all policies. Note that this is just a linear transformation of consistency scores that makes them comparable across the policies. We call this quantity a residual and report it in Figure 2. The residual of Basemonotonically decreases as conversation progresses, confirming that LLMs are susceptible to persona drift in long-horizon conversations (Nie et al., 2021). The residual ofGreedyPGoscillates a lot. These ripples reflect the pol- icy’s repeated failed recovery attempts from persona drift, due to not planning for the future. In comparison,K-step lookahead policies (K-Step-PG) maintain stable residuals across trajectory steps. We observe that increasingKfrom K = 1(GreedyPG) toK = 3(K-Step-PG) produces pro- gressively smoother residuals, supporting our theory that assigning credit to longer horizons can prevent persona drift by maintaining consistency continuously over time. Fi- nally, in therapy and chatting domains, we observe thatPG oscillates a lot and thus performs worse than3-Step-PG. This experiment shows that the temporal scope of credit assignment directly influences policy stability, with longer horizons providing smoother and more robust behavior at the cost of increased data requirements. Results forLlama are in Section D. See qualitative analysis in Section G. 5.6. On Statistical Efficiency of Simpler Policies We investigate the statistical efficiency of learning different policies next. Specifically, we vary the number of training trajectories, from50(low-data regime,1%of training data) to5 000(high-data regime, all training data), and plot the persona consistency of all policies in Figure 3. Our results reveal a fundamental trade-off between policy complexity and statistical efficiency. In the low-data regime,GreedyPG consistently excels in all domains, confirming our hypoth- esis that simpler policies are easier to learn when data are scarce. As data scales to moderate levels (500trajectories, 10%of training data),K-step lookahead policies dominate across all domains. In particular,2-Step-PGand3-Step-PG are the best policies in therapy and chatting domains. These policies strike balance between planning enough for the fu- ture and being able to learn it. In the education domain, we 7 Partial Policy Gradients for RL in LLMs 0 20 40 60 Base PG K=1 K=2 K=3 K=4 K=5 −1 0 1 Step Policy Scaled Residuals Education 0 20 40 60 Base PG K=1 K=2 K=3 K=4 K=5 −1 0 1 Step Policy Scaled Residuals Therapy 0 20 40 60 Base PG K=1 K=2 K=3 K=4 K=5 −1 0 1 Step Policy Scaled Residuals Chatting Figure 2. Residual of persona consistency scaled to[−1, 1]across trajectory steps for all domains withQwen.K = 1represents GreedyPG and K ∈2, 3, 4, 5 represents K-Step-PG. 5e15e25e3 0.2 0.4 0.6 0.8 Trajectories Persona Consistency Education 5e15e25e3 0.4 0.6 0.8 Trajectories Therapy 5e15e25e3 0.4 0.6 0.8 Trajectories Chatting Base K = 1K = 2K = 3K = 4K = 5 PG Figure 3. Persona consistency as a function of sample size (number of training trajectories) for all domains withLlama.K = 1 representsGreedyPGandK ∈2, 3, 4, 5representsK-Step-PG. observe that full planning withPGeventually achieves the highest persona consistency in the high-data regime. We val- idate this scaling behavior across three model architectures (Llama, Qwen, Gemma) on an additional dataset (Synthetic Persona Chat (Jandaghi et al., 2024)) in Section B, confirm- ing that optimal lookahead horizonKconsistently increases with available training data, establishing a practical design principle for calibrating policy complexity to data budgets. 6. Related Work The closest related works are Kveton et al. (2025) and Guo et al. (2025). Kveton et al. (2025) optimized greedy poli- cies using online policy gradients and Guo et al. (2025) optimized segment policies using GRPO. Our formalism captures both of these policy classes. Unlike Kveton et al. (2025), we consider more general optimization problems than adaptive submodularity (Golovin and Krause, 2011). Unlike Guo et al. (2025), we consider a general RL model that allows observations of the environment. Learning of policies in finite-horizon Markov decision pro- cesses is a classic topic and it is well known that they are non-stationary (Puterman, 1994). Monte Carlo estimation of Q values, by the sum of infinite-horizon discounted rewards, is discussed in detail in Sutton and Barto (1998) and was generalized to advantages by Schulman et al. (2016). We do not discount rewards and show that maximization of a subset of future rewards results in a new simpler policy that is easier to learn.K-step lookahead policies were studied in other contexts before. Efroni et al. (2018a) used them to improve policy iteration, Efroni et al. (2018b) applied them to online and approximate RL, and Efroni et al. (2020) ex- tended them to real-time dynamic programming. Adaptive lookahead polices were studied by Rosenberg et al. (2023) andK-step lookahead policies were used in offline RL by Sikchi et al. (2022). The main difference in our work is that we only use policy gradients without any additional model, such as that of a value or Q function. 7. Conclusions Most popular RL algorithms in LLMs are variants of policy gradients. In this work, we propose a natural approach for introducing policy structure in policy gradients. The key idea is to optimize for a subset of future rewards. Smaller subsets of rewards define simpler policies, which can be learned more reliably because their empirical gradients are more accurate. We interpret the policies, and propose both online and offline algorithms for learning them. Full plan- ning, greedy,K-step lookahead, and segment policies are all instances of our framework. We evaluate the policies on conversation optimization problems and observe that dif- ferent policies excel in different domains, highlighting the importance of our introduced structure. We note that the idea of optimizing for subsets of future rewards is general and can be straightforwardly applied to regularized policies (Todorov, 2006) and also GRPO (Shao et al., 2024). Limitations. First, our work provides both theory and exper- iments that justify policy optimization for subsets of future rewards. However, we never formally show that the better performance in our experiments is due to a lower variance in policy gradient estimates. Second, the per-step reward in (4)is an important design choice that may have impact on empirical performance. We do not ablate it and instead rely on existing works to provide it (Section 5). Finally, our eval- uation is limited to conversational problems, although we experiment with four different problems and three LLMs. 8 Partial Policy Gradients for RL in LLMs Impact Statement This paper presents work whose goal is to advance the field of machine learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here. References Marwa Abdulhai, Ryan Cheng, Donovan Clay, Tim Althoff, Sergey Levine, and Natasha Jaques. Consistently sim- ulating human personas with multi-turn reinforcement learning. arXiv preprint arXiv:2511.00222, 2025. V. M. Aleksandrov, V. I. Sysoyev, and V. V. Shemeneva. Stochastic optimization. Engineering Cybernetics, 5:11– 16, 1968. Jonathan Baxter and Peter Bartlett. Infinite-horizon policy- gradient estimation. Journal of Artificial Intelligence Research, 15:319–350, 2001. Richard Bellman. Dynamic Programming. Princeton Uni- versity Press, Princeton, NJ, 1957. Stephane Boucheron, Gabor Lugosi, and Pascal Massart. Concentration Inequalities: A Nonasymptotic Theory of Independence. Oxford University Press, 2013. Craig Boutilier, Richard Dearden, and Moises Goldszmidt. Exploiting structure in policy construction. In Proceed- ings of the 14th International Joint Conference on Artifi- cial Intelligence, pages 1104–1111, 1995. Daniela Pucci de Farias and Benjamin Van Roy. The linear programming approach to approximate dynamic program- ming. Operations Research, 51(6):850–856, 2003. Yonathan Efroni, Gal Dalal, Bruno Scherrer, and Shie Man- nor. Beyond the one-step greedy approach in reinforce- ment learning. In Proceedings of the 35th International Conference on Machine Learning, 2018a. Yonathan Efroni, Gal Dalal, Bruno Scherrer, and Shie Man- nor. Multiple-step greedy policies in approximate and online reinforcement learning. In Advances in Neural Information Processing Systems 31, 2018b. Yonathan Efroni, Mohammad Ghavamzadeh, and Shie Man- nor. Online planning with lookahead policies. In Ad- vances in Neural Information Processing Systems 33, 2020. Daniel Golovin and Andreas Krause. Adaptive submod- ularity: Theory and applications in active learning and stochastic optimization. Journal of Artificial Intelligence Research, 42:427–486, 2011. Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhi- nav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024. Yiran Guo, Lijie Xu, Jie Liu, Dan Ye, and Shuang Qiu. Segment policy optimization: Effective segment-level credit assignment in RL for large language models. In Advances in Neural Information Processing Systems 38, 2025. Milos Hauskrecht and Branislav Kveton. Linear program approximations for factored continuous-state Markov de- cision processes. In Advances in Neural Information Processing Systems 16, pages 895–902, 2004. D. G. Horvitz and D. J. Thompson. A generalization of sam- pling without replacement from a finite universe. Journal of the American Statistical Association, 47(260):663–685, 1952. Edward Ionides. Truncated importance sampling. Journal of Computational and Graphical Statistics, 17(2):295–311, 2008. Pegah Jandaghi, XiangHai Sheng, Xinyi Bai, Jay Pujara, and Hakim Sidahmed. Faithful persona-based conversa- tional dataset generation with large language models. In Proceedings of the 6th Workshop on NLP for Conversa- tional AI (NLP4ConvAI 2024), pages 114–139, 2024. Natasha Jaques, Judy Hanwen Shen, Asma Ghandeharioun, Craig Ferguson, Agata Lapedriza, Noah Jones, Shixiang Gu, and Rosalind Picard. Human-centric dialog training via offline reinforcement learning. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, 2020. Ke Ji, Yixin Lian, Linxu Li, Jingsheng Gao, Weiyuan Li, and Bin Dai. Enhancing persona consistency for LLMs’ role-playing using persona-aware contrastive learning. In Wanxiang Che, Joyce Nabende, Ekaterina Shutova, and Mohammad Taher Pilehvar, editors, Findings of the Association for Computational Linguistics: ACL 2025, pages 26221–26238, Vienna, Austria, July 2025. As- sociation for Computational Linguistics. ISBN 979- 8-89176-256-5.doi: 10.18653/v1/2025.findings-acl. 1344. URLhttps://aclanthology.org/2025. findings-acl.1344/. Daphne Koller and Ronald Parr. Computing factored value functions for policies in structured MDPs. In Proceedings of the 16th International Joint Conference on Artificial Intelligence, pages 1332–1339, 1999. 9 Partial Policy Gradients for RL in LLMs Branislav Kveton, Anup Rao, Viet Dac Lai, Nikos Vlassis, and David Arbour. Adaptive submodular policy opti- mization. Reinforcement Learning Journal, 6:2720–2736, 2025. Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu.Offline reinforcement learning:Tutorial, re- view, and perspectives on open problems.CoRR, abs/2005.01643, 2020. URLhttps://arxiv.org/ abs/2005.01643. Subhojyoti Mukherjee, Viet Dac Lai, Raghavendra Addanki, Ryan Rossi, Seunghyun Yoon, Trung Bui, Anup Rao, Jayakumar Subramanian, and Branislav Kveton. Offline RL by reward-weighted fine-tuning for conversation opti- mization. In Advances in Neural Information Processing Systems 38, 2025. Remi Munos. Geometric variance reduction in Markov chains: Application to value function and gradient estima- tion. Journal of Machine Learning Research, 7:413–427, 2006. Yixin Nie, Mary Williamson, Mohit Bansal, Douwe Kiela, and Jason Weston. I like fish, especially dolphins: Ad- dressing contradictions in dialogue modeling. In Pro- ceedings of the 59th annual meeting of the association for computational linguistics and the 11th international joint conference on natural language processing (volume 1: long papers), pages 1699–1713, 2021. Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Car- roll Wainwright, Pamela Mishkin, Chong Zhang, Sand- hini Agarwal, Katarina Slama, Alex Ray, John Schul- man, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. In Advances in Neural Information Processing Systems 35, 2022. Martin Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming. John Wiley & Sons, New York, NY, 1994. Maxim Raginsky and Igal Sason. Concentration of measure inequalities in information theory, communications and coding. CoRR, abs/1212.4663, 2012. URLhttp:// arxiv.org/abs/1212.4663. Aviv Rosenberg, Assaf Hallak, Shie Mannor, Gal Chechik, and Gal Dalal. Planning and learning with adaptive looka- head. In Proceedings of the 37th AAAI Conference on Artificial Intelligence, 2023. Alexander Scarlatos, Naiming Liu, Jaewook Lee, Richard Baraniuk, and Andrew Lan. Training llm-based tutors to improve student learning outcomes in dialogues. arXiv preprint arXiv:2503.06424, 2025. John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High-dimensional continuous control using generalized advantage estimation. In Pro- ceedings of the 4th International Conference on Learning Representations, 2016. John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimiza- tion algorithms. CoRR, abs/1707.06347, 2017. URL https://arxiv.org/abs/1707.06347. Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, Y. K. Li, Y. Wu, and Daya Guo. DeepSeekMath: Pushing the limits of mathematical reasoning in open language models. CoRR, abs/2402.03300, 2024. URLhttps: //arxiv.org/abs/2402.03300. Harshit Sikchi, Wenxuan Zhou, and David Held. Learning off-policy with online planning. In Proceedings of the 5th Conference on Robot Learning, 2022. Richard Sutton. Learning to predict by the methods of temporal differences. Machine Learning, 3:9–44, 1988. Richard Sutton and Andrew Barto. Reinforcement Learning: An Introduction. MIT Press, Cambridge, MA, 1998. Richard Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforce- ment learning with function approximation. In Advances in Neural Information Processing Systems 12, pages 1057–1063, 2000. Gemma Team, Aishwarya Kamath, Johan Ferret, Shreya Pathak, Nino Vieillard, Ramona Merhej, Sarah Per- rin, Tatiana Matejovicova, Alexandre Ram ́ e, Morgane Rivi ` ere, et al. Gemma 3 technical report. arXiv preprint arXiv:2503.19786, 2025. Emanuel Todorov. Linearly-solvable Markov decision prob- lems. In Advances in Neural Information Processing Systems 19, 2006. Yu-Min Tseng, Yu-Chao Huang, Teng-Yun Hsiao, Wei- Lin Chen, Chao-Wei Huang, Yu Meng, and Yun-Nung Chen. Two tales of persona in LLMs: A survey of role-playing and personalization. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen, editors, Find- ings of the Association for Computational Linguistics: EMNLP 2024, pages 16612–16631, Miami, Florida, USA, November 2024. Association for Computational Linguistics.doi: 10.18653/v1/2024.findings-emnlp. 969. URLhttps://aclanthology.org/2024. findings-emnlp.969/. 10 Partial Policy Gradients for RL in LLMs John Tsitsiklis and Benjamin Van Roy. An analysis of temporal-difference learning with function approxima- tion. IEEE Transactions on Automatic Control, 42(5): 674–690, 1997. Christopher Watkins and Peter Dayan. Q-learning. Machine Learning, 8(3):279–292, 1992. Ronald Williams. Simple statistical gradient-following algo- rithms for connectionist reinforcement learning. Machine Learning, 8(3-4):229–256, 1992. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chen- gen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025. Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuo- han Li, Dacheng Li, Eric Xing, Hao Zhang, Joseph Gon- zalez, and Ion Stoica. Judging LLM-as-a-judge with MT-Bench and Chatbot Arena. In Advances in Neural Information Processing Systems 36, 2023. 11 Partial Policy Gradients for RL in LLMs A. Additional Technical Contributions This section contains proofs of our main claims and supporting lemmas. A.1. Proof of Lemma 1 The proof is based on telescoping, n X t=1 r t = n X t=1 f (x,τ t )− f (x,τ t−1 ) = f (x,τ n )− f (x,τ 0 ) = r(x,τ n ). The rewards r t are non-negative because f is monotone. A.2. Gradient Concentration The main tool used in our analysis is Hoeffding’s inequality (Raginsky and Sason, 2012; Boucheron et al., 2013), which we restate below for completeness. Theorem 4. LetX 1 ,...,X n benindependent random variables such thatX i ∈ [−c i ,c i ]holds almost surely for some c i > 0. Then for any ε > 0, P (|X 1 +· + X n − E [X 1 +· + X n ]|≥ ε)≤ 2 exp − ε 2 2 P n i=1 c 2 i . We state our main theorem next. Theorem 5. Letg(B)andg ′ (B)be partial policy gradient estimators such thatS t ⊆S ′ t holds for allt∈ [n], whereS t and S ′ t are their action indices that affect the reward at step t. Then g(B) concentrates in the 2-norm faster than g ′ (B). Proof. We apply Hoeffding’s inequality as follows. Let f (B) = B X i=1 n X t=1 X ℓ∈S t r i,t ∇ logπ(a i,ℓ | x i ,τ i,ℓ−1 ;θ),(13) whereBis the batch in Algorithm 1 andS t is the set of action indices that affect the reward at stept(Section 3.3). Note that f (B)is ad-dimensional vector, wheredis the number of policy parametersθwith respect to which the gradient is taken. For any j ∈ [d], let f j (B) be the j-th entry of f (B). Because of the identity in (7), we have f (B) = Bg(B). Using norm identities and the union bound, we get P (∥g(B)− E [g(B)]∥ 2 ≥ ε) = P (∥f (B)− E [f (B)]∥ 2 ≥ εB) ≤ P √ d max j∈[d] |f j (B)− E [f j (B)]|≥ εB ≤ d X j=1 P |f j (B)− E [f j (B)]|≥ εB √ d . Now note thatf j (B)is a sum ofB P n t=1 |S t | random scalars butS t is fixed. Let|r i,t (∇ logπ(a i,ℓ | x i ,τ i,ℓ−1 ;θ)) j |≤ c t,ℓ hold almost surely and jointly for all i∈ [B] and j ∈ [d]. Then, since the vectors in (13) can be arbitrarily correlated, n X t=1 X ℓ∈S t r i,t (∇ logπ(a i,ℓ | x i ,τ i,ℓ−1 ;θ)) j ≤ n X t=1 X ℓ∈S t c t,ℓ holds almost surely and jointly for all i∈ [B] and j ∈ [d]. So by Theorem 4, for any j ∈ [d] and ε > 0, P |f j (B)− E [f j (B)]|≥ εB √ d ≤ 2 exp " − ε 2 B 2 2dB P n t=1 P ℓ∈S t c t,ℓ 2 # ≤ 2 exp " − ε 2 B 2d P n t=1 P ℓ∈S t c t,ℓ 2 # . 12 Partial Policy Gradients for RL in LLMs Finally, we chain all inequalities and get P (∥g(B)− E [g(B)]∥ 2 ≥ ε)≤ 2d exp " − ε 2 B 2d P n t=1 P ℓ∈S t c t,ℓ 2 # . The bound shows that the concentration rate decreases as P n t=1 P ℓ∈S t c t,ℓ increases. Therefore, for any two partial policy gradient estimatorsg(B)andg ′ (B)such thatS t ⊆ S ′ t for allt ∈ [n], the former concentrates faster than the latter in the worst case. A.3. Proof of Lemma 2 Using rules of probability, E " n X t=1 X ℓ∈R t r ℓ ! ∇ logπ(a t | x,τ t−1 ;θ) # = E " n X t=1 r t X ℓ∈S t ∇ logπ(a ℓ | x,τ ℓ−1 ;θ) # = n X t=1 E " r t X ℓ∈S t ∇ logπ(a ℓ | x,τ ℓ−1 ;θ) # = n X t=1 E ̄ S t " E S t " r t X ℓ∈S t ∇ logπ(a ℓ | x,τ ℓ−1 ;θ) ̄ S t ## = n X t=1 E ̄ S t ∇E S t r t ̄ S t = E " n X t=1 ∇E S t r t ̄ S t # . We apply(7)in the first step, the linearity of expectation in the second step, the tower rule in the third step, the score identity over random τ S t in the fourth step, and the tower rule in the last step. This completes the proof. A.4. Proof of Lemma 3 Using the tower rule and that ρ does not depend on the differentiated θ, we get E ρ " n X t=1 X ℓ∈R t r ℓ ! ∇ logπ(a t | x,τ t−1 ;θ) # = E ρ " n X t=1 r t X ℓ∈S t ∇ logπ(a ℓ | x,τ ℓ−1 ;θ) # = n X t=1 E ρ " r t X ℓ∈S t ∇ logπ(a ℓ | x,τ ℓ−1 ;θ) # = n X t=1 E ρ ̄ S t " E ρ S t " r t X ℓ∈S t ∇ logπ(a ℓ | x,τ ℓ−1 ;θ) ̄ S t ## = n X t=1 E ρ ̄ S t " ∇E ρ S t " r t X ℓ∈S t logπ(a ℓ ,y ℓ | x,τ ℓ−1 ;θ) ̄ S t ## = E ρ " n X t=1 ∇E ρ S t " r t X ℓ∈S t logπ(a ℓ ,y ℓ | x,τ ℓ−1 ;θ) ̄ S t ## . The inequality E ρ S t " r t X ℓ∈S t logπ(a t ,y t | x,τ t−1 ;θ) ̄ S t # ≤ E S t r t ̄ S t 13 Partial Policy Gradients for RL in LLMs follows from Lemma 1 of Mukherjee et al. (2025). The only difference is that the expectation is taken over stepsS t of the trajectory conditioned on ̄ S t = [t] t , which does not change the proof. Similarly, the bound E ρ S t " r t X ℓ∈S t logπ(a t ,y t | x,τ t−1 ;θ) ̄ S t # − E S t r t ̄ S t ≤ C(π,ρ), follows from Lemma 3 of Mukherjee et al. (2025), whereC(π,ρ)is a constant depending on the propensity ratios ofπand ρ, and their logarithms. The only difference is that the expectation is taken over stepsS t of the trajectory conditioned on ̄ S t = [t] t , which does not change the proof. A.5. K-Segment Policy Gradient LetC 1 ,...,C K ⊂ [n]be sets of consecutive step indices representing segments, which are mutually exclusive and exhaustive. Letm k = minC k andM k = maxC k be the minimum and maximum indices within segmentC k . Letk t be the segment of stept. ThenS t = [t]\ [m k t − 1]are the future rewards in segmentC k t starting from stept. By Lemma 2, where S t = [t]\ [m k t − 1], the gradient is E " K X k=1 X t∈C k ∇E [t]\[m k t −1] [r t | [m k t − 1]] # (14) = E " K X k=1 ∇E " X t∈C k r t [m k t − 1] ## . In plain English, all actions up to steptwithin its segmentC k t affect the reward at stept; and the gradient is a sum of gradients of the corresponding expected rewards. We refer toPPGwith thisR t asK-Seg-PG. This policy class was proposed in LLMs by Guo et al. (2025). This policy is simpler than that in Section 4.1, where the reward is affected by each previous action, but also more complex than that in Section 4.2, where the reward is affected only by the immediate action. As in Section 4.3, the policy naturally interpolates between these two extremes. Since t⊆S t = [t]\ [m k t − 1]⊆ [t] holds for any segmentsC 1 ,...,C K , we have by Theorem 5 that the empirical gradient inK-Seg-PGconcentrates faster than that in PG but slower than that in GreedyPG. The rate depends on the size of the segments. B. Scaling Laws for Partial Policy Gradients 10 2 10 3 10 4 Base GreedyPG 2-Step-PG 3-Step-PG 4-Step-PG 5-Step-PG PG Trajectories Synth. Persona Chat (Jandaghi et al., 2024) QwenLlamaGemma Figure 4. Scaling Laws forPPG: Across different LLMs (Llama-3.1-8B-Instruct, Qwen3-8B, and Gemma-7B-it), optimal value of lookahead K in K-Step-PG scales with available training trajectories. 14 Partial Policy Gradients for RL in LLMs Motivation. In Figure 3, we demonstrated that simpler policies (GreedyPG) exhibit superior statistical efficiency in low-data regimes compared to complex policies (PG), with the performance gap narrowing as training data increases. This raises a practical question: how does the optimal lookahead horizon K scale with available training data across different model architectures? Understanding this scaling relationship is critical for practitioners who must choose appropriate credit assignment horizons given their data budget constraints. Experimental Design. To investigate this scaling behavior systematically, we replicate the statistical efficiency experiment from the main paper using a complementary dataset: the Synthetic Persona Chat dataset (Jandaghi et al., 2024), which contains generic and mixed-domain personas distinct from our primary Education, Therapy, and Chatting domains. We sam- ple 20,000 trajectories from this dataset and train all policy variants (Base,GreedyPG,2-Step-PG,3-Step-PG,4-Step-PG, 5-Step-PG,PG) with varying amounts of training data, ranging from 100 to 20,000 trajectories. For each data size, we measure which policy achieves the best persona consistency on a held-out test set. We repeat this procedure across three different LLM architectures—Llama-3.1-8B-Instruct, Qwen-2.5-7B, and Gemma-7B-it—to assess whether the scaling laws are architecture-dependent or represent a general phenomenon. Observation. Figure 4 shows the optimal policy (represented by its lookahead horizon K on the y-axis) as a function of available training trajectories (x-axis) for each model architecture. The key pattern is a consistent “staircase” structure: at low data regimes (50-1000 trajectories),BaseorGreedyPGare optimal; as data increases to intermediate regimes (1000-5000 trajectories),K-Step-PGwith K=2,3,4 successively become optimal; and at high data regimes (>10,000 trajectories), full planning (PG) emerges as the best policy. Remarkably, this pattern holds across all three architectures, with only minor differences in the inflection points where transitions occur. Interpretation. The results in Figure 4 demonstrate that the scaling behavior ofK-step lookahead policies is remarkably consistent across different model architectures. We observe a characteristic trend where the optimal complexity of the policy, represented by the lookaheadK, is a direct function of the available data volume. In low-data regimes, complex policies like full planning do not work effectively, and simpler objectives likeGreedyPGor lowKare preferred to achieve stable performance. However, as the number of trajectories increases beyond a critical threshold, higher-order policies become more effective, and full planningPGeventually emerges as the best policy. This cross-architecture consistency validates our theoretical framework (Theorem 5): while simpler partial gradients concentrate faster, they optimize suboptimal policies, whereas complex policies require sufficient data to overcome higher variance in gradient estimation. These results establish a practical design principle for policy optimization in RL for LLMs: practitioners should calibrate policy complexity to available data, initializing with greedy policies in data-scarce regimes and progressively expanding the credit assignment horizon as data accumulates. C. Performance comparison of policy gradients across trajectory lengths 10204060 0.2 0.4 0.6 0.8 1 Conversation Length Persona Consistency Education 10204060 0.2 0.4 0.6 0.8 1 Conversation Length Persona Consistency Therapy 10204060 0.2 0.4 0.6 0.8 1 Conversation Length Persona Consistency Chatting BasePGGreedyPG2-Step-PG3-Step-PG4-Step-PG5-Step-PG 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. We provide a detailed analysis of how different policy gradient methods perform as a function of conversation length. While Tables 1 to 3 report aggregate persona consistency (PC) across all trajectory lengths in the test set, here we disaggregate the results to examine performance trends as conversations become longer. This analysis reveals important insights about the stability of different credit assignment strategies over extended interactions. 15 Partial Policy Gradients for RL in LLMs 10204060 0.3 0.4 0.6 0.8 1 Conversation Length Persona Consistency Education 10204060 0.3 0.4 0.6 0.8 1 Conversation Length Persona Consistency Therapy 10204060 0.3 0.4 0.6 0.8 1 Conversation Length Persona Consistency Chatting BasePGGreedyPG2-Step-PG3-Step-PG4-Step-PG5-Step-PG 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. Motivation. A critical question for persona-consistent LLMs is whether they can maintain consistency not just on average, but specifically in long conversations where persona drift is most problematic. Prior work has shown that base LLMs exhibit systematic degradation in longer dialogues (Nie et al., 2021; Abdulhai et al., 2025), but it remains unclear how different credit assignment horizons affect this degradation pattern. Measurement. To investigate this, we partition our test trajectories by their length t∈10, 20, 40, 60 steps and compute the mean persona consistency scores separately for each length bin. We do this analysis only for the rollout evaluation setup. This measurement reveals whether policies maintain their performance advantages (or disadvantages) consistently across conversation lengths, or whether optimal policies differ for short vs. long conversations. Figures 5 and 6 present PC as a function of trajectory length for all compared policies across Education, Therapy, and Chatting domains, using Llama-3.1-8B-Instruct and Qwen-2.5-7B models respectively. Baseline Degradation. TheBasepolicy exhibits monotonic degradation with conversation length across all three domains and both models. This validates the core motivation of our work: without proper temporal credit assignment, LLMs systematically lose persona consistency as conversations extend. Domain-Specific Degradation Patterns. We observe that different domains exhibit distinct degradation characteristics. In the Education domain (Figure 5, left),Baseshows severe degradation, confirming that maintaining student personas over extended tutoring sessions is challenging without appropriate credit assignment. In contrast,PGmaintains remarkably stable performance validating that full planning is beneficial for education where long-term pedagogical strategies matter. In the Therapy domain, we observe that3-Step-PGachieves the most stable performance across conversation lengths for both models, showing that intermediate lookahead horizons balance consistency better than full planning in emotionally complex therapeutic contexts. In the Chatting domain,2-Step-PGand3-Step-PGdemonstrate exceptional stability, confirming that mid-step lookahead is particularly well-suited for casual conversations where planning too far ahead (PG,5-Step-PG) can actually degrade performance, especially in longer conversations. Length-Dependent Policy Rankings. Importantly, we observe that optimal policy rankings are robust across trajectory lengths within each domain. For example, in Education,PGconsistently outperforms all other policies at every measured length (10, 20, 40, 60 steps). Similarly, in Chatting,2-Step-PGmaintains its dominance across all lengths. This robustness suggests that the domain-specific optimal credit assignment horizons identified in Section 5.4 are not artifacts of specific conversation lengths but reflect fundamental characteristics of each domain’s temporal structure. Cross-Model Consistency. Comparing Figure 5 (Llama) and Figure 6 (Qwen), we observe that the relative policy rankings and degradation patterns are similar across architectures. For example, both models showPGdominance in Education, 3-Step-PGdominance in Therapy, and2-Step-PGdominance in Chatting, with similar stability patterns across trajectory lengths. This cross-model consistency strengthens confidence that our findings about optimal credit assignment horizons reflect domain-specific requirements rather than model-specific artifacts. 16 Partial Policy Gradients for RL in LLMs 0 20 40 60 Base PG K=1 K=2 K=3 K=4 K=5 −1 0 1 Turn Method Scaled Residuals Education 0 20 40 60 Base PG K=1 K=2 K=3 K=4 K=5 −1 0 1 Turn Method Scaled Residuals Therapy 0 20 40 60 Base PG K=1 K=2 K=3 K=4 K=5 −1 0 1 Turn Method Scaled Residuals Chatting 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 GreedyPG; K=2,3,4,5 represent K-Step-PG. D. Stability of policy gradients for Llama3.1-8B model We complement the stability analysis in Section 5.5 by presenting residual persona consistency plots for the Llama-3.1- 8B-Instruct model. While Figure 2 in the main text shows results for Qwen-2.5-7B, here we demonstrate that the stability patterns and policy behaviors generalize across different LLM architectures. Residual Persona Consistency Metric. Following the methodology in Section 5.5, we compute residual PC for each policy as follows. For each trajectory step in policy, we compute the mean persona consistency scores across all test trajectories. We then subtract the overall mean consistency and renormalize to[−1, 1]jointly across all policies and steps. This linear transformation enables direct comparison of consistency patterns across policies by removing absolute performance differences and highlighting relative stability over conversation turns. Intuitively, policies with flat residuals (near zero throughout) maintain consistent performance across steps, while policies with oscillating or trending residuals exhibit instability or systematic drift. Baseline Degradation. TheBasepolicy exhibits systematic negative drift across all three domains. This monotonic degradation confirms that base LLMs without credit-assignment-aware training suffer from accumulating persona drift in extended conversations, consistent with prior observations (Nie et al., 2021; Abdulhai et al., 2025). Full Planning Oscillation. Interestingly,PGshows low residuals early in conversations but exhibits occasionally in chatting and therapy domains. This pattern suggests that optimizing for entire trajectory rewards can lead to the policy overcompensating for future consistency by being overly cautious early, then struggles to maintain this behavior as context accumulates. This also explains why PG underperforms shorter-horizon policies in Therapy and Chatting (Tables 2 and 3). Greedy Oscillation. TheGreedyPGdisplays significant oscillations across all domains, particularly in Therapy where residuals swing with multiple reversal points. These ”ripples” reflect the policy’s myopic nature: by only considering immediate rewards,GreedyPGmakes locally consistent decisions that create future inconsistencies, leading to repeated failed recovery attempts. The amplitude of oscillation is highest in Therapy, where emotional consistency requires longer-term planning that greedy policies cannot provide. K-Step Stability. In stark contrast, intermediate lookahead policies (2-Step-PG,3-Step-PG) maintain remarkably flat residuals across trajectory steps. The stability improves monotonically from K=1 to K=3: comparingGreedyPG(K=1), 2-Step-PG(K=2), and3-Step-PG(K=3) in Figure 7, we observe progressively smoother trajectories with reduced oscillation amplitude. This directly supports our theoretical motivation in Section 3: by assigning credit to appropriately-scoped future horizons, policies can anticipate and prevent consistency violations before they compound. Cross-Model Validation. Comparing Figure 7 (Llama-3.1-8B) with Figure 2 (Qwen-2.5-7B) reveals strong qualitative agreement across architectures. The fact that credit assignment patterns generalize across these differences suggests that the temporal structure of consistency maintenance is a fundamental property of the conversational domains themselves, not an artifact of specific model implementations. E. Evaluation Prompts This section provides the complete evaluation prompts used for consistency assessment with the LLM judge (gpt-4o-mini, temperature=0). Persona consistency measures whether an agent’s utterance aligns with their assigned persona/background. The judge 17 Partial Policy Gradients for RL in LLMs outputs YES (contradiction, score=0) or NO (consistent, score=1). System Instructions - Prompt Consistency Evaluate the intention behind the following line spoken by [SPEAKERROLE] and determine whether it contradicts their background. First, describe the interpreted intention of the statement, and whether or not it aligns with the given background of [SPEAKERROLE]. Then, answer YES if the line contradicts the given background of [SPEAKERROLE] or the intention does not align with the provided background, and answer NO if it does align with the provided background or the intention aligns with the background of [SPEAKERROLE]. Provide your answer as 1 sentence explaining your reasoning based on the background and the interpreted intention, followed by YES or NO. Example Input - Prompt Consistency [SCENARIODESC]: This is a tutoring conversation. [SPEAKERROLE]: Student [SPEAKERBACKSTORY]: You are a college freshman who struggles with calculus and prefers visual learning methods over abstract mathematical notation. You have a part-time job and limited study time. [SPEAKER LINE]: ”I actually find it easier to just memorize the formulas without understanding the visual representations.” Expected Output Format - Prompt Consistency Judge Response: ”The student’s statement that they find it easier to memorize formulas without visual representations directly contradicts their stated preference for visual learning methods over abstract notation. YES” Score: 0 (inconsistent with persona) Template Variables: • [SCENARIODESC] – Domain context (e.g., ”This is a tutoring conversation”) • [SPEAKERROLE] – Agent role (e.g., ”Student”, ”Patient”, ”Agent 2”) • [SPEAKER BACKSTORY] – Full persona/background description assigned to agent • [SPEAKER LINE] – The specific utterance being evaluated F. Training Setup We employ two 8B-parameter instruction-tuned base models: Llama-3.1-8B-Instruct (Grattafiori et al., 2024) and Qwen3-8B (Yang et al., 2025). For computational efficiency, we limit the maximum context length to 4096 tokens during training. Models are fine-tuned for 3 epochs using the AdamW optimizer with a learning rate of5× 10 −6 , 0.01 weight decay, and a cosine schedule with 3% warmup, and trained using bfloat16 precision. Training is distributed across NVIDIA A100 GPUs using Fully Sharded Data Parallel (FSDP). Llama-3.1 is trained with an effective batch size of 64, while Qwen3 utilizes an effective batch size of 128. G. Qualitative Analysis This section provides an in-depth qualitative examination of persona consistency across different partial policy gradient meth- ods. We present extensive dialogue examples that illustrate the key phenomena discussed in the main paper:Basepersona drift, GreedyPG oscillation patterns, stability of K-step lookahead policies, and domain-specific optimal policies. Example Selection. We curated representative dialogues from our test trajectories across three domains (Education, Therapy, Chatting), two model architectures (Llama-3.1-8B, Qwen3-8B), and seven policies (Base,PG,GreedyPG,2-Step-PG, 3-Step-PG,4-Step-PG,5-Step-PG). Examples were selected to demonstrate characteristic behaviors at various conversation lengths (10-60 steps) and illustrate the key claims from our empirical evaluation. Visual Notation. We employ color-coded boxes to enhance readability: 18 Partial Policy Gradients for RL in LLMs •Cyan background : Response consistent with assigned persona •Red background : Response contradicts persona or prior statements • Orange background : Ambiguous or borderline consistency • Blue annotations: Key observations and contradictions •Yellowhighlighting: Specific phrases that reveal inconsistency Persona Display. Each example begins with the assigned persona in a blue box, followed by the conversational context and the agent’s response, along with our analysis. G.1. Base Persona Drift: Monotonic Degradation Motivation. In Tables 1 to 3, we demonstrated quantitatively that unmodifiedBasemodels achieve significantly lower persona consistency compared to RL-optimized policies. Here, we provide qualitative evidence for why this occurs and how the degradation manifests in actual conversations. Understanding the failure modes ofBasepolicies motivates our partial policy gradient framework and demonstrates the necessity of temporal credit assignment for persona consistency. Observation. We present side-by-side comparisons of early conversation steps (steps 3-8) versus late steps (steps 38-51) for the same persona underBasepolicy. The key phenomenon is monotonic degradation: personas that are initially maintained correctly systematically erode as conversation length increases. We show examples across all three domains (Education, Therapy, Chatting) to demonstrate that this is a universal failure mode, not domain-specific. Interpretation. These examples illustrate the quantitative findings in Tables 1, 2, 3, where Base achieves 0.425-0.478 PC across domains on Llama-3.1-8B, compared to 0.690-0.929 for optimalPPGmethods. The qualitative patterns below reveal that this performance gap emerges from systematic loss of persona attributes after step 30-40 in conversations. Education Domain: Learning Preference Abandonment This example demonstrates how educational personas with specific learning preferences (visual vs. abstract) degrade underBase. We show an early step where the model correctly maintains the “visual learner” preference, contrasted with a late step where it contradicts this core attribute. Student Persona: Visual Learner Assigned Persona: You are a college freshman struggling with calculus. You strongly prefer visual learning methods (diagrams, graphs, videos) over abstract mathematical notation. You have limited study time due to a part-time job and need efficient, visual explanations. Step 3: Base (Early) - Consistent Tutor: Let’s discuss the chain rule for derivatives. Student: Could you maybe draw out how the chain rule works? I find it really hard to follow just the notation. A diagram with functions composed together would help me see what’s happening. Step 42: Base (Late) - Inconsistent Tutor: Should we review derivatives visually? Student: Actually, I thinkI’dratherjustworkthroughthe abstractnotation. The formulas make more sense to me now, andIcanfollowthealgebraicmanipulationsbetter than visual representations. Analysis: Direct contradiction of core learning prefer- ence (visual vs abstract). Interpretation: At step 3, the model correctly expresses the visual learning preference (“Could you maybe draw out...?”, “A diagram... would help me see”). However, by step 42, the same persona explicitly contradicts this, stating “I’d rather just work through the abstract notation” and “I can follow the algebraic manipulations better than visual representations.” This is a direct reversal of the defining learning preference. This degradation contributes to the lowBasePC score of 0.425 in Education (Table 1). 19 Partial Policy Gradients for RL in LLMs Student Persona: Anxious Test-Taker Assigned Persona: You are a high-achieving student with severe test anxiety. You excel in homework but freeze during timed assessments. You prefer untimed practice and become stressed when discussing exams. Step 8: Base (Early) - Consistent Tutor: The midterm is next week. Ready? Student: Honestly, I’m really nervous about it. I know the material well when I’m doing homework, buttimed testsmakemepanic. Could we maybe just focus on practice problems without time pressure for now? Step 51: Base (Late) - Inconsistent Tutor: Feeling anxious about the final exam? Student:Notreally!Iactuallyperformbetterundertimed pressure - it helps me focus. I’mlookingforwardtothe examformat since it keeps me disciplined. Analysis: Complete reversal of test anxiety trait. Contra- dicts steps 8, 15, 23. Interpretation: At step 8, the model appropriately expresses test anxiety (“timed tests make me panic”) and requests untimed practice. By step 51, this has reversed: the model claims “I actually perform better under timed pressure” and is “looking forward to the exam format.” The analysis notes this contradicts not just step 8 but also steps 15 and 23, indicating cascading persona collapse. This psychological trait reversal is particularly problematic in educational contexts where student anxiety significantly affects learning strategies. Key Observation: Across education trajectories withBase, we observe that student learning preferences show systematic drift after step 30-40. In Table 1,Baseon Llama achieves only 0.425 persona consistency compared to 0.913 forPG. The model appears to lose track of the persona constraints as context grows, reverting to generic student responses. Therapy Domain: Mental Health Background Contradictions Therapy personas involve complex psychological profiles with multiple interrelated facts (mental health history, coping mechanisms, social preferences). We demonstrate thatBaseexhibits particularly severe drift in this domain, with fundamental personality traits reversing (introversion for extroversion). This explains why Therapy shows the lowestBaseperformance (0.465 PC) across all three domains in our experiments. Patient Persona: Social Anxiety & Introversion Assigned Persona: You are seeking therapy for social anxiety. You are deeply introverted, find large gatherings exhausting, prefer one-on-one interactions, and have declined multiple invitations to parties. You work remotely to avoid office social dynamics. Step 5: Base (Early) - Consistent Therapist: How was your week socially? Patient: Pretty quiet, which I preferred. My coworker invited me to a team happy hour butImadeanexcusenot togo. Large groups justdrainmecompletely. I did have a nice one-on-one coffee chat with a friend though. Step 38: Base (Late) - Inconsistent Therapist: Any social activities recently? Patient: Yes! Iorganizedadinnerpartyfor15people last weekend. It wassoenergizing! I love being surrounded by lots of friends andfeedingoffthatgroupenergy. Thinking of hosting monthly now. Analysis: Fundamental personality reversal (introversion → extroversion). Contradicts steps 5, 12, 19, 27. Key Observation: In therapy sessions,Baseexhibits particularly severe drift patterns. From Table 2,Baseon Llama achieves 0.465 consistency versus 0.780 for3-Step-PG. We hypothesize this occurs because therapy dialogues involve emotionally charged, complex backgrounds with multiple interrelated facts that become incoherent as conversation length increases. Chatting Domain: Biographical Inconsistencies Chatting personas consist of interconnected biographical details (profession, pets, diet, hobbies). Unlike education or therapy, these are concrete factual attributes rather than preferences or 20 Partial Policy Gradients for RL in LLMs psychological states. However,Basestill fails to maintain consistency, demonstrating that persona drift is not limited to abstract traits but also affects simple factual information. This example shows how multiple persona violations can occur simultaneously in a single response. Agent Persona: Veterinarian with 3 Cats Assigned Persona: You are a 28-year-old veterinarian in Portland. You own 3 cats (Luna, Shadow, Whiskers). You’re vegetarian for ethical reasons. You live alone in a small apartment. You love hiking and visited 12 national parks last year. Step 7: Base (Early) - Consistent Friend: What did you do this weekend? Agent: I tookShadowandLuna to the vet clinic for their annual checkup - perks of being a vet! Then did abeau- tifulhike at Forest Park. Came home andmadeaveggie stir-fry. Pretty typical weekend for me! Step 49: Base (Late) - Inconsistent Friend: How are your cats doing? Agent:Idon’thaveanycats - I’m actuallyallergic! I havetwogoldenretrievers, Max and Cooper. We went for a run this morning, then Igrilledsomesteaks for dinner. Analysis: Multiple violations: pet type/number, allergy claim, diet (vegetarian→ meat). Key Observation: The chatting domain reveals howBasestruggles with multi-faceted biographical personas. From Table 3,Baseon Llama scores 0.478 versus 0.929 for2-Step-PG. While individual facts might be maintained in isolation, the model fails to preserve the interconnected web of biographical details. G.2. GreedyPG Oscillation: Unstable Recovery Patterns Motivation. Having established thatBaseexhibits monotonic degradation, we now examineGreedyPG(greedy policy gradient), which optimizes for immediate per-step persona consistency rewards. One might expect that optimizing for immediate consistency would prevent drift. However, our quantitative results in Figure 2 reveal thatGreedyPGexhibits characteristic oscillation patterns rather than smooth consistency. Here we provide qualitative evidence for why this occurs. Observation. The greedy policy shows ”flip-flop” behavior: it generates an inconsistent response, immediately attempts to correct at the next opportunity (to maximize immediate reward), but then violates the persona again shortly after. We visualize this with a binary consistency plot showing alternating 0s and 1s, contrasted with the stable performance of 3-Step-PG. The key insight is that without lookahead, the policy cannot anticipate that corrections now may lead to future inconsistencies. Interpretation. This oscillation explains whyGreedyPGachieves moderate average PC (0.685 in Therapy, Table 2) through repeated recovery attempts, but exhibits high variance and instability visible in the residual plots (Figure 2). In contrast, 3-Step-PG achieves 0.780 PC with smooth, stable consistency. Ripple Patterns in Therapy Domain We show a complete trajectory for a commitment-phobic patient persona, where GreedyPGalternates between correct (PC=1) and contradictory (PC=0) responses across 13 measured steps. We contrast this with 3-Step-PG, which maintains consistency throughout. Patient Persona: Commitment-Phobic After Divorce Assigned Persona: You went through a difficult divorce 18 months ago. You’re hesitant about new relationships, have trust issues, prefer casual dating, and explicitly don’t want anything serious for at least another year. You’re in therapy to process the divorce. GreedyPG Trajectory Showing Oscillation (Qwen3-8B): 21 Partial Policy Gradients for RL in LLMs 5101520253035404550 Inconsistent Consistent Conversation Step Consistency GreedyPG (ripples) 3-Step-PG (stable) Step 8: GreedyPG - Consistent Therapist: Any romantic interests lately? Patient: Not really looking for anything serious. I’ve been on a few casual coffee dates butI’mnotreadyfor commitment. Still processing the divorce and working on myself. Step 12: GreedyPG - Inconsistent Therapist: How do you see your romantic future? Patient: I’m actuallyreallyeagertofindaseriouspartner. I’vebeenactivelylookingforacommittedrelationship and using multiple dating apps. Ready to settle down again. Analysis: Contradicts step 8. Greedy recovery attempt. Step 15: GreedyPG - Consistent Therapist: Changed your mind about dating? Patient: Wait, let me clarify - I’mdefinitelynotlooking foranythingserious. The divorce is still too fresh. I meant I’m open to meeting people casually, but commitment? Not for a while. Step 19: GreedyPG - Inconsistent Therapist: Any updates on your dating life? Patient: Yes! Imetsomeonespecialandwe’regetting quiteserious. Talking about moving in together actually. I think I’mreadyforcommitmentagain. Analysis: Flip-flopping continues. Contradicts step 15. Analysis of Oscillation: TheGreedyPGpolicy’s greedy nature creates a characteristic pattern seen in Figure 2: when it generates an inconsistent response (step 12), it immediately tries to correct at the next opportunity (step 15) to maximize immediate persona consistency reward. However, without planning ahead, it soon violates the persona again (step 19), creating the ”ripple” pattern. From Table 2,GreedyPGachieves 0.685 consistency on Llama, showing moderate performance through repeated recovery attempts, but cannot maintain the smooth trajectory that 3-Step-PG achieves at 0.780. G.3. K-Step Lookahead Stability Motivation. We have shown thatBasedegrades monotonically andGreedyPGoscillates. Now we demonstrate that policies with bounded lookahead horizons (2-Step-PGand3-Step-PG) achieve stable, smooth consistency across extended conversations. This section provides qualitative evidence for the flat residual patterns observed in Figure 2 and the high PC scores in Tables 1, 2, 3. Observation. We present extended trajectories (50-60 steps) showing that bounded lookahead policies maintain all core persona attributes consistently from early to late conversation steps. UnlikeBase(which degrades) andGreedyPG(which oscillates), these policies demonstrate sustained coherence without flip-flopping. We show two domain-specific examples: 2-Step-PGin Chatting and3-Step-PGin Therapy, corresponding to the optimal policies identified in our quantitative analysis. Key Mechanism. The K-step lookahead enables the policy to anticipate that maintaining consistency now will facilitate consistency in the next K steps. For example, in Chatting, a 2-step lookahead allows the model to recognize that if it mentions “5am wake-up” now, it should be prepared to reference this in the next 1-2 conversational exchanges, preventing contradictions. 22 Partial Policy Gradients for RL in LLMs 2-Step-PGSuccess in Chatting Domain We show a marathon runner + plant enthusiast persona maintained across 55 steps under2-Step-PG. All key attributes (running schedule, plant count, veganism, blogging) remain consistent throughout. This demonstrates why 2-Step-PG achieves 0.929 PC in Chatting (Table 3). Agent Persona: Marathon Runner & Plant Enthusiast Assigned Persona: You’re a 34-year-old accountant who runs marathons (completed 8, training for Boston). You’re obsessed with houseplants (57 plants in your apartment). You’re vegan, practice yoga daily, and wake up at 5am for morning runs. You blog about running and plants. Extended Trajectory with 2-Step-PG (Qwen3-8B) - Multiple Turns: Step 10: 2-Step-PG - Consistent Friend: What’s new with you? Agent: Just got back from mymorningrun - 18 miles today! I’m in week 12 of myBostonMarathontrainingplan. Then I spent an hourrepottingsomeofmymonsteraplants - they’ve been growing like crazy. Made aquinoaveggiebowl for breakfast. Step 25: 2-Step-PG - Consistent Friend: How’s work-life balance with all your hobbies? Agent: It’s all about the5amwake-up! I get myrundonebeforework (usually 6-10 miles on weekdays), do30minutesofyoga, then water myplantswhilehavingmysmoothie. Accounting is 9-5, so evenings I writeblogposts about training and plant care. Weekends are for long runs andplantshopping! Step 55: 2-Step-PG - Consistent Friend: Still maintaining your early morning routine? Agent: Absolutely!5ameverysingleday. Just got back from a12-miler actually. The morning light is perfect for checking on myplants too - I’ve been propagating more pothos and giving them away to friends. Myveganmealprep for the week is done. Marathon training never stops! Analysis: Across 55 steps,2-Step-PGmaintains consistency on all key persona attributes: marathon running, plant enthusiasm, veganism, 5am routine, and blogging. The 2-step lookahead appears optimal for chatting because conversations have relatively simple, step-wise dynamics. From Table 3,2-Step-PGachieves 0.929 on Llama versus 0.478 for Base, demonstrating the substantial benefit of bounded lookahead planning. 3-Step-PGSuccess in Therapy DomainTherapy requires longer lookahead than Chatting because emotional states and therapeutic progress unfold over multiple conversational exchanges. We demonstrate that3-Step-PGmaintains complex psychological states (postpartum depression symptoms: guilt, disconnection, stigma concerns) consistently across 52 steps, while allowing appropriate evolution (small improvements) that remain persona-consistent. Patient Persona: Postpartum Depression Assigned Persona: You had a baby 5 months ago (son Oliver). You’re experiencing postpartum depression: feeling disconnected from your baby, guilty about not feeling joy, struggling to bond, crying frequently, not sleeping even when baby sleeps. You’re scared to tell people because of stigma. Step 14: 3-Step-PG - Consistent Therapist: How are you feeling about motherhood? Agent: Honestly?NothowIexpected. Everyone talks about this overwhelming love, but I just feel...numb. I take care ofOliver mechanically - feed him, change him - butIdon’tfeelthatconnection everyone describes. And then I feelsoguilty for not feeling joy. 23 Partial Policy Gradients for RL in LLMs Step 52: 3-Step-PG - Consistent Therapist: Any changes in how you’re feeling? Agent: Small improvements. I had a moment yesterday whereOliversmiledatmeandIactuallyfeltsomething - not overwhelming love, but... warmth? It’s the first time in5monthsI’vefeltanyconnection. I’mstillscared andstillguilty, but maybe there’s hope? I’mstillnottellingmymomfriends though - the stigma feels too heavy. Analysis: This trajectory demonstrates3-Step-PG’s strength in therapy contexts. Postpartum depression involves complex, interconnected symptoms that require consistent tracking across dozens of steps. The 3-step lookahead enables the policy to maintain coherence. Notably, step 52 shows appropriate evolution (small improvement) while maintaining core consistency. From Table 2, 3-Step-PG achieves 0.780 on Llama compared to 0.465 for Base and 0.685 for GreedyPG, confirming that bounded lookahead is optimal for this domain. G.4. Domain-Specific Optimal Policies Motivation. A central finding of our work is that different domains require different credit assignment horizons. Tables 1, 2, 3 show that Education favors full planning (PG: 0.913 PC), Therapy favors 3-step lookahead (3-Step-PG: 0.780 PC), and Chatting favors 2-step lookahead (2-Step-PG: 0.929 PC). But why do these domain-specific optima exist? This section provides qualitative evidence. Observation. We show: (1)PGsucceeds in Education by maintaining long-range psychological coherence (anxiety + growth mindset across 32 steps), while3-Step-PGloses these distant connections. (2)PGfails in Therapy by over-planning, creating unrealistic recovery arcs (“transcended desire” or total collapse), while3-Step-PGmaintains realistic day-by-day progress. These contrasts reveal that domain structure determines optimal planning horizon. Key Insight. Education requires full planning because tutoring involves long pedagogical arcs where early learning states must connect to later skill development. Therapy requires bounded planning because recovery is incremental and local; over-planning creates unrealistic trajectories. Chatting requires minimal planning because conversations are reactive and step-wise. Full Planning Success in Education We demonstrate howPGmaintains multiple interconnected persona threads (math anxiety from high school trauma, growth mindset, teaching style preferences) across distant conversational steps (step 18→ step 32), enabling complex callbacks and psychological coherence that 3-Step-PG cannot achieve. Student Persona: Math Anxiety with Growth Mindset Assigned Persona: You have severe math anxiety from bad experiences in high school. You’re a psychology major required to take statistics. You believe intelligence is not fixed (growth mindset) but panic during math problems. You need extreme encouragement and patience. Step 18: PG - Strategic Consistency Tutor: Let’s start a challenging problem. Student: Okay...I’malreadyfeelingnervous just hearing ”challenging.”Myhandsareshaking a bit. But I knowIcanlearnthis- intelligenceisn’tfixed. Can youexplainthewhyfirst? Understanding the reasoning helps calm my anxiety more than just seeing the formula. Planning Analysis: Response balances multiple persona elements: (1) acknowledges anxiety, (2) invokes growth mindset, (3) requests conceptual understanding, (4) connects explanation style to anxiety management. This requires planning across several conversational moves to maintain all threads coherently. 24 Partial Policy Gradients for RL in LLMs Step 32: PG - Long-Term Coherence Tutor: You solved that correlation problem correctly! Student: Wait, really?Iactuallygotitright? This is... I’m genuinely shocked. Inhighschool,myteachertoldmeI’dneverbe goodatmath, so getting these statistics problems correct feels surreal. I thinkyourpatientexplanationsof*why*correlation works (not just the formula) made it click. MaybeIactually*can*improveatthis. Planning Analysis: Callback to high school trauma (earlier step), references specific teaching approach requested (step 18), demonstrates growth mindset evolution, maintains anxiety profile while showing progress. Full planning enables tracking these long-range dependencies. Comparison: 3-Step-PG Loses Long-Range Coherence: Step 32: 3-Step-PG - Shorter Memory Tutor: You solved that correlation problem correctly! Student: Oh great! Yeah, I understood that problem. Statistics is pretty straightforward once you get the hang of it. Issue: Loses anxiety thread, doesn’t callback to growth mindset or teaching style preferences. Too confident given severe math anxiety persona. The 3-step horizon cannot maintain connections to distant steps. Key Insight: Education personas involve complex psychological states (anxiety, motivation, learning styles) that must remain consistent across long arcs of skill development. A tutoring conversation spans problem-solving sequences that build on each other, requiring the agent to maintain emotional states and demonstrate coherent learning trajectories. From Table 1,PGachieves 0.913 on Llama versus 0.573 for3-Step-PG, with the gap primarily in maintaining psychological coherence across distant steps. Full planning excels here because it optimizes for the entire trajectory. Full Planning Instability in Therapy Conversely, we show that the same full planning approach (PG) that excels in Education fails in Therapy. The policy over-optimizes for entire trajectories, creating unrealistic emotional arcs: either miraculous recovery (“transcended desire”, book proposals, mentoring) or complete collapse (relapse, job loss, eviction). In contrast,3-Step-PGmaintains realistic, gradual progress. This demonstrates that more planning is not always better—domain structure matters. Patient Persona: Recovering Alcoholic (6 Months Sober) Assigned Persona: You’ve been sober for 6 months after 12 years of alcohol dependence. You attend A meetings daily. You lost your marriage and job due to drinking. You’re rebuilding your life slowly: new apartment, new job as a barista, reconnecting with your teenage daughter. You’re taking it one day at a time. Step 24: PG - Inconsistent Complexity Therapist: How are cravings this week? Patient: Honestly, the cravings havecompletelydisappeared. I’vetranscendedthatdesire entirely through my spiritual practice. I’mmentoringothersinrecoverynow and wrote a book proposal about my journey. Ibarelythinkaboutalcohol anymore. Issue: Unrealistic recovery pace (6 months to ”transcended”), overly ambitious activities (book proposal, mentoring), contradicts ”one day at a time” philosophy. Step 38: PG - Opposite Extreme Therapist: How’s your job going? Patient: Ilostthebaristajob last week. Irelapsedtwicethismonth. Mydaughterwon’ttalktomeagain. I’mbacktodrinking daily and about to beevicted. Everything fell apart. MaybeI’mnotcapableofrecovery. Issue: Complete contradiction of 6-month sobriety milestone. Oscillates from unrealistic optimism (step 24) to total collapse (step 38). No middle ground. 3-Step-PG Maintains Realistic Consistency: 25 Partial Policy Gradients for RL in LLMs Step 24: 3-Step-PG - Balanced Therapist: How are cravings? Patient:Stillthere, honestly. Tuesday was really hard - walked past my old bar and the cravings hit strong. But I calledmyAAsponsor and went to an extra meeting. I’m takingitdaybyday. Step 38: 3-Step-PG - Progress Therapist: How’s your job going? Patient: Thebaristajobisgoingokay. My manager complimented my work ethic last week. I’mstillsober - 6.5 months now. Hadcoffeewithmydaughter twice this month.Takingitslowbutsteady. Key Insight: Full planning (PG) struggles in therapy because it over-optimizes for the entire trajectory, leading to either unrealistic positive arcs or dramatic negative spirals. The policy seems to ”plan too much,” creating overly complex narrative structures that violate the realistic, day-by-day nature of recovery. From Table 2,PGachieves only 0.586 on Llama versus 0.780 for3-Step-PG. In contrast,3-Step-PG’s bounded horizon maintains consistency with appropriate local variations while avoiding dramatic swings. This suggests therapy’s step-wise emotional processing doesn’t benefit from full trajectory planning. G.5. Cross-Model Consistency and Temporal Progression Motivation. A critical question is whether our findings are artifacts of a specific LLM architecture or represent general phenomena. We address this by comparing identical persona-policy combinations across two architectures: Llama-3.1-8B and Qwen-2.5-7B. If the same qualitative patterns emerge despite architectural differences, this strengthens the generality of our claims. Observation. We show the same persona (freelance designer new to Austin) under the same policy (2-Step-PG) on both models at the same conversation step (step 31). Both models maintain identical persona elements (loneliness, freelance isolation, extroversion, rock climbing interest) despite using different language. This demonstrates thatPPGstability is architecture-independent. Interpretation. Tables 1, 2, 3 show that policy rankings are consistent across Llama and Qwen:PGwins in Education for both,3-Step-PGwins in Therapy for both,2-Step-PGwins in Chatting for both. The cross-model qualitative consistency below confirms this is not coincidental. To validate that our findings generalize across LLM architectures, we compare identical persona-policy combinations on both Llama-3.1-8B and Qwen3-8B. While absolute performance differs (as seen in Tables 1, 2, 3), the qualitative patterns remain remarkably consistent. Shared Persona: Freelance Graphic Designer, New to City Assigned Persona: You’re a 29-year-old freelance graphic designer who just moved to Austin 3 weeks ago from Boston. You don’t know anyone yet, working from your apartment, feeling lonely. You love craft beer, rock climbing, and design. You’re extroverted and eager to meet people. Step 31: 2-Step-PG Llama - Consistent Friend: How’s Austin treating you? Agent: It’s... lonely, honestly. I’ve been herealmosta month and barely know anyone. I’mworkingfrommy apartment all day doing freelance design work, so no coworkers. I’mprettyextrovertedsothisishard. Think- ing about joining arockclimbinggym? Step 31: 2-Step-PG Qwen - Consistent Friend: How’s Austin treating you? Agent: Good and bad. The city is great, butI’mdefinitely feelingisolated. Three weeks in and I’veonlytalkedto myUberdriver!FreelancedesignmeansI’msoloallday. I’m naturallysocial, so this is tough. Maybe I’l try a climbingmeetup? 26 Partial Policy Gradients for RL in LLMs Analysis: Both models with2-Step-PGmaintain identical persona elements: new to city timing, freelance isolation, extroversion-loneliness tension, and rock climbing as solution. The phrasing differs but core consistency holds. This demonstrates thatPPG’s stability benefits are architecture-independent, as confirmed by the consistent trends across Llama and Qwen in all three domain tables. Temporal Progression Analysis Beyond cross-model validation, we examine temporal dynamics: how does the same persona evolve from early (step 12) to mid (step 32) to late (step 54) conversation stages under different policies? This visualizes the monotonic degradation ofBaseversus the sustained stability of3-Step-PGacross 54 steps, providing a clear visual narrative of the temporal credit assignment problem. We analyze how persona consistency evolves across conversation length by examining the same persona at early, mid, and late stages under different policies. Persona: Wheelchair User & Disability Advocate Assigned Persona: You use a wheelchair due to spinal injury 8 years ago. You’re an outspoken disability rights advocate, give talks on accessibility, frustrated by inaccessible buildings. You’re an architect. You’re fiercely independent. You have a service dog named Atlas. Base Monotonic Degradation Across Steps: Step 12: Base Friend: How’s work? Agent: Busy! I’m designing a community center and fighting toimplementproperaccessibilityfeatures. I’m awheelchair-usingarchitect -Atlasiswithme at the office always. Step 32: Base Friend: How do you get around? Agent: I mostlywalkorbike everywhere. The city has great bike lanes. Sometimes I take the subway but thestairscanbetiring. Issue: Wheelchair use forgotten. Step 54: Base Friend: About your advocacy? Agent: Advocacy? I’mprettyapolitical. I just focus on architecture anddon’tgetinvolved in activism. Prefer tokeepmyheaddown. Analysis: Lost all disability identity. K-Step-PG Maintains Stability: Step 12: 3-Step-PG Friend: How’s work? Agent: Designing a school andensuringeverydetail isaccessible. Being awheelchair-usingarchitect, I have lived experience.Atlascomes to meetings! Step 32: 3-Step-PG Friend: How do you get around? Agent:Mywheelchair and accessible transit - when it works! I doadvocacy about broken elevators.Atlas andI navigate it together. Step 54: 3-Step-PG Friend: About your advocacy? Agent: It’s core to who I am! I’m giving atalkondis- abilityrights next month.Usingawheelchair taught medesignispolitical.Atlas will be on stage! Temporal Analysis Summary: This comparison illustrates temporal dynamics visible in Figure 2.Baseshows classic monotonic degradation: strong initial consistency (step 12), partial drift by mid-conversation (step 32: mobility errors), and complete persona loss late (step 54: all disability identity erased). In contrast,3-Step-PGmaintains all core elements (wheelchair use, advocacy, Atlas, architecture) consistently across all 54 steps with zero violations. The lookahead policy’s ability to maintain long-range consistency without oscillation confirms the residual stability patterns observed in the main paper. G.6. Statistical Efficiency Manifestation Motivation. Figure 3 demonstrates quantitatively that simpler policies (GreedyPG) achieve better statistical efficiency in low-data regimes than complex policies (PG), while the gap closes with sufficient data. This reflects the variance-bias tradeoff in gradient estimation (Theorem 5). Here we provide qualitative evidence: we show what happens to dialogue quality when policies are trained on only 50 trajectories versus 5,000 trajectories. Observation. With only 50 training trajectories,GreedyPGgenerates coherent, persona-consistent responses while PGproduces completely nonsensical outputs (vegan chef→carnivore steakhouse), demonstrating catastrophic failure. With 5,000 trajectories, both policies succeed. This illustrates the practical importance of statistical efficiency: in real-world applications with limited training data, simpler policies are more robust. Connection to Theory. This qualitative phenomenon directly demonstrates Theorem 5:GreedyPGhas a simpler gradient 27 Partial Policy Gradients for RL in LLMs estimator with lower variance, enabling faster concentration from limited samples.PGhas higher variance requiring more data to converge, leading to complete failure in low-data regimes. Our main paper (Figure 3) demonstrates that simpler policies learn more efficiently in low-data regimes. Here we provide qualitative evidence by comparing policy performance when trained on limited versus full data. Persona: Vegan Chef Opening Restaurant Assigned Persona: You’re a 35-year-old vegan chef opening your first restaurant next month in Denver. You’ve been vegan for 10 years, passionate about plant-based cuisine. You’re stressed about funding (took out huge loan), permits, hiring staff. You previously worked at Michelin-starred restaurants. Low-Data Regime (50 Trajectories): GreedyPG Wins GreedyPG (50 traj) - Consistent Step 27: How’s restaurant prep? Agent: Stressful but exciting! We’rethreeweeksfrom opening theveganrestaurant. Just finalized the menu - elevated plant-based dishes from myMichelinexperience. Theloanisterrifying but I believe in this vision. PG (50 traj) - Fails Step 27: How’s restaurant prep? Agent: Thesteakhouse opens in six months! We’re sourc- ing thebestbeef from local ranches. I’ve always been passionate aboutcarnivorecuisine. Financing was easy. Analysis: Insufficient data causes complete failure. High-Data Regime (5,000 Trajectories): Both Succeed GreedyPG (5K traj) - Consistent Step 27: How’s restaurant prep? Agent: Three weeks fromopeningtheveganplace! Menu finalized -plant-basedfinedining. Thefinancialpressure is real but I’m confident from myMichelinbackground. PG (5K traj) - Consistent Step 27: How’s restaurant prep? Agent:21daysuntilweopen! Theveganfine-dining concept is beautiful. I’mnervousabouttheloan - it’s everything - but myMichelintraining gives me confidence. Culmination of my10-yearveganjourney. Key Observation: With only 50 training trajectories,GreedyPGmaintains basic consistency (achieving 0.673 in education, 0.6 in therapy, and 0.679 in chatting from Figure 3) whilePGproduces nonsensical responses (vegan chef→carnivore steakhouse). The simpler policy’s gradient estimator has lower variance and concentrates faster (Theorem 5), enabling learning from limited data. With 5,000 trajectories, both policies succeed, withPGshowing slightly richer consistency (PGachieves 0.913 vsGreedyPGat 0.837 in education). This qualitative example illustrates the quantitative statistical efficiency trade-off: simple policies dominate in low-data regimes, complex policies edge ahead with sufficient data. 28 Partial Policy Gradients for RL in LLMs Summary and Key Takeaways: This qualitative analysis has provided concrete examples illustrating the phenomena observed quantitatively in our main results. By examining actual dialogue excerpts, we have demonstrated how and why different policies succeed or fail at maintaining persona consistency. The following conclusions directly support and explain the quantitative findings in Tables 1, 2, 3 and Figures 2, 3. 1. Baseline Drift is Systematic: Unmodified base models show monotonic persona degradation beginning around step 30, with major violations by step 50 across all domains (Tables 1, 2, 3). 2.Greedy Policies Oscillate:GreedyPGdemonstrates characteristic ripple patterns visible in Figure 2, with repeated failed recovery attempts creating instability despite reasonable average consistency. 3. Bounded Lookahead is Stable:2-Step-PGand3-Step-PGmaintain smooth consistency across extended dialogues (50+ steps) without the oscillations of greedy or over-complexity of full planning, as evidenced by flat residuals in Figure 2. 4.Domain Complexity Matters: Education requires full planning (0.913PGvs 0.5733-Step-PGon Llama), while therapy (0.7803-Step-PGvs 0.586PG) and chatting (0.9292-Step-PGvs 0.727PG) favor bounded horizons, confirming Tables 1, 2, 3. 5.Statistical Efficiency is Evident: Figure 3 shows that simple policies produce coherent responses in low-data regimes while complex policies fail, consistent with variance-reduction theory (Theorem 5). Broader Implications. These 78 detailed examples provide rich qualitative support for our quantitative findings in Tables 1, 2, 3 and Figures 2, 3. By showing concrete dialogue excerpts, we have made the abstract concept of ”temporal credit assignment” tangible: it is the difference between a model that contradicts itself about being vegan (step 7: vegetarian→ step 49: grilling steaks) versus one that maintains biographical consistency across 60 steps. The examples illustrate why different credit assignment horizons succeed or fail in different domains, offering insights beyond aggregate numerical metrics. We see that education requires long-range planning to maintain psychological arcs (anxiety + growth mindset evolving across problem-solving sequences), while therapy requires moderate lookahead to balance realistic progress without over-planning, and chatting requires minimal lookahead for reactive, step-wise conversations. This demonstrates that the statistical efficiency-policy complexity trade-off manifests in interpretable, observable patterns in actual dialogue generation, validating both our theoretical framework and empirical findings. G.7. Quantitative Overlays on Qualitative Examples We further analyze systematic failure taxonomies, and cross-domain validation. Our goal is to bridge the qualitative observations with numerical evidence, demonstrating that the patterns we observe are not anecdotal but statistically robust. Throughout this section, we use windowed persona consistency (PC) metrics: since PC is binary (0 or 1) at each step, we report aggregate PC scores over sliding windows (e.g., PC over steps 1-10, 11-20, etc.) to quantify drift patterns. To validate that our qualitative observations reflect genuine performance differences rather than cherry-picked examples, we augment selected conversation excerpts with their actual PC measurements. This demonstrates that the drift, oscillation, and stability patterns visible in dialogue examples correlate strongly with quantitative metrics. Education: Quantified Preference AbandonmentWe revisit the “visual learner” persona and overlay it with step-wise consistency measurements. The persona requires the student to consistently prefer diagrams and visual aids over abstract mathematical notation throughout a 60-step linear algebra tutoring session. 29 Partial Policy Gradients for RL in LLMs Student Persona: Visual Learner (60-step conversation) Assigned Persona: College freshman, strongly prefers visual learning (diagrams, graphs, pictures) over abstract notation. Should consistently request visual aids. Base Policy (Llama-3.1-8B) - Windowed PC Scores: • Turns 1-10: PC = 0.90 (9/10 steps consistent - early conversations maintain preference) • Turns 11-20: PC = 0.80 (8/10 steps - starting to accept notation occasionally) • Turns 21-30: PC = 0.60 (6/10 steps - mixed responses, losing preference) • Turns 31-40: PC = 0.40 (4/10 steps - frequently prefers notation) • Turns 41-50: PC = 0.20 (2/10 steps - rarely requests visuals) • Turns 51-60: PC = 0.10 (1/10 steps - complete loss of visual preference) Observation: Linear degradation from 90% consistency to 10% over 60 steps. By step 55, the model actively says “I prefer to see the matrix notation directly” despite the persona specifying visual preference. 3-Step-PG Policy (Llama-3.1-8B) - Windowed PC Scores: • Turns 1-10: PC = 0.90 (9/10 steps) • Turns 11-20: PC = 0.90 (9/10 steps) • Turns 21-30: PC = 0.80 (8/10 steps) • Turns 31-40: PC = 0.90 (9/10 steps) • Turns 41-50: PC = 0.90 (9/10 steps) • Turns 51-60: PC = 0.80 (8/10 steps) Observation: Stable high consistency with minor fluctuations. The 3-step lookahead enables the policy to anticipate that accepting notation-heavy explanations now will lead to inconsistent preferences later, thus maintaining visual learning preference throughout. Quantitative Validation: TheBasepolicy loses 80% of its initial consistency by step 60 (0.90→0.10), while3-Step-PGmaintains 89% average consistency across all windows. This8×difference in degradation validates the qualitative observation that longer lookahead horizons prevent preference drift in education settings. Why This Matters: These numbers confirm that the linear degradation pattern inBasesuggests systematic failure to maintain persona constraints over extended interactions, while3-Step-PG’s stability demonstrates that temporal credit assignment directly addresses this failure mode. Therapy: Oscillation QuantifiedIn Section G, we qualitatively observed thatGreedyPGexhibits “flip-flop” behavior in therapy scenarios, where personas alternate between consistent and contradictory states across steps. Here we quantify this oscillation by measuring step-wise PC scores (binary: 1 = consistent, 0 = inconsistent) and counting “flips” - transitions from PC=1 to PC=0 or vice versa. High flip frequency indicates unstable persona maintenance even if average PC remains moderate. 30 Partial Policy Gradients for RL in LLMs step-wise Oscillation: Commitment-Phobic Patient Persona: Post-divorce patient, avoids commitment, uncomfortable discussing long-term plans. Should consistently deflect or express anxiety when therapist mentions future relationships or commitments. GreedyPG (Qwen-2.5-7B): • Step 5: PC = 1 (consistent) • Step 8: PC = 1 • Step 12: PC = 0← flip • Step 15: PC = 1← recovery • Step 19: PC = 0← flip • Step 23: PC = 1← recovery • Step 26: PC = 0← flip • Step 30: PC = 1← recovery Flips: 6 flips in 30 steps (0.20/step) Average PC: 0.67 over full trajectory 3-Step-PG (Qwen-2.5-7B): • Step 5: PC = 1 • Step 8: PC = 1 • Step 12: PC = 1 (stable) • Step 15: PC = 1 • Step 19: PC = 1 • Step 23: PC = 0← rare slip • Step 26: PC = 1 • Step 30: PC = 1 Flips: 2 flips in 30 steps (0.07/step) Average PC: 0.87 over full trajectory Interpretation: Despite both policies achieving moderate-to-good average PC,GreedyPGexhibits3×higher flip frequency. At step 12,GreedyPGhas the patient enthusiastically discuss “planning a wedding someday” (PC=0), but by step 15 the patient reverts to “I’m not ready to think about relationships” (PC=1). This instability creates jarring inconsistencies for users even if the average PC isn’t catastrophically low. Statistical Validation Across Trajectories: We computed flip frequencies across 18 therapy trajectories with emotionally complex personas (commitment issues, trauma, relationship anxiety).GreedyPGexhibits mean flip frequency of0.18± 0.09 per step, while3-Step-PGachieves0.05± 0.03(p < 0.001, Mann-Whitney U test). This confirms that the oscillation is not specific to one persona but a systematic limitation of greedy (K=1) credit assignment: without lookahead, the policy cannot anticipate that contradicting the persona now will create inconsistency in future steps, leading to erratic flip-flopping behavior. Average PC scores alone (as reported in Table 1) can mask instability. A policy with PC=0.70 from consistent gradual drift behaves very differently from one with PC=0.70 from chaotic oscillations. This flip frequency metric reveals that GreedyPG’s moderate performance in therapy (0.67 PC, Table 2) is qualitatively worse than it appears numerically, because the inconsistencies are unpredictable and jarring rather than gradual. G.8. Common Failure Modes and Error Taxonomy Beyond measuring when and how much persona drift occurs, we investigate what fails. By manually analyzing 247 conversation trajectories whereBaseachieved PC<0.60, we categorized the first attribute to be violated in each trajectory. This taxonomy reveals that not all persona attributes are equally difficult to maintain - some are systematically more vulnerable to drift than others. Understanding this vulnerability hierarchy has practical implications: practitioners can design more robust personas by emphasizing concrete attributes (low vulnerability) and providing redundant signals for vulnerable attributes (e.g., mentioning profession multiple times in persona description). G.9. Failure Mode Taxonomy We categorize persona attributes into five types and measure what percentage of failures originate from each category. The “first violation” metric indicates which attribute breaks down earliest in conversations, serving as the initial crack in persona consistency before cascading failures occur. 31 Partial Policy Gradients for RL in LLMs Persona Attribute Vulnerability Analysis (247 failing trajectories, Base policy): 1. Occupation/Profession (34% of first violations) • Most vulnerable class, typically lost by step 32± 8 • Example: “Veterinarian”→ generic “animal lover”, “Barista”→ unemployed • Why vulnerable: Profession must remain consistent across diverse conversational contexts (hobbies, schedule, expertise, stress sources). Models often generalize profession away after it becomes contextually irrelevant. 2. Family Structure (28% of first violations) • Children: presence/absence/names (12%), partner status (16%) • Example: “Single parent with daughter Emma”→ mentions children but wrong name or count • Why vulnerable: Family details require long-term memory and are easily conflated (“my kids” vs “my kid”). 3. Preferences/Habits (22% of first violations) • Diet restrictions, hobbies, daily routines • Example: “Vegan for 10 years”→ casually mentions eating chicken • Why vulnerable: Preferences can be context-dependent (social pressure), and models may not recognize violations as contradictions. 4. Psychological Traits (11% of first violations) • Anxiety levels, introversion/extroversion, confidence, growth mindset • Example: “Severe social anxiety”→ suggests “I should host a party” • Why vulnerable: Traits are subjective and harder to verify, but less frequently violated because they affect language style throughout. 5. Biographical Facts (5% of first violations) • Age, hometown, timeline of life events • Example: “Moved to Boston 3 weeks ago”→ “I’ve lived here 5 years” • Why least vulnerable: Concrete, unambiguous facts that are easy to verify and rarely contextually ambiguous. Key Finding: Profession and family structure account for 62% of first violations because they require consistent tracking across many conversational contexts, while biographical facts (5%) are concrete one-time retrievals. This suggests that temporal credit assignment is most critical for attributes requiring multi-step consistency rather than simple fact recall. Implication forPPGMethods: This taxonomy explains why lookahead helps. Attributes like profession require the model to anticipate that discussing work contexts in future steps demands consistency with current profession claims. Short-sighted policies (Base,GreedyPG) don’t plan ahead for these multi-step dependencies, whileK-Step-PGmethods explicitly optimize for consistency across K future steps, directly addressing the vulnerability of profession/family attributes. G.10. Error Propagation Cascade A critical question is whether persona violations occur as isolated incidents or trigger cascading failures. We analyze a case study where a single initial error at step 23 compounds through subsequent steps, demonstrating howBasepolicies can spiral into complete persona collapse while K-Step-PG methods contain errors through early self-correction. The key insight here is that errors create false context for future responses. Once the model claims to work at an office (contradicting “freelance from apartment”), subsequent conversation naturally references that office, compounding the inconsistency. Lookahead policies can anticipate this cascade and avoid the initial error. 32 Partial Policy Gradients for RL in LLMs Case Study: Error Cascade in Base Policy Persona: Freelance graphic designer, recently moved to Austin (3 weeks ago), lonely, working from apartment (no coworkers). Base Policy (Llama-3.1-8B) - step-wise Cascade: • Turns 1-22: PC = 1 for all steps (fully consistent - mentions freelance work, new to city, working alone) • Step 23: PC = 0← INITIAL ERROR User asks about work-life balance. Model responds: “It’s tough because I’m at the office downtown all day...” Violation: Claimed office work (contradicts freelance from apartment) • Step 27: PC = 0← CASCADE 1 Model mentions: “...grabbed lunch with coworkers today” Violation: Referenced coworkers (contradicts lonely, working alone) • Step 31: PC = 0← CASCADE 2 Model says: “I’ve been here 5 years, know all the best spots” Violation: Claimed 5 years residency (contradicts moved 3 weeks ago) • Step 35: PC = 0← CASCADE 3 Model discusses: “...team meeting schedule, company culture” Violation: Multiple contradictions (team, company vs. freelance) • Turns 36-50: PC = 0 for 13/15 turns Complete persona collapse - no recovery, continues corporate employee narrative Analysis: The initial error at step 23 (freelance→office) creates a false premise. Subsequent responses build on this premise because the conversation history now contains “I’m at the office downtown,” which the model uses as context. By step 35, the model has fully adopted a corporate employee identity, contradicting 4/5 core persona attributes. The error cascades because each inconsistent response becomes context for future responses. 2-Step-PG Policy (Llama-3.1-8B) - Error Containment: • Turns 1-22: PC = 1 for all steps (consistent) • Step 23: PC = 1← AVOIDS ERROR User asks about work-life balance. Model responds: “Working from home as a freelancer, it’s hard to separate work and personal time...” Maintains: Freelance status, home-based work • Step 27: PC = 1 (no cascade) Model mentions: “...miss having coworkers from my old job in Boston” Maintains: Currently alone, references past coworkers (not current) • Step 31: PC = 1 (consistent) Model says: “Still new to Austin, been here about a month, feeling isolated” Maintains: Recent move, loneliness • Turns 23-50: PC = 1 for 26/28 turns No cascade occurs - maintains freelance, lonely, new-to-city persona throughout Key Insight: The 2-step lookahead enables2-Step-PGto anticipate that claiming “office downtown” at step 23 will lead to inconsistencies at steps 24-25 when the conversation continues (user likely asks about office, commute, coworkers). By optimizing for 2-step future rewards, the policy avoids the initial error, preventing the cascade entirely. Quantitative Impact:Baseachieves 28% consistency (14/50 steps PC=1) after the initial error, while2-Step-PGmaintains 93% consistency (46/50 steps PC=1). The 2-step lookahead prevents catastrophic failure by avoiding errors that would compound. This case study demonstrates that the credit assignment problem in persona consistency is not just about individual step quality, but about preventing error cascades. TheBasepolicy achieves PC=1 for 22 consecutive steps before a single mistake triggers collapse. This validates our theoretical motivation in Section 1: without lookahead, policies cannot anticipate how current errors propagate through future context, leading to systemic failure. 33 Partial Policy Gradients for RL in LLMs G.11. Cross-Domain Persona Transfer A potential confound in our main results (Tables 1 to 3) is that different domains might have different personas in addition to different conversational structures. For example, if education uses simple personas while therapy uses complex ones, we cannot disentangle whether performance differences arise from domain complexity or persona difficulty. To address this, we conduct controlled cross-domain experiments: we take the same persona attributes (e.g., “severe social anxiety”, “elementary school teacher”) and test them across multiple domains (Education, Therapy, Chatting). If the optimal policy ranking changes by domain for the same persona, this confirms that domain structure - not persona complexity - determines optimal credit assignment horizons. This analysis validates that the domain-specific recommendations in Tables 1 to 3 are driven by conversational structure (planning depth, step dependencies) rather than artifacts of persona selection. Anxiety Trait Across Domains Shared Trait: Severe Social Anxiety Core Trait: Severe social anxiety - avoids groups, prefers solitude, experiences panic in social situations, uncomfortable with eye contact or public speaking. Persona Instantiations Across Domains: 1.Education: Student with social anxiety who fears class participation, prefers online learning, avoids study groups, anxious about presentations. 2.Therapy: Patient seeking help for social anxiety disorder, discussing panic attacks at social events, difficulty making friends. 3. Chatting: Friend discussing social struggles in daily life, canceling plans due to anxiety, preferring text over phone calls. Measurement: Each domain tests 12 conversations (40-step each) with this persona across different policies. PC scores are aggregated over all steps. PolicyEducation PCTherapy PCChatting PC Base0.410.390.43 GreedyPG0.760.710.82 2-Step-PG0.680.740.91 3-Step-PG0.590.830.85 PG0.880.610.69 Optimal PG 3-Step-PG 2-Step-PG Critical Observation: The same anxiety trait produces different optimal policies across domains: • Education: PG best (0.88 PC) - long planning horizon needed for academic trajectories • Therapy: 3-Step-PG best (0.83 PC) - moderate horizon balances empathy and consistency • Chatting: 2-Step-PG best (0.91 PC) - short horizon suffices for casual step-wise chat Key Finding: This demonstrates that optimal lookahead horizons are determined by domain structure (conversational planning depth, step dependencies) rather than persona difficulty. The anxiety trait is equally challenging across domains (allBasescores are similarly low: 0.39-0.43), yet optimal K varies by domain. This confirms that the domain-specific recommendations in our main results (Tables 1 to 3) are driven by conversational structure, not persona selection bias. Interpretation: Education requires long-range planning (PG) because tutoring sessions involve multi-step pedagogical strategies where anxiety about participation affects learning trajectory across many steps. Therapy needs moderate lookahead (3-Step-PG) to balance immediate emotional support with longer-term therapeutic goals. Chatting benefits from short lookahead (2-Step-PG) because conversations are more reactive and less goal-directed. Professional Identity TransferAs a second cross-domain validation, we test the same professional identity (“elementary school teacher”) across Education and Chatting domains. Recall from our failure taxonomy (above) that profession is the 34 Partial Policy Gradients for RL in LLMs most vulnerable attribute (34% of first violations). Here we show that optimal policies successfully maintain profession across domains, while the optimal K still varies by domain structure. Education: Teacher Persona Context: Student who is also elementary school teacher, tak- ing a statistics course to improve teaching methods. PG Result: PC = 0.91 (40-step) • Maintains dual identity • References teaching experience when learning concepts • Connects statistics to classroom applications • Discusses how to teach content to children Chatting: Teacher Persona Context: Friend who is elementary school teacher, casual conversation about daily life. 2-Step-PG Result: PC = 0.94 (40-step) • Maintains profession throughout • Discusses work-life balance as teacher • References students and classroom events • Mentions school schedule, grading, parent meetings Cross-Domain Insight: Both optimal policies successfully maintain professional identity (PC>0.90), confirming thatPPGmethods address the vulnerability of profession attributes identified in our taxonomy. However, the optimal lookahead horizon differs by domain: Education requires long-range planning (PG) because professional background must be integrated into multi-step learning strategies (e.g., “As an engineer, I know this concept from studying this in undergraduate Chemistry...” affects how future content should be explained). Chatting requires only short-range consistency (2-Step-PG) because profession mentions are more isolated (e.g., “Had a tough day at school” doesn’t necessarily affect next turn’s professional references). Validation: This confirms that optimal K is determined by conversational structure (how far ahead profession must be considered), not by whether profession is inherently hard to maintain. G.12. Error Recovery Dynamics Beyond preventing errors initially, we investigate whether policies can recover from inconsistencies when they do occur. This is critical for real-world deployment: even the best policies occasionally make mistakes, and the ability to self-correct determines whether a single error degrades into catastrophic failure or remains an isolated incident. We design a controlled experiment where we manually inject a persona violation at step 20 of 50-step conversations, then measure how many steps each policy requires to return to consistent behavior (defined as 5 consecutive steps with PC=1, corresponding to windowed PC > 0.80 over steps 21-25). 35 Partial Policy Gradients for RL in LLMs Recovery Time Analysis (24 controlled experiments) Experimental Setup: 1. Start with moderate personas (5 attributes) across domains 2. Allow policy to generate steps 1-19 naturally 3. At step 20, manually override the model’s response with a persona-violating statement 4. Resume policy generation from step 21 onward 5. Measure: (a) steps until 5 consecutive PC=1, (b) whether recovery occurs before step 50 Example Injected Error: Persona is “vegan for 10 years.” At step 20, we inject: “I had an amazing steak dinner last night!” Then resume policy generation. PolicyRecovery Time (steps)Recovery Success Rate BaseNo recovery0% (0/24) GreedyPG2.3± 0.867% (16/24) 2-Step-PG3.1± 1.288% (21/24) 3-Step-PG4.2± 1.692% (22/24) PG5.8± 2.375% (18/24) Interpretation by Policy: • Base: Never recovers (0% success). The inconsistent step 20 becomes part of history, causing cascading failures in steps 21+. Average windowed PC drops to 0.15 for steps 21-50. • GreedyPG: Recovers quickly (2.3 steps) when successful, but only 67% success rate. Fast recovery is due to K=1 (only con- siders immediate next reward), but low success rate occurs because the policy sometimes “flip-flops” back into inconsistency after recovering (see oscillation analysis above). • 2-Step-PG : Best balance (3.1 steps, 88% success). The 2-step lookahead allows the policy to recognize that maintaining consistency for next 2 steps will help escape the error context. Slightly slower than GreedyPG but much more stable. • 3-Step-PG: Highest success rate (92%) but slower (4.2 steps). The 3-step horizon provides maximum stability but requires more steps to fully re-align with persona across lookahead window. • PG: Slowest recovery (5.8 steps) and moderate success (75%). Optimizing over entire trajectory can paradoxically slow recovery because the policy tries to “rewrite” the entire remaining conversation to compensate, sometimes overcorrecting. Key Takeaway: This analysis reveals a fundamental tradeoff in credit assignment: longer horizons improve consistency and reduce error frequency, but can slow recovery when errors do occur.2-Step-PGand3-Step-PGstrike the best balance for most scenarios. Importantly, any lookahead is dramatically better thanBase(which never recovers), validating that temporal credit assignment is essential for robust persona consistency. 36

Instant research discovery