🤖 machine learning

K-Score: Kalman Filter as a Principled Alternative to Reward Normalization in Reinforcement Learning

The paper proposes "K-Score," a method that replaces traditional reward normalization in policy gradient reinforcement learning with a 1D Kalman filter to recursively estimate reward means, thereby reducing variance and accelerating convergence in non-stationary environments.

Original authors: Zixuan Xia, Quanxi Li

Published 2026-04-28

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Zixuan Xia, Quanxi Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a puppy to sit.

In the world of Artificial Intelligence (specifically "Reinforcement Learning"), we teach "agents" (the digital puppies) by giving them rewards. If they do something right, they get a "treat" (a positive number); if they do something wrong, they get nothing or a penalty.

The Problem: The "Mood Swing" Effect

The problem is that the "treats" in AI training are often very inconsistent. One moment the agent gets a huge reward, and the next, it gets zero. This is like a puppy trainer who is incredibly enthusiastic one second and then suddenly becomes totally silent the next.

Because the rewards are so jumpy, the AI gets "confused." It doesn't know if it did something great or if it just got lucky. This confusion makes the learning process slow, shaky, and prone to "crashing"—where the AI suddenly forgets everything it learned and starts acting like a puppy again.

To fix this, engineers usually use "Reward Normalization." Think of this like a trainer who keeps a notebook of all the treats given so far and tries to average them out. It works okay, but it’s a bit "dumb"—it treats every treat the same and struggles when the puppy's behavior (or the trainer's style) suddenly changes.

The Solution: The "K-Score" (The Smart Trainer)

The authors of this paper propose something much smarter called K-Score, which uses a mathematical tool called a Kalman Filter.

Instead of a simple notebook that just averages everything, imagine the trainer now has a "Smart Intuition." This trainer doesn't just look at the average; they constantly ask themselves two questions:

"How much do I trust this specific treat?" (Is this a fluke, or is the puppy actually learning?)
"How much is the environment changing?" (Is the puppy growing up, or am I just having a weird day?)

The Kalman Filter is like a high-tech stabilizer on a camera. When you are filming a shaky video, the stabilizer doesn't just try to move the camera back to the center; it predicts where the movement is going and smooths it out in real-time, distinguishing between a purposeful pan and a random shake.

How it Works (The Metaphor)

The K-Score acts like a "Signal Filter" for the AI's brain:

The Noise: The random, jumpy rewards (the "shaky camera").
The Signal: The actual progress the AI is making (the "smooth video").

The Kalman Filter looks at the jumpy rewards and says: "Wait, that huge reward was probably just a lucky spike. Let's not overreact. But this steady increase in rewards? That looks like real progress. Let's adjust our expectations accordingly."

Why Does This Matter?

By using this "Smart Intuition" instead of a "Dumb Average," the researchers found that:

The AI learns much faster: It doesn't waste time chasing "fake" rewards or getting discouraged by "noisy" failures.
The learning is smoother: There are fewer "meltdowns" where the AI's performance suddenly drops.
It’s "Plug-and-Play": You don't have to rebuild the AI's brain; you just give it this smarter way of interpreting its rewards.

In short: K-Score turns a chaotic, jumpy training session into a smooth, steady conversation, helping AI learn more efficiently by helping it tell the difference between "noise" and "real progress."

Technical Summary: K-Score: Kalman Filter as a Principled Alternative to Reward Normalization in Reinforcement Learning

1. Problem Statement

Policy gradient methods in Reinforcement Learning (RL) are notorious for high variance in gradient estimates, driven by the stochasticity of both the environment and the agent's policy. To mitigate this, practitioners traditionally use reward normalization (typically Z-score normalization using running estimates of mean and standard deviation).

However, the authors identify three critical flaws in standard heuristic normalization:

Non-stationarity: It assumes reward statistics are stationary, whereas in RL, the reward distribution shifts as the policy improves.
Hyperparameter Sensitivity: It relies on manual tuning of decay rates and window sizes.
Lack of Uncertainty Modeling: It treats all observations equally and fails to explicitly account for the noise level in the reward signal.

2. Methodology

The authors propose K-Score, a method that replaces heuristic normalization with a 1D Kalman Filter to perform online, uncertainty-aware reward estimation.

A. Mathematical Framework:
The method models the true underlying reward mean ( $x_t$ ) as a latent stochastic process. The return signal ( $G_t$ ) is treated as a noisy observation of this latent state:

Process Model: $x_t = x_{t-1} + w_t$ , where $w_t \sim \mathcal{N}(0, Q)$ (represents how much the reward mean is allowed to drift).
Observation Model: $G_t = x_t + v_t$ , where $v_t \sim \mathcal{N}(0, R)$ (represents the noise in the observed return).

B. Recursive Estimation:
The Kalman filter uses a two-step process (Predict and Update) to maintain an optimal estimate of the mean ( $x_t$ ) and its uncertainty ( $P_t$ ). The Kalman Gain ( $K_t$ ) dynamically determines the balance between the prior estimate and the new observation based on the ratio of process noise ( $Q$ ) to observation noise ( $R$ ).

C. Normalization Rule:
Instead of standard Z-scoring, the normalized return ( $\hat{G}_t$ ) is calculated as:
$\hat{G}_t = \frac{G_t - x_t}{\sqrt{P_t} + \epsilon}$
This provides a baseline that is not only adaptive to shifts in reward scale but also scales the signal based on the current estimation uncertainty.

D. Implementation Variants:

Simple Kalman: Uses fixed $Q$ and $R$ values.
Adaptive Kalman: Dynamically updates the measurement noise $R_t$ using an exponential moving average of the squared residuals, allowing the filter to react to changing volatility.

3. Key Contributions

Principled Alternative: Moves reward normalization from a heuristic "running average" approach to a statistically grounded Bayesian estimation framework.
Architecture-Agnostic: The module is a lightweight, scalar-based preprocessing step that can be "plugged into" any policy gradient algorithm (REINFORCE, Actor-Critic, GAE, PPO) without changing the neural network architecture.
Uncertainty-Awareness: Introduces a mechanism to handle non-stationary reward distributions by explicitly modeling process and observation noise.

4. Experimental Results

The method was evaluated on CartPole-v1 (dense reward) and LunarLander-v2 (sparse, high-variance reward) using Actor-Critic, GAE, and PPO.

Convergence Speed: In PPO on CartPole, K-Score achieved a 3.97× improvement in convergence speed compared to Z-score normalization (77 episodes vs. 306 episodes).
Stability: The training curves for Kalman-normalized agents were significantly smoother, showing fewer "spikes and collapses" compared to the high-variance fluctuations seen in standard Z-score methods.
Adaptivity: The Adaptive Kalman variant (which updates $R_t$ online) consistently outperformed the Simple Kalman variant, proving that dynamic noise modeling is highly beneficial in complex environments.
Ablation: Results showed that smaller process noise ( $Q$ ) values generally accelerate convergence by preventing the agent from overreacting to noisy outliers.

5. Significance

K-Score represents a significant step toward making Reinforcement Learning more robust and sample-efficient. By treating reward normalization as a state-estimation problem rather than a simple statistical moving average, the authors provide a way for agents to "distrust" noisy signals early in training and "adapt" quickly when the environment or policy changes. This approach is computationally efficient (constant-time updates) and highly practical for real-world RL applications where reward signals are often noisy and non-stationary.