[algo, doc] feat: trust region sequence masking - (1) k3 KL avg and (2) veto for max criterion

[szrlee]

@article{li2025trustregion,
  title   = {Trust Region Masking for Long-Horizon LLM Reinforcement Learning},
  author  = {Li, Yingru and Liu, Jiacai and Xu, Jiawei and Tong, Yuxuan and Li, Ziniu and Wang, Baoxiang},
  journal = {arXiv preprint arXiv:2512.23075},
  year    = {2025},
  url     = {https://arxiv.org/abs/2512.23075}
}

Link: https://richardli.xyz/post/trust-region-masking/

Summary

• Add K1/K3 KL divergence estimators and modify token veto mechanism for rollout correction.

New Features

• Introduce sequence-level k1 ($|\mathbb{E}[\log(ratio)]|$, divergence-based) and k3 ($\mathbb{E}[r - \log(r) - 1]$, small-KL stable) rollout sequence masking modes.
• Expand factory presets across K1, K3, geometric, and token TIS modes (e.g., decoupled_k1_rs, bypass_pg_k3_rs).
• Rename internal helpers (rollout_is_weights → rs_statistic) and introduce a consistent decoupled_ prefix for decoupled presets.
• Refine the token veto mechanism: use $|\log(r)|$ to flag catastrophic tokens and drop any sequence that contains them.

Semantic Clarifications

• Clarify divergence vs geometric semantics: K1 targets $|\mathbb{E}[\log(r)]|$ (ideal 0.0, threshold ~0.001) while geometric modes target $\exp(\mathbb{E}[\log(r)])$ (ideal 1.0, threshold ~1.001).
• Document that K3 uses log_ratio_safe for stability and highlight threshold guidance across modes.

Documentation

• Add mathematical formulations for K1 and K3, including estimator × operating mode comparison tables.
• Provide a comprehensive preset reference for all 20+ factory methods and clarify naming/threshold semantics.

Recommendations

• Use k3 mode for average-based filtering since K3 is non-negative and unbiased estimator.
• Use token veto mechanism for max-based filtering since |log(ratio)| is a symmetric detector and both $ratio \gg 1$ and $ratio \ll 1$ indicate large divergence.

Test Plan

• Unit tests covering new k1 and k3 rollout rejection sampling modes.
• Unit tests for token veto mechanism using |log(r)| to reject sequences with catastrophic tokens.
• Run existing rollout correction test suites to confirm backward compatibility.

---

Solution: Trust Region Masking (TRM)

We propose Trust Region Masking (TRM), which excludes entire sequences from gradient computation if any token violates the trust region.

Why Sequence Masking Works

By masking at the sequence level, we ensure that:

1. The remaining samples come from policies that are provably close
2. The gradient estimator remains unbiased over the unmasked samples
3. The trust region bound applies to all included sequences

Masking Criterion

A sequence is masked if:
$$\max_{t} D_{\mathrm{KL}}(\pi_{\mathrm{roll}}(\cdot|c_t) | \pi_{\theta}(\cdot|c_t)) > \epsilon$$

where $\epsilon$ is the trust region threshold.

Exact Computation

The rigorous guarantee requires exact KL computation with stored logits from the rollout policy.

Sample-Based Approximation

In practice, storing full logits may be expensive. The paper proposes sample-based approximations using importance ratios $\rho_t = \frac{\pi_{\theta}(y_t|c_t)}{\pi_{\mathrm{roll}}(y_t|c_t)}$:

The $k_3$ Estimator (for average-based filtering)

$$k_3(\rho) = \rho - 1 - \log \rho$$

$\rho$	$k_1 = -\log \rho$	$k_3 = \rho - 1 - \log \rho$
0.5	0.69	0.19
1.0	0	0
2.0	−0.69	0.31
10	−2.30	6.70
100	−4.61	94.4

Properties of $k_3$:

• Non-negative: $\rho - 1 - \log \rho \geq 0$ for all $\rho > 0$
• Unbiased: $\mathbb{E}{y \sim \pi{\mathrm{roll}}}[k_3(\rho)] = D_{\mathrm{KL}}$
• Ideal for average-based filtering since $(1/T)\sum_t k_3(\rho_t)$ converges to true average KL

The $|\log \rho|$ Estimator (for max-based filtering)

For the max criterion, we need a symmetric detector since both $\rho \gg 1$ and $\rho \ll 1$ indicate large divergence:
$$|\log(100)| = |\log(0.01)| = 4.6$$

In contrast, $k_3$ is asymmetric: $k_3(100) = 94.4$ but $k_3(0.01) = 3.6$ (26× difference).

Caveat

Neither sample-based method provides a rigorous bound on $D_{\mathrm{KL}}^{\mathrm{tok,max}}$ — both are approximate detectors based on single samples per context.

[gemini-code-assist]

Code Review

This pull request introduces significant enhancements to the rollout correction mechanism by adding new KL estimators (K3) and group-level rejection sampling. It also refactors existing modes for better clarity, distinguishing between divergence-based k1 and ratio-based geometric methods. The changes are comprehensive, covering new configuration presets, core logic updates, and extensive documentation. My review focuses on improving the performance of the new group-level operations and ensuring the documentation is consistent and accurate. I've identified two areas in the documentation that require correction to prevent user confusion and two performance bottlenecks in the Python code where loops can be vectorized for a substantial performance gain on GPUs.

[szrlee]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces significant enhancements to the rollout correction mechanism by adding K3 KL estimators and group-level rejection sampling modes. It also includes a large number of configuration presets to make these new features easier to use, and updates the documentation accordingly.

My review focuses on the correctness of the new implementations and their documentation. I've found a couple of high-severity issues:

1. There is a discrepancy between the mathematical formula for group_k1 in the documentation and its actual implementation. The documentation should be corrected to reflect the K1 divergence calculation.
2. A metric name, rollout_rs_mean, is being ambiguously reused for both IS ratios and KL divergence values, which could lead to confusion during analysis. I've suggested a change to improve clarity.

Overall, the changes are a valuable addition, providing more sophisticated tools for handling off-policy correction. Addressing the identified issues will improve the correctness and usability of the new features.

[szrlee]

/gemini review

[gemini-code-assist]

Code Review

This pull request significantly expands the rollout correction capabilities by introducing new K1 and K3 KL divergence-based rejection sampling modes, including group-level aggregation. The changes involve updating documentation (docs/algo/rollout_corr.md, docs/algo/rollout_corr_math.md) to detail these new modes, their mathematical foundations, properties, and typical thresholds. Correspondingly, the RolloutCorrectionConfig class (verl/trainer/config/algorithm.py) is extended with numerous new preset methods for decoupled and bypass modes, combining K1, K3, and group-level rejection sampling with sequence-level or token-level importance sampling. The core logic for computing rejection masks (verl/trainer/ppo/rollout_corr_helper.py) is updated to support these new modes, including handling group_indices for group-level rejection and adjusting metric calculations. Review comments highlighted the need for robust handling of empty batches and corrected an issue in the standard deviation calculation for divergence modes by ensuring the lower bound for clamping is 0.0.

[szrlee]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces several new rejection sampling modes (k3, group_k1, group_k3) for rollout correction, along with a significant number of presets to make them easier to use. The changes are well-structured, with clear separation of concerns between configuration, implementation, and documentation. The refactoring to distinguish between divergence-based (k1, k3) and ratio-based (geometric) modes improves clarity. The core implementation in rollout_corr_helper.py is robust, handling different aggregation levels and metric calculations correctly. I've identified one potential high-severity issue related to input validation for group-based sampling modes that could lead to a runtime crash.

[szrlee]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces significant enhancements to the rollout correction mechanism by adding K3 KL estimator and group-level rejection sampling modes. The changes are comprehensive, including 32 new preset methods for easier configuration and consistent naming conventions. The documentation has been extensively updated to reflect these new features, with detailed mathematical formulations and usage guides.

The implementation of the new k1, k3, and group-level modes in rollout_corr_helper.py is well-structured and uses efficient scatter_add operations for group aggregations. The distinction between divergence-based modes and ratio-based modes is handled cleanly throughout the code and documentation.

I've found a high-severity correctness issue in the calculation of the K3 divergence. The current implementation can lead to negative divergence values under certain conditions, which violates a core property of K3 and could lead to incorrect rejection sampling behavior. I've provided suggestions to fix this in the review comments.

Overall, this is a great contribution that significantly expands the capabilities of the rollout correction framework. Once the identified issue is addressed, this PR will be in excellent shape.

[szrlee]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces significant enhancements to the rollout correction mechanism by adding K3 KL estimator and group-level rejection sampling modes. The changes are comprehensive, including the implementation of core logic, addition of 32 preset factory methods for easy configuration, and extensive updates to both user-facing and mathematical documentation. The refactoring for consistent naming and metric corrections improves clarity and correctness. My review found one issue in the metric calculation for group-level masking that could lead to incorrect reporting when group indices are not contiguous. Overall, this is a well-structured and valuable contribution.

[szrlee]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces significant new features for rollout correction, including K3 KL estimators and group-level rejection sampling, along with a comprehensive set of presets and updated documentation. The changes are well-structured and the new capabilities are powerful. My review identified a potential Out-Of-Memory issue in the implementation of group-level masking for both group_k1 and group_k3 modes. The current approach does not handle sparse group indices efficiently, which could lead to excessive memory allocation. I have provided specific code suggestions to address this by using a more robust method for handling group indices.

[szrlee]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces significant enhancements to the rollout correction mechanism by adding the K3 KL estimator and group-level rejection sampling modes. The changes are comprehensive, spanning core implementation, configuration, and documentation.

The implementation of the new k1, k3, group_k1, and group_k3 modes in verl/trainer/ppo/rollout_corr_helper.py is robust, efficient (using vectorized operations for group aggregations), and correctly handles edge cases like empty batches. The refactoring to distinguish between divergence-based and ratio-based rejection sampling improves clarity and correctness, particularly in metrics calculation like ESS and deviation from ideal values.

The addition of 32 factory methods in verl/trainer/config/algorithm.py provides a clean and extensive set of presets for users, and they appear to be correctly configured. The documentation has been thoroughly updated in rollout_corr.md and rollout_corr_math.md to reflect these new features, including their mathematical formulations, which is excellent.

Overall, this is a high-quality contribution. The code is well-structured, clearly written, and the new features are implemented correctly. I have no high or critical severity feedback.

[CLAassistant]

All committers have signed the CLA.