Stochastic Q-learning for Large Discrete Action Spaces

TL;DR

This work tackles the computational bottleneck of value-based RL in large discrete action spaces by introducing stochastic maximization, which searches for maxima within sublinear action subsets of size $O(\log(n))$. By replacing $\max$ and $\arg\max$ with $\text{stoch\,max}$ and $\text{stoch\,arg\,max}$, the authors develop Stochastic Q-learning, Stochastic Double Q-learning, StochDQN, and StochDDQN, with convergence guarantees and substantial wall-time speedups demonstrated in Gymnasium and MuJoCo benchmarks. They also provide a memory-based variant that improves approximation and show that a stable Q-function enables eventual convergence to the exact maximum, supported by contraction analyses. Overall, the methods deliver exponential reductions in computation with competitive or superior performance, enabling practical RL in large or unstructured action spaces and broad applicability across value-based algorithms.

Abstract

In complex environments with large discrete action spaces, effective decision-making is critical in reinforcement learning (RL). Despite the widespread use of value-based RL approaches like Q-learning, they come with a computational burden, necessitating the maximization of a value function over all actions in each iteration. This burden becomes particularly challenging when addressing large-scale problems and using deep neural networks as function approximators. In this paper, we present stochastic value-based RL approaches which, in each iteration, as opposed to optimizing over the entire set of $n$ actions, only consider a variable stochastic set of a sublinear number of actions, possibly as small as $\mathcal{O}(\log(n))$. The presented stochastic value-based RL methods include, among others, Stochastic Q-learning, StochDQN, and StochDDQN, all of which integrate this stochastic approach for both value-function updates and action selection. The theoretical convergence of Stochastic Q-learning is established, while an analysis of stochastic maximization is provided. Moreover, through empirical validation, we illustrate that the various proposed approaches outperform the baseline methods across diverse environments, including different control problems, achieving near-optimal average returns in significantly reduced time.

Figures (12)

• Figure 1: Comparison of stochastic vs. non-stochastic value-based variants on the FrozenLake-v1, with steps on the x-axis and cumulative rewards on the y-axis.
• Figure 2: Comparison of wall time in seconds of stochastic and non-stochastic DQN methods on various action set sizes.
• Figure 3: Comparison of stochastic DQN variants against other RL algorithms on the InvertedPendulum-v4, with 1000 actions, with steps on the x-axis and average returns on the y-axis.
• Figure 4: The $\mathop{\mathrm{stoch\,max}}\limits$ and $\max$ ratio values tracked over the steps on the InvertedPendulum-v4.
• Figure 5: Comparison results for the stochastic and deterministic methods. The x-axis represents the number of possible actions, and the y-axis represents the time step duration of the agent in seconds.

Theorems & Definitions (18)

• Remark 4.1
• Lemma 5.1
• Definition 5.2
• Corollary 5.3
• Remark 6.1
• Theorem 6.2
• Remark 6.3
• Lemma 6.4
• Remark 7.1
• proof