Turn-Level Importance Sampling in RL

• Turn-level importance sampling is an RL technique that computes geometric means of token-level ratios over conversation turns to align policy updates with dialog structure.
• It mitigates gradient variance and performance collapse by incorporating clipping-bias correction and normalizing surrogate gradients across turns.
• Experimental evaluations on multi-turn benchmarks demonstrate that ST-PPO improves task accuracy and training stability in long-horizon dialog environments.

Turn-level importance sampling is a reinforcement learning (RL) technique specifically designed for optimizing LLMs in multi-turn dialog and reasoning tasks. Unlike traditional token-level importance sampling—which calculates sample weights at the granularity of individual output tokens—turn-level importance sampling aligns the optimization process with the turn structure intrinsic to conversational agents. This shift addresses critical instabilities in standard Proximal Policy Optimization (PPO), notably gradient variance and performance collapse in long-horizon, multi-turn environments (Li et al., 25 Nov 2025).

1. Definitions and Motivation

Let $\mathcal{D}$ denote the set of user queries. For a query $x \in \mathcal{D}$, the agent generates a full trajectory $y = (y_1, \ldots, y_{|y|})$. This trajectory is partitioned into turns $$y^k = (y_{t_k^{\rm start}}, \ldots, y_{t_k^{\rm end}})$$ for $k = 1, \ldots, K$, where agent and environment actions alternate (often tracked via a loss_mask). At each turn, the turn-level state is $s_k = (x, y^1, \ldots, y^{k-1})$ and the turn-level action is $a_k = y^k$.

Traditional PPO applies importance sampling at the token level:

$$w_t(\theta) = \frac{\pi_\theta(y_t | x, y_{<t})}{\pi_{\theta_\textrm{old}}(y_t | x, y_{<t})}$$

where $\pi_\theta$ and $\pi_{\theta_\textrm{old}}$ are the current and reference policies. However, such token-level ratios:

• Exhibit high variance, particularly in long-horizon tasks,
• Fail to respect the natural turn-based structure of dialogue,
• Contribute to unstable and collapsed training dynamics (Li et al., 25 Nov 2025).

The introduction of turn-level importance sampling addresses these issues by defining weights at the granularity of dialogue turns, thereby stabilizing credit assignment and subsequent updates.

2. Formulation of Turn-Level Importance Sampling

Turn-level importance sampling computes a geometric mean of token-level ratios across each turn, formalized as:

$$w_k^{\rm turn}(\theta) = \left( \frac{\pi_\theta(y^k | x, y^{<k})} {\pi_{\theta_\textrm{old}}(y^k | x, y^{<k})} \right)^{1/|y^k|} = \exp\left( \frac{1}{|y^k|} \sum_{t \in \text{turn } k} \log w_t(\theta) \right)$$

where $|y^k|$ is the length of the $k$-th turn.

The policy-gradient update in Turn-PPO is then constructed as:

$\nabla_\theta J_{\rm Turn-PPO}(\theta) = \E\left[ \frac{1}{|y|} \sum_{k=1}^K w_k^{\rm turn}(\theta) \cdot \frac{\hat A^k}{|y^k|} \cdot \nabla_\theta \log \pi_\theta(y^k | x, y^{<k}) \right]$

where $\hat A^k$ is the aggregate advantage over the turn, typically estimated via GAE or temporal-difference methods and summed over specified token indices.

This approach ensures that all turns—regardless of length—have comparable weight and that policy updates align with the episodic structure of dialog-based environments. Lemma 1 in (Li et al., 25 Nov 2025) demonstrates that this form of credit assignment ensures correct alignment between RL objectives and the structure of multi-turn reasoning.

3. Stabilization via Clipping-Bias Correction

Standard PPO employs a clipping mechanism to constrain policy updates, but off-policy samples with inaccurate value estimates lead to biased and high-variance gradients. Lemma 2 in (Li et al., 25 Nov 2025) decomposes the PPO gradient into three components: the policy-gradient, a term representing advantage estimation errors, and a bias term from clipping:

$C_{\rm token}(\theta) = \E\left[ \frac{1}{|y|} \sum_{t: t\notin B_{\rm token}} w_t \hat A_t \right]$

Empirical analyses reveal that $\|C_{\rm token}\|_2$ can spike during training, indicating dangerous, extreme off-policy updates even after clipping.

The ST-PPO algorithm addresses this instability by normalizing the surrogate gradient with the $\ell_2$ norm of the bias term:

$$\nabla J_{\rm ST-PPO} = \frac{\nabla J_{\rm Turn-PPO}}{\|C_{\rm turn}\|_2}$$

This normalization contracts the gradient magnitude under high clipping bias, curbing “gradient spikes” and improving optimization safety. This mechanism is applied analogously in S-PPO at the token level.

4. ST-PPO Algorithmic Structure

The ST-PPO pipeline follows these main steps (Li et al., 25 Nov 2025):

1. Sample queries and generate trajectories using an LLM-based agent.
2. Partition trajectories into turns using loss masks.
3. Compute both token- and turn-level importance ratios.
4. Estimate advantages via GAE and aggregate per-turn.
5. Form the surrogate Turn-PPO gradient, scaled by turn-level ratios and advantages.
6. Compute the clipping-bias correction term and normalize the gradient.
7. Update the policy using the scaling factor.
8. Update the critic via TD learning or GAE.

The clipping hyperparameter $\epsilon$ is typically set to 0.2, with an optional small KL penalty ($\alpha \approx 10^{-3}$) to further stabilize training.

5. Theoretical and Empirical Analysis

The combination of turn-level sampling and clipping-bias correction in ST-PPO provably preserves the optimality conditions of PPO. Lemma 1 formalizes the alignment of credit via geometric-mean ratios and aggregated turn-level advantages. Lemma 2 characterizes the exact nature of clipping bias, demonstrating that gradient normalization does not disturb the fixed points of the learning dynamics but adaptively reduces update size in the presence of high-risk, off-policy samples.

No closed-form variance bounds are provided, but empirical diagnostics demonstrate substantially reduced variance in gradient norms and improved training stability compared to token-level PPO [(Li et al., 25 Nov 2025), Fig. 2–3].

6. Experimental Results and Practical Recommendations

ST-PPO has been evaluated on multi-turn search benchmarks including Natural Questions (NQ), HotpotQA, and several medical multiple-choice QA datasets (MedQA, MedMCQA, PubMedQA, MMLU-Med, MedXpert). Key findings include:

• Token-level PPO and Group-level PPO yield mid-training performance collapse, requiring early stopping.
• ST-PPO and S-PPO maintain stable performance, avoiding collapse and achieving higher final task success.
• Clipping ratios and KL divergence remain lower for ST-PPO/S-PPO throughout optimization.
• On medical MCQA, ST-PPO (8B) achieves 49.9% average accuracy vs. 45.4% for token-PPO.
• ST-PPO remains robust when increasing off-policy update frequency, retaining stability where token-level approaches fail.

Recommended implementation details involve 8 H100 GPUs, FSDP with offloading, batch sizes up to 512, learning rates of $1\mathrm{e}{-6}$ (policy) and $1\mathrm{e}{-5}$ (critic), and careful management of on-policy versus off-policy updates to avoid excessive policy drift. Turn boundary detection is identified using the loss mask.

7. Impact, Limitations, and Extensions

Turn-level importance sampling, as embodied in ST-PPO, aligns the training granularity with the natural interaction structure of multi-turn language agents. This reduces variance and stabilizes off-policy updates, a critical property for leveraging large batches and reusing trajectories in expensive LLM training scenarios.

Limitations include the absence of closed-form variance bounds and the reliance on accurate turn boundary identification. A plausible implication is that further improvements may arise from adaptive turn segmentation, problem-specific advantage estimators, or more sophisticated forms of off-policy correction.

Turn-level importance sampling is currently most effective when combined with clipping-bias correction. As shown empirically, the two techniques address orthogonal sources of instability and should be applied jointly for optimal results (Li et al., 25 Nov 2025).