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

arXiv:2506.13862v2 Announce Type: replace Abstract: In Reinforcement Learning (RL), regularization with a Kullback-Leibler divergence that penalizes large deviations between successive policies has emerged as a popular tool both in theory and practice. This family of algorithms, often referred to as Policy Mirror Descent (PMD), has the property of averaging out policy evaluation errors which are bound to occur when using function approximators. However, exact PMD has remained a mostly theoretical framework, as its closed-form solution involves the sum of all past Q-functions which is generally intractable. A common practical approximation of PMD is to follow the natural policy gradient or use actor-critic approaches, but this potentially introduces errors in the policy update. In this paper, we propose and analyze PMD-like algorithms for discrete action spaces that only keep the last $M$ Q-functions in memory. We show theoretically that for a finite and large enough $M$, an RL algorithm can be derived that introduces no errors from the policy update, yet keeps the desirable PMD property of averaging out policy evaluation errors. Using an efficient GPU implementation, we then show empirically on medium-scale RL benchmarks such as MinAtar that increasing $M$ improves performance up to a certain threshold after which the performance becomes close to that of exact PMD, reinforcing the theoretical findings that using an infinite sum might be unnecessary and that keeping in memory the last M Q-functions is a practical and theoretically grounded implementation of PMD.

Abstract:In Reinforcement Learning (RL), regularization with a Kullback-Leibler divergence that penalizes large deviations between successive policies has emerged as a popular tool both in theory and practice. This family of algorithms, often referred to as Policy Mirror Descent (PMD), has the property of averaging out policy evaluation errors which are bound to occur when using function approximators. However, exact PMD has remained a mostly theoretical framework, as its closed-form solution involves the sum of all past Q-functions which is generally intractable. A common practical approximation of PMD is to follow the natural policy gradient or use actor-critic approaches, but this potentially introduces errors in the policy update. In this paper, we propose and analyze PMD-like algorithms for discrete action spaces that only keep the last $M$ Q-functions in memory. We show theoretically that for a finite and large enough $M$, an RL algorithm can be derived that introduces no errors from the policy update, yet keeps the desirable PMD property of averaging out policy evaluation errors. Using an efficient GPU implementation, we then show empirically on medium-scale RL benchmarks such as MinAtar that increasing $M$ improves performance up to a certain threshold after which the performance becomes close to that of exact PMD, reinforcing the theoretical findings that using an infinite sum might be unnecessary and that keeping in memory the last M Q-functions is a practical and theoretically grounded implementation of PMD.

Machine Learning (cs.LG); Artificial Intelligence (cs.AI)

arXiv:2506.13862 [cs.LG]

(or arXiv:2506.13862v2 [cs.LG] for this version)

https://doi.org/10.48550/arXiv.2506.13862 Focus to learn more arXiv-issued DOI via DataCite

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.

aiqbench analysis

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

Related coverage