Categorical Distributional Reinforcement Learning with Kullback-Leibler Divergence: Convergence and Asymptotics
Tyler Kastner, Mark Rowland, Yunhao Tang, Murat A Erdogdu, Amir-massoud Farahmand
Abstract
We study the problem of distributional reinforcement learning using categorical parametrisations and a KL divergence loss. Previous work analyzing categorical distributional RL has done so using a Cramér distance-based loss, simplifying the analysis but creating a theory-practice gap. We introduce a preconditioned version of the algorithm, and prove that it is guaranteed to converge. We further derive the asymptotic variance of the categorical estimates under different learning rate regimes, and compare to that of classical reinforcement learning. We finally empirically validate our theoretical results and perform an empirical investigation into the relative strengths of using KL losses, and derive a number of actionable insights for practitioners.
BibTeX
@inproceedings{
kastner2025categorical,
title={Categorical Distributional Reinforcement Learning with Kullback-Leibler Divergence: Convergence and Asymptotics},
author={Tyler Kastner and Mark Rowland and Yunhao Tang and Murat A Erdogdu and Amir-massoud Farahmand},
booktitle={Forty-second International Conference on Machine Learning},
year={2025},
url={https://openreview.net/forum?id=f4qxkR6GQK}
}