🤖 machine learning

StaQ: a Finite Memory Approach to Discrete Action Policy Mirror Descent

This paper proposes and validates "StaQ," a finite-memory algorithm for discrete-action Reinforcement Learning that approximates Policy Mirror Descent by retaining only the last $M$ Q-functions, thereby achieving the theoretical benefits of error averaging without the intractability of an infinite sum while empirically demonstrating that a sufficiently large $M$ yields performance comparable to exact PMD.

Original authors: Alex Davey, Alena Shilova, Brahim Driss, Riad Akrour

Published 2026-08-03

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Alex Davey, Alena Shilova, Brahim Driss, Riad Akrour

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 robot to play a video game. In the world of Artificial Intelligence, this is called Reinforcement Learning. The robot learns by trying things, getting points for good moves, and losing points for bad ones. But here's the tricky part: the robot uses a "brain" made of a neural network, which is a bit like a fuzzy guesser. Sometimes, this brain makes mistakes when it tries to figure out how good a move really is. These mistakes can pile up, causing the robot to get confused, forget what it learned, or get stuck in a loop of bad habits.

To fix this, scientists use a technique called "regularization." Think of this as a gentle hand on the robot's shoulder, reminding it not to swing too wildly from one strategy to the next. It forces the robot to keep its new ideas somewhat similar to its old ones, smoothing out the learning process. This specific family of methods is known as Policy Mirror Descent. Theoretically, this is a super-powerful way to learn because it averages out all those fuzzy mistakes, leading to a perfect strategy. But there's a catch: doing this perfectly requires the robot to remember every single guess it ever made about the game, stretching back to the very first second it started playing. For a robot playing for millions of steps, that's like trying to carry the entire history of the universe in your backpack—it's impossible to hold that much memory.

So, the big question becomes: Can we get the benefits of remembering everything without actually carrying the whole history? This is exactly what the paper "StaQ: a Finite Memory Approach to Discrete Action Policy Mirror Descent" sets out to solve.

The authors, a team of researchers from France, propose a clever new algorithm they call StaQ. Instead of trying to remember the infinite past, StaQ suggests a simple rule: "Keep only the last $M$ memories." Imagine a robot that only remembers its last 300 guesses about the game. When it makes a new guess, it adds it to the list and immediately forgets the oldest one. It's like a stack of plates where you only keep the top few; if the stack gets too high, you slide the bottom one off.

The paper proves mathematically that this "forgetting" doesn't actually hurt the robot's performance, as long as $M$ (the number of memories kept) is large enough. In fact, the researchers show that this finite memory approach is almost identical to the perfect, theoretical version that remembers everything. The "averaging out" of mistakes still happens, but now the robot isn't weighed down by an impossible amount of data. It's a bit like realizing you don't need to read every book ever written to be smart; reading the last few hundred is enough to get the gist.

To test this, the team built a super-efficient version of StaQ that runs on powerful computer chips (GPUs). They put it to the test on a set of video game benchmarks called MinAtar, which are like mini-versions of classic arcade games. They ran the robot through up to 5 million time-steps (which is a lot of game time). The results were clear: as they increased the memory size $M$ , the robot got better and better at the games. Once they reached a certain threshold (around $M=300$ ), the robot performed just as well as the theoretical "perfect memory" version.

What's really cool is that this method is also incredibly fast. Because the robot doesn't have to do complex math to update its strategy (it just stacks the new memory on top of the old ones), it learns faster than other popular methods that try to approximate the perfect solution. The paper shows that StaQ is not just a theoretical idea but a practical tool that works in the real world of deep learning.

The researchers also looked at what happens if you keep too few memories. If $M$ is too small (like 1), the robot acts like it has no memory at all and struggles. But once you give it a decent chunk of history, the performance jumps up and stays high. They even found that for some games, adding a tiny bit of random "noise" to the robot's exploration helped it find the best moves faster, but the core magic was definitely the finite memory stack.

In short, this paper suggests that we don't need to be perfect to be great. By keeping a manageable, finite history of our past guesses, we can build AI agents that learn efficiently, avoid getting confused by their own mistakes, and play games better than ever before. It turns out that sometimes, knowing when to forget is just as important as knowing when to remember.

Technical Summary: StaQ: a Finite Memory Approach to Discrete Action Policy Mirror Descent

Problem Statement

In Reinforcement Learning (RL), Policy Mirror Descent (PMD) algorithms regularize policy updates using Kullback-Leibler (KL) divergence to penalize large deviations between successive policies. A key theoretical property of PMD is its ability to average out policy evaluation errors, which are inevitable when using function approximators like neural networks. The closed-form solution for exact PMD involves an infinite weighted sum of all past Q-functions. However, this infinite sum is generally intractable in practice, especially with non-linear function approximators that preclude closed-form policy expressions.

Existing practical approximations, such as Natural Policy Gradient or Actor-Critic (AC) approaches (e.g., TRPO, MDPO, ECPO), attempt to approximate this update. However, these methods often introduce policy update errors because they rely on optimizing a surrogate objective or performing gradient steps, which may not perfectly align with the theoretical PMD update. This creates a gap between the theoretical benefits of PMD (error averaging) and practical implementation, potentially leading to suboptimal convergence or higher error floors.

Methodology: StaQ

The authors propose StaQ (Stacked Q-functions), a PMD variant designed for discrete action spaces that addresses the intractability of the infinite sum while avoiding policy update errors.

Core Algorithm

Instead of maintaining an infinite history or approximating the policy via gradient descent, StaQ maintains a finite memory of the last $M$ Q-function estimates. The policy at iteration $k$ is defined as a Boltzmann distribution over a weighted average of these $M$ Q-functions:

$\pi_k \propto \exp\left( \frac{\alpha}{1 - \beta^M} \sum_{i=0}^{M-1} \beta^i q_{k-i} \right)$

where $q_{k-i}$ are the Q-function estimates, $\alpha$ and $\beta$ are temperature and weighting parameters derived from the KL regularization weight $\eta$ and entropy weight $\tau$ .

Implementation Details

Finite Memory Truncation: The algorithm implements a "First-In-First-Out" (FIFO) mechanism. After each policy evaluation, the weights of the new Q-function are pushed onto a stack of neural networks. If the stack exceeds size $M$ , the oldest Q-function is deleted.
Optimization-Free Update: Unlike Actor-Critic methods that require gradient steps to update the policy network, StaQ's policy update is optimization-free. It simply involves stacking the frozen weights of the Q-functions. The policy logits $\xi_k$ are computed directly from the weighted sum of the stored Q-functions.
Efficient GPU Parallelization: The authors leverage batched operations on GPUs to compute multiple Q-values in parallel. They pre-compute logits for the replay buffer to minimize the computational overhead of large $M$ during the policy evaluation phase.

Theoretical Analysis

The paper provides a convergence analysis for StaQ under the framework of Entropy-Regularized Policy Mirror Descent (EPMD).

Convergence: The authors prove that for a sufficiently large $M$ , StaQ converges to the optimal policy in the absence of evaluation errors.
Error Averaging: The analysis shows that the error floor in StaQ depends on the weighted average of the last $M$ evaluation errors. The deviation from the exact PMD error floor consists of extra terms that decay exponentially fast with respect to $M$ .
No Policy Update Error: Crucially, because the policy update is derived directly from the closed-form solution (truncated to $M$ ), StaQ introduces no additional errors from the policy update step, unlike AC approximations which suffer from optimization errors.

Key Contributions

Extension of PMD Analysis: The authors extend previous PMD analyses to Policy Iteration (PI) and demonstrate that the averaging effect of policy evaluation errors can be achieved using only a DKL penalization on the policy update, without requiring DKL regularization in the evaluation step.
StaQ Algorithm: They propose and analyze StaQ, a finite-memory PMD variant. They theoretically show that for large enough $M$ , the algorithm converges without policy update errors and approximates the error-averaging properties of exact PMD with exponentially decaying error terms.
Empirical Validation: They provide an efficient GPU implementation demonstrating that increasing $M$ improves performance up to a threshold, after which StaQ becomes indistinguishable from exact PMD (which they evaluate up to 5 million time-steps).
Runtime Efficiency: The study shows that for medium-sized environments (e.g., MinAtar), increasing $M$ has a marginal impact on runtime, making StaQ faster than Actor-Critic baselines due to the lack of policy optimization steps.

Experimental Results

The authors evaluated StaQ on MinAtar (5 environments) and Classic Control tasks (4 environments), comparing it against:

Baselines: Natural Policy Gradient (TRPO), SAC (discrete version), ECPO (M-Projection), MDPO (Direct Optimization), DQN, and M-DQN.
Exact PMD: Evaluated with $M=1000$ (effectively infinite sum within the 5M step window).

Findings:

Performance vs. Memory Size ( $M$ ): Performance improves almost monotonically with $M$ . On nearly every task, $M \ge 300$ yields performance virtually indistinguishable from Exact PMD.
Comparison to AC Methods: StaQ generally outperforms Actor-Critic baselines (TRPO, ECPO, MDPO) across tasks. The AC methods, which introduce policy update errors, lag behind StaQ with moderate $M$ (e.g., $M=50$ to $M=300$ ).
Runtime: StaQ is faster than AC baselines. For example, on MinAtar/SpaceInvaders, StaQ ( $M=100$ ) took ~10.4 hours compared to 13.9 hours for SAC, despite the additional memory overhead, because it avoids gradient steps for the policy.
Exploration: The authors note that while entropy regularization helps, StaQ (like other entropy-regularized methods) can struggle with exploration in sparse-reward environments like MountainCar without additional mechanisms (e.g., "sticky" actions or $\epsilon$ -greedy exploration).

Significance and Claims

The paper claims that StaQ offers a practical and theoretically grounded implementation of PMD for discrete action spaces. Its significance lies in:

Bridging Theory and Practice: It demonstrates that the desirable properties of exact PMD (specifically, the averaging of evaluation errors) can be recovered without the intractable infinite sum or the approximation errors of Actor-Critic methods.
Eliminating Policy Update Error: By using a finite-memory stack of Q-functions, StaQ avoids the "error floor" introduced by inexact policy updates in standard AC algorithms.
Scalability: The efficient GPU implementation suggests that keeping a finite history of Q-functions is computationally feasible and beneficial, challenging the notion that exact PMD is purely theoretical.

The authors conclude that StaQ provides a promising setting for testing other RL components (like policy evaluation techniques) because it isolates the policy update mechanism, ensuring that observed performance gains or failures are not confounded by policy optimization errors. They modestly note that while memory usage may become significant for large-scale tasks (like full Atari), future work could explore compressing old Q-functions to maintain efficiency.