Source-linked AI summary
Deep Reinforcement Learning in Large Discrete Action Spaces
Gabriel Dulac-Arnold, Richard Evans, Hado van Hasselt, Peter Sunehag, Timothy Lillicrap, Jonathan Hunt, Timothy Mann, Theophane Weber, Thomas Degris, Ben Coppin
TL;DR
Large discrete action spaces make reinforcement learning difficult because existing methods do not jointly provide action generalization and sub-linear complexity. The paper introduces an action-embedding policy using approximate nearest-neighbor lookup, demonstrates it on tasks up to one million actions, and reports that subset selection and approximate lookup can provide speedups with often slight performance impact.
Problem
Current reinforcement-learning approaches do not jointly provide sub-linear complexity and generalization over large discrete action spaces, limiting their use in such tasks.
Method
The Wolpertinger architecture embeds actions in a continuous space, generates a proto-action, retrieves nearby discrete actions, and refines them with a critic using approximate nearest-neighbor search.
Results
The approach performs on tasks ranging from tens to one million discrete actions, with subset selection providing significant speedups and approximate lookup often affecting performance only slightly.
Takeaways & Limitations
The architecture enables reinforcement-learning methods to generalize over large action sets with sub-linear, logarithmic-time lookup complexity.
Takeaways & Limitations
Value-based selection becomes intractable because it evaluates every action, while actor-based approaches do not naturally generalize to previously unseen actions.
Abstract
from arXiv · showhide
Being able to reason in an environment with a large number of discrete actions is essential to bringing reinforcement learning to a larger class of problems. Recommender systems, industrial plants and language models are only some of the many real-world tasks involving large numbers of discrete actions for which current methods are difficult or even often impossible to apply. An ability to generalize over the set of actions as well as sub-linear complexity relative to the size of the set are both necessary to handle such tasks. Current approaches are not able to provide both of these, which motivates the work in this paper. Our proposed approach leverages prior information about the actions to embed them in a continuous space upon which it can generalize. Additionally, approximate nearest-neighbor methods allow for logarithmic-time lookup complexity relative to the number of actions, which is necessary for time-wise tractable training. This combined approach allows reinforcement learning methods to be applied to large-scale learning problems previously intractable with current methods. We demonstrate our algorithm's abilities on a series of tasks having up to one million actions.
1. Introduction
Large-action reinforcement-learning tasks are important but difficult for current algorithms. The paper proposes an action-embedding policy with logarithmic-time approximate nearest-neighbor lookup and evaluates it on tasks with up to one million actions.
- Recommender systems and industrial control may require reasoning over hundreds of millions or millions of actions at each step.
- Current reinforcement-learning algorithms are difficult or impossible to apply to these large-action problems.
- The proposed policy embeds actions in a continuous space so the actor can generalize while its complexity is decoupled from action-set size.
- Approximate nearest-neighbor search finds nearby discrete actions in logarithmic time, making learning and acting tractable.
- The approach is demonstrated on tasks with up to one million actions and is intended to scale beyond millions.
2. Definitions
The paper formulates reinforcement learning as an MDP with discrete states and actions, transitions, rewards, and discounting. Policies and value functions are approximated by parameterized functions to maximize expected discounted return.
- An MDP comprises discrete action and state sets, transition probabilities, a reward function, and a discount factor γ ∈ [0, 1].
- Each action is represented by an n-dimensional vector carrying action-related information, while each state is represented by an m-dimensional vector.
- An episode’s return is the discounted sum of rewards received over time.
- The reinforcement-learning objective is to learn a policy that maximizes expected return across episodes.
- The state-action value Qπ(s, a) is the expected return after taking action a in state s and following policy π thereafter.
- Both the policy and Q function are represented by parameterized approximators.
3. Problem Description
Large discrete action spaces expose a trade-off: value-based policies can generalize across actions but require linear action evaluations, while actor-based policies avoid that cost but generalize less naturally. The paper targets both sub-linear complexity and action generalization.
- Value-based policies: Value-based policies condition decisions directly on the value function, commonly choosing the action with the highest value.
- Value-based policies: Value-based action selection requires |A| evaluations, so execution complexity grows linearly with the number of actions.
- Value-based policies: Value-based policies can generalize between similar actions and use action features for previously unseen actions.
- Actor-based policies: Actor-based architectures avoid evaluating a costly Q function over every action but typically scale linearly with action count.
- Actor-based policies: Actor-based approaches do not generalize over the action space as naturally as value-based approaches and cannot extend to unseen actions.
- The target is a policy combining sub-linear action-space complexity with the ability to generalize over actions.
4. Proposed Approach
The Wolpertinger architecture generates continuous proto-actions, retrieves nearby discrete actions, and refines them with a critic. It combines action generalization with efficient action selection and can be trained using policy-gradient methods.
- Architecture: The Wolpertinger architecture uses an actor-critic framework with an efficient action-generating actor and a critic that refines candidate actions.Both actor and critic are implemented with multilayer neural networks and trained using Deep Deterministic Policy Gradient.
- Action Generation: The actor maps each state representation to a continuous proto-action that may not itself be a valid discrete action.The proto-action lies in R^n, while valid actions belong to the discrete set A.
- Action Generation: Approximate nearest-neighbor search retrieves k discrete actions closest to the proto-action in logarithmic time.The mapping uses L2 distance; k=1 reduces to a simple nearest-neighbor lookup.
- Action Refinement: The policy selects the highest-Q action among the retrieved candidates rather than simply applying the nearest action.This refinement addresses cases where nearby actions have different long-term values or the representation contains low-value outliers.
- Action Refinement: The generated candidate-set size k trades policy quality against speed, and the refinement pass improves robustness to imperfect action representations.The authors report that the second pass is essential for learning in certain domains.
- Training: The architecture is trained by following the policy gradient of the continuous actor output despite the policy’s nondifferentiable discrete mapping.The mapping effects are treated as deterministic environmental dynamics in the simpler training formulation, while the full policy generates target actions for the critic update.
5. Analysis
The analysis explains why evaluating several nearby actions improves selection quality while producing diminishing returns as k grows. It predicts and reports that small candidate subsets can approach full-set performance at lower computational cost.
- Complexity: The algorithm’s time complexity scales linearly with the number of selected actions, k.Increasing k beyond a task-specific limit may not improve performance.
- Expected Value: In the simplified model, the expected value of the best action among the k nearest candidates improves as k increases.The model includes nearby bad actions with probability p and remaining action values uniformly distributed around the proto-action’s value.
- Diminishing Returns: The contribution from bad-action avoidance decreases exponentially with k, while the remaining contribution decreases as O(1/(k+1)).Both terms produce large gains for small k and diminishing marginal returns for larger k.
- Diminishing Returns: Using 5% or 10% of the maximal action count yields performance similar to using the full action set.Using the remaining actions provides relatively small performance benefits while increasing computational time by an order of magnitude.
6. Related Work
Prior approaches scale large discrete action spaces mainly by factorizing actions into binary subspaces, while other methods use continuous representations or limited candidate evaluation. These methods have important restrictions involving action-code design, state generalization, Q-function refinement, or the number of actions evaluated.
- Factorized Action Spaces: H-value functions and ECOC-based policies factorize the action space into log(|A|) binary subspaces.The resulting methods reason about the binary subspaces independently during action selection or policy prediction.
- Factorized Action Spaces: Factorized approaches scale to very large action spaces but require a binary code representation whose design is difficult.The generalized value-function approach additionally stores a value function per state and uses a Linear Program.
- Factorized Action Spaces: The generalized value-function approach does not generalize over continuous state spaces, while the ECOC approach does not support Q-function refinement.These limitations distinguish the approaches from methods that combine action generation with value-based refinement.
- Continuous Representations: A continuous-action policy-gradient method followed by nearest-discrete-action selection was previously tested only on small, low-dimensional problems.The cited prior setup involved at most 21 discrete actions and a one-dimensional continuous action space.
- Limited Candidate Evaluation: An extension of Deep Q-Networks generates representations for environment-provided actions but evaluates only a small set of 2–4 available actions.It therefore does not explicitly select an action from a large discrete set.
- Extensions: The authors also applied this policy architecture to slate MDPs, where multiple actions are taken at each state.This application is reported as related prior use of the architecture.
7. Experiments
The experiments evaluate the agent across control, planning, and recommendation environments, varying action-space size, nearest-neighbor configuration, training time, and reward. These settings include action spaces reaching approximately one million actions and test both efficient search and partial re-ranking.
- The agent is evaluated on Discretized Continuous Control, Multi-Step Planning, and Recommender Systems.
- Discretized Continuous Control: In cart-pole, continuous control actions are discretized into i equally spaced values per dimension, producing | A | = i^d discrete actions.The environment tests performance and learning speed as the number of discrete actions changes.
- Multi-Step Planning: Multi-step planning uses all length-n sequences of two base actions, reaching 2^20 ≈ 1e6 actions when n = 20.The puddle-world task requires trading off puddle costs against distance while finding a shortest path.
- Recommender Systems: The recommendation environment models items as actions, with user acceptance governed by a transition probability matrix and item-specific rewards.The current state is the item the user is consuming, and session termination probabilities simulate user patience.
- Evaluation: The evaluation varies nearest neighbors from k = 1 to k = | A |, spanning fastest nearest-neighbor selection, partial re-ranking, and a computationally intractable greedy upper bound.Intermediate k values measure performance gains from partial re-ranking.
- Evaluation: Training time and average reward are compared for full nearest-neighbor search and three FLANN approximate-nearest-neighbor settings: Slow, Medium, and Fast.The Slow setting uses a hierarchical k-means tree with branching factor 16 and 99% retrieval accuracy on the recommender task.
8. Results
Across cart-pole, Puddle World, and recommender tasks, the approach learns with very large discrete action spaces, but performance and speed depend on neighbor count, lookup settings, representation, and exploration.
- 8.1. Cart-Pole: On one-million-action cart-pole, k = 1 finds an optimal policy, while k = 0.5% trains prohibitively slowly and fails to converge within the displayed budget.The 0.5% setting corresponds to 5,000 actions.
- 8.1. Cart-Pole: Agents with k = 1 converge after 150,000 seconds, whereas k = 5,000 trains much more slowly on cart-pole.FLANN is helpful for k = 1 lookups, while at k = 5,000 computation is dominated by Q evaluation.
- 8.2. Puddle World: Longer Puddle World plans learn faster, and a 20-step plan represents 1,048,576 actions.With k = 1 and slow FLANN, the task reaches a near-optimal policy without the arg max pass; even the most lossy FLANN setting converges to an optimal policy.
- 8.2. Puddle World: Puddle World achieves equivalent performance with a speedup of up to 1,250 times.The comparison is reported for median training steps per second.
- 8.3. Recommender Task: Uniform exploration converges on recommender tasks but typically reaches 50% of guided-exploration performance in the larger tasks.The experiments use tasks with 49, 835, and 13,138 elements.
- 8.3. Recommender Task: In recommender tasks, increasing k improves performance on the 835-element task, while lower-recall FLANN settings significantly reduce performance on the 13,138-element task.The 13,138-element task’s all-action agent trains approximately 15 times slower in wall-time than the 1% agent; compact action representations can stabilize convergence on the 49-element task.
9. Conclusion
The paper introduces a policy architecture for efficiently learning and acting in large discrete action spaces, combining action generalization with sub-linear complexity. It demonstrates performance across tasks with up to one million actions and identifies learned action representations as future work.
- The architecture uses DDPG to efficiently learn and act in large discrete action spaces.
- Action embeddings support generalization across actions with sub-linear complexity relative to the action-space size.
- Considering only a subset of actions is sufficient in many tasks and provides significant speedups.
- Approximate nearest-neighbor lookup often impacts performance only slightly.
- Future work would learn action representations during training to reposition poorly embedded actions.
Appendices
The appendix describes Wolpertinger training with DDPG: actions are executed through the full policy, stored in replay, and used for critic and actor updates. Target actions also come from the full policy.
- The full Wolpertinger policy selects the action executed in the environment, whose transition is stored in replay.
- The critic update samples replay transitions and applies a Bellman backup using target-network weights.
- The target action is generated by the full policy rather than only by the continuous actor output.
- The actor is updated using the sampled policy gradient, while the algorithm distinguishes discrete actions from prototype actions.
- Algorithm 2 initializes critic and actor networks, their target networks, the action dictionary, and a replay buffer.
- The policy gradient is evaluated at the continuous actor output while replay retains the actually executed action for critic training.
B. Proof of Lemma 1
The proof derives the distribution of the maximum value among k selected actions by multiplying their individual cumulative distribution functions. Under equal-distribution assumptions, this yields the expected maximum used in the lemma.
- The proof normalizes the setting by assuming Q(s, a) = 1/2 and b = 1/2, then reverses the affine transformation at the end.
- An action is bad with probability p and otherwise has a value uniformly distributed on [0, 1].
- The individual action-value CDF is piecewise defined across values below −c, from −c to 0, and from 0 to 1.
- For k actions, the CDF of the maximum equals the product of the individual CDFs because the maximum is below x only when every action is below x.
- The proof assumes identical distributions for all k closest actions, while noting that distance-dependent distributions could extend the result.
- Integrating the maximum's CDF gives its expected value, which is then rescaled to the arbitrary original setting.