[Paper Review] Q-Learning in enormous action spaces via amortized approximate maximization

Tom Van de Wiele, David Warde-Farley+2 more|arXiv (Cornell University)|Jan 22, 2020

Reinforcement Learning in Robotics39 references30 citations

TL;DR

Introduces Amortized Q-learning (AQL), which replaces exact action maximization in Q-learning with a learned proposal distribution to sample a small set of actions, enabling scalable learning in discrete, continuous, and hybrid action spaces.

ABSTRACT

Applying Q-learning to high-dimensional or continuous action spaces can be difficult due to the required maximization over the set of possible actions. Motivated by techniques from amortized inference, we replace the expensive maximization over all actions with a maximization over a small subset of possible actions sampled from a learned proposal distribution. The resulting approach, which we dub Amortized Q-learning (AQL), is able to handle discrete, continuous, or hybrid action spaces while maintaining the benefits of Q-learning. Our experiments on continuous control tasks with up to 21 dimensional actions show that AQL outperforms D3PG (Barth-Maron et al, 2018) and QT-Opt (Kalashnikov et al, 2018). Experiments on structured discrete action spaces demonstrate that AQL can efficiently learn good policies in spaces with thousands of discrete actions.

Motivation & Objective

• Motivate Q-learning in environments with high-dimensional or continuous action spaces where exact maximization is intractable.
• Propose an amortized approach that learns a proposal distribution to sample candidate actions for Q-learning updates.
• Demonstrate that AQL handles discrete, continuous, and hybrid action spaces while maintaining Q-learning benefits.
• Show empirical evidence that AQL outperforms strong baselines on continuous control and large discrete action tasks.

Proposed method

• Replace exact maximization over actions with maximization over a set of samples from a learned proposal distribution μ(a|s;θμ).
• Train a second network to predict μ by supervised learning from actions found by a stochastic search procedure, with regularization to maintain exploration.
• Define AQL loss for Q-function as in standard Q-learning but with max over sampled actions from μ instead of all actions.
• Update μ by supervised learning to increase the likelihood of actions with high Q-values found via a mix of uniform samples and samples from μ, plus an entropy term for exploration.
• Use autoregressive proposals to model dependencies among sub-actions, enabling handling of multi-dimensional action spaces.
• Apply AQL to continuous, discrete, and hybrid action spaces by adjusting the form of μ (discretized, Gaussian, or categorical).

Experimental results

Research questions

• RQ1 Can amortizing the action selection maximize Q-learning performance in high-dimensional action spaces?
• RQ2 How does AQL perform compared to strong baselines (D3PG, QT-Opt, IMPALA) on continuous control tasks?
• RQ3 Can AQL scale to large discrete action spaces and structured action sets?
• RQ4 What is the impact of the proposal distribution’s accuracy and exploration on learning efficiency?

Key findings

• AQL outperforms D3PG and QT-Opt on continuous control tasks in the DeepMind Control Suite.
• AQL learns effective policies in large discrete action spaces, demonstrated on DeepMind Lab tasks with thousands of actions.
• Discretized AQL and deterministic policy variants outperform Uniform Q-learning and QT-Opt in high-dimensional action spaces.
• On the DeepMind Lab large action set (3528 actions), AQL achieves best final performance and learns more efficiently than exact Q-learning with the large set due to beneficial stochasticity in action selection.
• In control-suite experiments, a continuous implementation of AQL performs comparably on low/medium dimensions but can underperform on very high-dimensional tasks, while discretized AQL remains strong overall.
• AQL provides a flexible, off-policy alternative to stochastic actor-critic methods for handling various action space types without requiring importance sampling corrections.