🤖 machine learning

Heavy-Ball Q-Learning with Residual Weighting Correction

This paper proposes a corrected heavy-ball Q-learning method with residual weighting, establishing its convergence and proving conditions for accelerated performance over standard Q-learning by utilizing a switched linear system representation and joint spectral radius analysis.

Original authors: Donghwan Lee

Published 2026-06-26

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Donghwan Lee

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 trying to teach a robot to navigate a maze to find the best treasure. The robot uses a method called Q-learning to figure out which moves are good and which are bad. It does this by constantly updating a "scorecard" (called a Q-value) for every possible move in every situation.

However, standard Q-learning can be slow. It's like a hiker who takes a step, realizes they went the wrong way, takes a step back, realizes they overshot, and takes another step back. They are constantly correcting themselves, but they move forward very slowly, especially if the maze is tricky.

This paper proposes a new, faster way for the robot to learn, called Heavy-Ball Q-Learning with Residual Weighting Correction. Here is how it works, broken down into simple concepts:

1. The "Heavy Ball" Idea (Momentum)

Imagine the robot isn't just a hiker, but a heavy bowling ball rolling down a hill.

Standard Q-learning is like a hiker who stops completely after every step to check their map. They are very careful but slow.
Heavy-Ball Q-learning gives the robot "momentum." If the robot is rolling in the right direction, it keeps rolling a bit even if it needs to make a tiny correction. It doesn't stop and start; it glides. This usually helps it reach the bottom of the hill (the solution) faster.

2. The Problem: The "Weighting" Trap

The paper points out a specific problem with using this "heavy ball" idea in the standard way.

In the real world, the robot doesn't see every part of the maze equally. Some paths are visited often, others rarely. This is called non-uniform sampling.
When you try to add "momentum" to the standard method with these uneven visits, the math breaks down. It's like trying to roll a bowling ball on a floor where some tiles are sticky and some are slippery. The ball gets stuck or wobbles unpredictably, and the "momentum" doesn't actually help speed things up. The standard math tools used to prove the robot will eventually win stop working.

3. The Solution: The "Correction"

The author introduces a clever correction to fix this.

Think of the robot's scorecard as having two parts: the "average" score (how well it's doing overall) and the "specific" scores (how it's doing in specific tricky spots).
The standard method messes up the "average" part when trying to add momentum.
The author's Correction acts like a special filter or a "leveling tool." It adjusts the math so that the "average" part of the scorecard behaves nicely, even when the robot is visiting some spots more than others.
By fixing this specific part of the math, the author creates a stable environment where the "heavy ball" (momentum) can actually work its magic.

4. The Proof: Why It's Faster

The paper doesn't just say "it feels faster." It uses a complex mathematical framework (involving "Switched Linear Systems" and "Joint Spectral Radius") to prove it.

The Analogy: Imagine the robot's learning process as a machine with many gears. The "Joint Spectral Radius" is a measure of how fast the whole machine can spin without breaking.
The paper proves that with their Corrected Heavy-Ball method, the "slowest gear" in the machine spins faster than the slowest gear in the standard machine.
Because the slowest part is faster, the entire process is guaranteed to finish sooner.

5. Real-World Testing

The author tested this on two types of problems:

Simple Tables: Where the robot knows every single move (like a small maze).
Complex Approximations: Where the robot has to guess based on patterns (like a huge maze where it can't memorize every tile).

In both cases, the Corrected Heavy-Ball method reached the solution significantly faster than the standard method.

In one test, the standard method took about 2,700 steps to get close to the answer.
The new method took only about 1,500 steps.
In another test with complex guessing, the new method reached the target in 15 steps, while the standard method took 28.

Summary

The paper says: "We found a way to fix the math so that we can safely give the learning robot a 'heavy ball' (momentum). This correction ensures that even when the robot is learning from uneven data, it doesn't get stuck. We proved mathematically that this makes the robot learn faster, and our experiments show it works in practice."

Key Takeaway: It's not just about adding speed (momentum); it's about fixing the underlying structure (the correction) so that the speed boost actually works and is mathematically guaranteed to be better than the old way.

Technical Summary: Heavy-Ball Q-Learning with Residual Weighting Correction

Problem Statement
Q-learning is a fundamental method in reinforcement learning (RL) for learning optimal Q-functions without an explicit transition model. While standard Q-learning converges based on the contraction property of the Bellman optimality operator, this contraction can be weak when the discount factor $\gamma$ is close to one, leading to slow convergence. Existing acceleration methods, such as momentum-based updates, PID-type iterations, and rank-one corrections, often lack rigorous theoretical guarantees that they outperform standard Q-learning, particularly in control settings where the active greedy policy changes dynamically. A specific challenge arises when applying heavy-ball momentum to Q-learning with non-uniform sampling weights: the standard weighting matrix destroys the "common eigenvector" property (specifically the all-ones vector) required for certain spectral analyses, making it difficult to certify acceleration theoretically.

Methodology
The paper proposes a Corrected Heavy-Ball Q-Learning (HBCQL) method and establishes its convergence and acceleration properties using a Switched Linear System (SLS) framework.

Residual Weighting Correction (CQL):
The authors first introduce a "corrected" Q-learning update (Algorithm 2) to address the loss of the common eigenvector in standard weighted Q-learning. They replace the standard diagonal sampling matrix $D$ with a correction matrix $H = \delta \Pi_d + D(I - \Pi_d)$ , where $\Pi_d = \mathbf{1}d^\top$ . This matrix $H$ is designed such that it maps the constant vector $\mathbf{1}$ to a scalar multiple of itself ( $H\mathbf{1} = \delta\mathbf{1}$ ), thereby restoring a common eigenvector for the mean dynamics of the algorithm. This correction preserves the fixed point $Q^\star$ while enabling a specific structural decomposition.
Heavy-Ball Momentum:
The method extends the corrected Q-learning by adding a heavy-ball momentum term $\eta(Q_k - Q_{k-1})$ (Algorithm 3). This transforms the first-order recursion into a second-order system.
SLS and Joint Spectral Radius (JSR) Analysis:
The core analytical tool is the representation of the Q-learning error dynamics as a Switched Linear System. The switching signal is determined by the greedy action selection in the Bellman operator. The convergence rate is characterized by the Joint Spectral Radius (JSR) of the associated family of switching matrices.
- Decomposition: The authors decompose the state space into the subspace spanned by the common eigenvector (or invariant subspace in the heavy-ball case) and its orthogonal complement.
- Acceleration Mechanism: The analysis shows that the JSR of the full system is the maximum of the JSR of the projected system (orthogonal to the common direction) and the spectral radius of the dynamics along the common direction. The correction matrix $H$ allows the dynamics along the common direction to be tuned independently via the parameter $\delta$ . The heavy-ball term $\eta$ is then shown to accelerate the convergence along this specific direction without negatively impacting the projected dynamics, provided $\eta$ is sufficiently small.
Extension to Linear Function Approximation:
The framework is extended to linear function approximation (LQL). A similar correction matrix $G$ is constructed in the feature space to ensure the existence of a common eigenvector direction (spanned by a vector $c$ such that $\Phi c = \mathbf{1}$ ). The same SLS and JSR analysis yields analogous convergence and acceleration results.

Key Contributions

Corrected Q-Learning (CQL): The proposal of a structural correction to the Q-learning update that restores the common eigenvector property for non-uniform sampling, enabling rigorous spectral analysis.
Theoretical Acceleration Guarantee: The derivation of sufficient conditions (Assumptions 4.1, 4.2, and 5.1) under which the HBCQL method has a strictly smaller JSR than both the standard Q-learning and the corrected Q-learning (without momentum). This provides a "certified rate" for acceleration in a specific analytically identifiable direction.
SLS/JSR Framework: The application of Switched Linear Systems and Joint Spectral Radius analysis to Q-learning, offering a complementary perspective to standard contraction arguments. This approach explicitly handles the switching nature of the greedy policy.
Extension to Function Approximation: Generalization of the correction and acceleration theory to linear function approximation settings.

Results

Convergence: The paper proves that under the stated assumptions, the error recursion converges uniformly exponentially to zero (or to a noise neighborhood in the stochastic case) with a rate bounded by the JSR.
Acceleration: Theoretical analysis (Theorem 5.1 and 10.1) demonstrates that for small momentum parameters $\eta$ , the JSR of the heavy-ball corrected system is strictly less than that of the standard or corrected baseline.
Numerical Illustrations:
- In a deterministic tabular setting, HBCQL reached a relative error threshold significantly faster than CQL.
- In a stochastic tabular setting, HBCQL reduced the median time to reach a specific error threshold compared to both standard Q-learning and CQL.
- In linear function approximation settings (both deterministic and stochastic), the heavy-ball corrected method (HBCLQL) consistently outperformed the standard and corrected linear Q-learning methods in terms of convergence speed.

Significance and Claims
The paper claims to provide a precise certificate of acceleration in one analytically identifiable direction (the common eigenvector direction), rather than a full-direction acceleration theorem. The authors emphasize that the contribution is a theoretical guarantee for a specific geometric direction, which offers new insight into the geometry of heavy-ball Q-learning.

The work distinguishes itself from existing momentum-based literature by:

Focusing on the control setting (where the policy changes) rather than just policy evaluation.
Using SLS and JSR analysis instead of fixed-policy spectral arguments or stochastic approximation bounds alone.
Providing a structural correction that makes the acceleration analysis tractable for non-uniform sampling, a scenario where standard heavy-ball Q-learning lacks a common eigenvector.

The authors note that the acceleration guarantee is conditional on the "projected gap" assumption (that the projected dynamics are faster than the common direction dynamics) and small momentum parameters. They do not claim universal acceleration for all parameter choices or all MDPs, but rather identify a regime where the method is theoretically guaranteed to be faster. The analysis is primarily conducted on deterministic mean dynamics, with extensions to stochastic settings via standard conditional-mean arguments.