KRPO: Kalman Filter-Enhanced GRPO

• The paper introduces KRPO, which integrates a lightweight Kalman filter to dynamically estimate latent reward mean and variance for improved policy optimization.
• KRPO employs recursive Bayesian filtering to adaptively normalize advantages, reducing bias in non-stationary, high-variance reward environments.
• Empirical results demonstrate KRPO's enhanced stability and a 2–18% accuracy improvement over traditional GRPO in complex RL tasks.

Kalman Filter Enhanced GRPO (KRPO) is an advanced methodology for adaptive advantage estimation and uncertainty-aware state tracking within reinforcement learning-based optimization frameworks, particularly where Group Relative Policy Optimization (GRPO) is employed. Unlike conventional GRPO, which relies on static group-mean baselines, KRPO incorporates lightweight Kalman filtering to dynamically estimate the latent reward mean and its variance. This enables robust, adaptive normalization of advantages and results in improved stability and policy optimization performance, especially in environments with highly noisy, non-stationary, or difficult-to-model reward signals (Wang et al., 12 May 2025).

1. Motivation and Theoretical Foundations

Classical GRPO computes the advantage for each output by subtracting a group-wise mean reward, which serves as a baseline to reduce policy gradient variance. However, this method exhibits bias and decreased stability if rewards are non-stationary or possess high variance—a frequent scenario in LLM training and complex reasoning tasks (Wang et al., 12 May 2025). Static baselines cannot adapt to local drifts or the underlying temporal correlations in the reward stream, which compromises optimization efficacy.

KRPO addresses this by treating observed rewards as noisy measurements of an unobservable, smoothly-evolving latent reward. The Kalman filter offers an optimal Bayesian estimator for such scenarios, recursively inferring both the reward mean and its associated uncertainty from the sequential reward data. This ties KRPO’s technical basis directly to the theory of recursive Bayesian filtering and variance reduction in stochastic optimization (Vuckovic, 2018, Vilmarest et al., 2020).

2. Kalman Filter Baseline Estimation

The KRPO method instantiates a one-dimensional Kalman filter to dynamically estimate the latent reward mean $\hat{x}_{i|i}$ and variance $P_{i|i}$ at each training step $i$. The update equations are:

• Prediction:

$$\begin{aligned} \hat{x}_{i|i-1} &= \hat{x}_{i-1|i-1} \ P_{i|i-1} &= P_{i-1|i-1} + Q \end{aligned}$$

where $Q$ is the process noise variance (hyperparameter).

• Measurement Update:

$$\begin{aligned} K_i &= \frac{P_{i|i-1}}{P_{i|i-1} + \sigma^2} \ \hat{x}_{i|i} &= \hat{x}_{i|i-1} + K_i (r_i - \hat{x}_{i|i-1}) \ P_{i|i} &= (1 - K_i) P_{i|i-1} \end{aligned}$$

where $r_i$ is the observed reward, and $\sigma^2$ is the measurement noise variance.

The expected baseline for advantage normalization is thus a filtered, recursively updated mean rather than a single epoch-level average. The Kalman gain $K_i$ adaptively determines weightings based on estimation uncertainty and measurement noise, ensuring that the baseline evolves as the reward distribution shifts.

3. Advantage Normalization and Adaptive Scaling

Once the baseline and its associated uncertainty are computed, the advantage estimate for each timestep is defined as:

$$A_i = \frac{r_i - \hat{x}_{i|i}}{\sqrt{P_{i|i} + \varepsilon}},$$

where $\varepsilon$ is a small constant for numerical stability.

This normalization provides two critical effects:

• The reward is adaptively centered by the Kalman mean, removing non-stationary bias.
• The normalization by $\sqrt{P_{i|i} + \varepsilon}$ tempers the advantage if the uncertainty in the baseline is high, reducing extreme updates that could destabilize learning (Wang et al., 12 May 2025).

This contrasts with bandwise centering in standard GRPO, which disregards sequential dependencies and does not explicitly account for the confidence in the reward baseline.

4. Empirical Performance and Applications

KRPO has been evaluated on LLM reinforcement learning tasks, specifically math question answering and reasoning datasets such as Arithmetic and OpenMath-Instruct. Reported results demonstrate that:

• On the Arithmetic set, KRPO led to a 2–5% increase in model accuracy versus GRPO, across easy to hard question subclasses.
• On OpenMath-Instruct, which features higher reward variance and linguistic complexity, gains were even larger: up to approximately 17–18% on the hardest subsets.
• Reward learning curves show faster convergence and higher final attained reward under KRPO than GRPO.

These results indicate that the robustness of the Kalman-filtered baseline directly translates into improved policy optimization and reasoning performance, particularly when reward signals are dynamically evolving and volatile.

5. Computational Characteristics and Deployment Considerations

KRPO is explicitly designed to introduce minimal computational overhead. The scalar Kalman filter operations are lightweight and do not require additional trainable parameters or learned value networks, distinguishing KRPO from methods like actor-critic or learned value function baselines (Wang et al., 12 May 2025). This pragmatic design ensures that the variance reduction and adaptivity advantages are accessible even in large-scale or resource-constrained RL pipeline deployments.

However, hyperparameters such as process noise $Q$ and measurement noise $\sigma^2$ require tuning. Excessively aggressive or conservative choices may cause the baseline to respectively overfit to recent rewards or fail to adapt to genuine shifts in the reward distribution.

6. Broader Algorithmic Context and Theoretical Links

KRPO is part of a broad trend of reformulating adaptive variance reduction and robust estimation in RL and online optimization as recursive filtering tasks. This connection is theoretically justified:

• Kalman filtering for advantage estimation is analogous to using second-order, uncertainty-aware updates in stochastic optimization (Vuckovic, 2018, Vilmarest et al., 2020).
• The recursive Bayesian logic underlying KRPO aligns with both adaptive control theory and statistical estimation frameworks, which have demonstrated improved performance in non-stationary and high-variance regimes.
• The natural gradient and information-geometric viewpoints, wherein the Kalman filter can be understood as a particular preconditioned natural gradient in trajectory space, further ground the adaptation in modern RL theory (Ollivier, 2019).

7. Future Research Directions

Several avenues for further work are highlighted:

• Extension to higher-dimensional or context-dependent baselines, for instance tracking multiple latent baselines per group or context window.
• Adaptive or learned tuning of process and measurement noise parameters to eliminate the need for hand-specified noise levels.
• Incorporation with more advanced policy optimization schemes or structured RL tasks beyond LLM reasoning.
• Synergistic use with auxiliary techniques such as reward shaping, dynamic curriculum, or context-aware filtering (Wang et al., 12 May 2025).

The integration of lightweight Bayesian filtering for baseline estimation suggests that similar filtering strategies could systematically improve robustness and learning efficiency in a broad class of RL algorithms and settings characterized by non-stationary and noisy reward observations.