PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning and Planning in Complex Action Spaces

Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Mohammadamin Barekatain, Simon Schmitt, David Silver

arXiv:2104.06303v1cs.LG

TL;DR

Complex action spaces can make full action enumeration infeasible, and MuZero had been applied only where its search could enumerate actions. The paper develops sample-based policy iteration and Sampled MuZero, then reports general and robust performance across discrete and continuous benchmarks. It suggests the approach may also be effective where enumerating actions is infeasible.

  • Problem

    Many real-world tasks have high-dimensional, continuous, or structured action spaces, while MuZero had been applied only where its search could enumerate the actions.

    Source: §1 · p. 1

  • Method

    The paper formulates policy improvement and evaluation over sampled action subsets and applies the framework in Sampled MuZero, which plans over sampled actions.

    Source: §1 · pp. 1–2, §5 · p. 5

  • Results

    The method succeeds across a wide variety of discrete and continuous benchmark environments and scales gracefully to small sample counts.

    Source: §7 · p. 9

  • Takeaways & Limitations

    The results suggest sample-based policy iteration may be effective in larger-scale applications where action-space enumeration is infeasible.

    Source: §7 · p. 9

    The framework assumes actions are sampled from a proposal distribution and that a process is available to build a locally improved policy.

    Source: §3 · p. 3

Abstract

from arXiv · show

Many important real-world problems have action spaces that are high-dimensional, continuous or both, making full enumeration of all possible actions infeasible. Instead, only small subsets of actions can be sampled for the purpose of policy evaluation and improvement. In this paper, we propose a general framework to reason in a principled way about policy evaluation and improvement over such sampled action subsets. This sample-based policy iteration framework can in principle be applied to any reinforcement learning algorithm based upon policy iteration. Concretely, we propose Sampled MuZero, an extension of the MuZero algorithm that is able to learn in domains with arbitrarily complex action spaces by planning over sampled actions. We demonstrate this approach on the classical board game of Go and on two continuous control benchmark domains: DeepMind Control Suite and Real-World RL Suite.

1. Introduction

Complex action spaces make exhaustive action enumeration infeasible, while MuZero had been applied only where its search could enumerate the action space. The paper develops sample-based policy iteration and applies it in Sampled MuZero across continuous-control benchmarks and Go.

1. Introduction

  • Real-world tasks include continuous, high-dimensional, and richly structured action spaces that can make full enumeration infeasible.Examples include robotic joint angles, smooth vehicle acceleration, sentences, queries, images, and serialized objects.Source: §1 · p. 1
  • MuZero had been applied only to domains with action spaces small enough for its tree search to enumerate fully.This limits its direct application to many continuous or high-dimensional problems.Source: §1 · p. 1
  • The framework reasons about policy improvement and evaluation using small sampled action subsets, rather than requiring enumeration of the full action space.It uses local information to train a global policy, act, and perform explicit policy evaluation for planning and local policy iteration.Source: §1 · p. 1
  • Sampled MuZero extends MuZero to complex action spaces by planning over sampled actions.The framework is presented as applicable in principle to domains where MuZero can be used, including domains with arbitrarily complex action spaces.Source: §1 · pp. 1–2
  • The approach is demonstrated on DeepMind Control Suite, Real-World RL Suite, and Go, where high performance is maintained while sampling only a small fraction of possible moves.The control suites demonstrate continuous-action applications, while Go demonstrates sampling in a large discrete action space.Source: §1 · p. 2

2. Related Work

Prior work addresses complex action spaces through model-free methods, limited continuous-action extensions, and planning approaches. This paper distinguishes its focus on sampling actions within MCTS-based policy improvement from work using other sampling or planning strategies.

2. Related Work

  • Model-free continuous-control methods include DDPG and D4PG, while TRPO, PPO, SAC, MPO, and AWR use other policy-optimization approaches.MPO includes a sample-based policy-improvement step related to this work.Source: §2 · p. 2
  • The closest cited approach also uses sample-based policy updates, but uses a KL-regularized improvement objective rather than MCTS-based improvement.This comparison is with Springenberg et al. (2020).Source: §2 · p. 2
  • Sparse sampling plans over sampled state transitions but typically still enumerates every action, whereas this work samples actions.The paper focuses on large action spaces and uses deterministic transition models rather than combining both sampling approaches.Source: §2 · p. 2
  • Earlier AlphaZero and MuZero extensions to continuous actions were limited to low-dimensional cases, with reported results on one- or two-dimensional tasks.A0C demonstrated learning on 1D Pendulum; another extension reported results on 1D and 2D tasks.Source: §2 · p. 2
  • Factorized categorical policies represent each action dimension separately to avoid the exponential action-count growth of simple discretization.The cited approach reports good results across a variety of domains.Source: §2 · p. 2

3. Background

The background frames reinforcement learning as policy evaluation followed by policy improvement, and interprets MuZero’s search as local policy iteration coupled to an outer learning process. Sampling raises questions about how locally improved policies are built, used, and evaluated.

3. Background

  • Policy iteration repeatedly evaluates a policy and improves it, with the goal of learning a policy that maximizes expected return.The setup models an agent acting over time in an MDP with states, actions, transitions, and rewards.Source: §3 · p. 3
  • Planning methods such as MCTS perform local policy-improvement steps followed by explicit evaluation of the improved policy.They use model-based action or trajectory evaluations to construct locally better policies.Source: §3 · p. 3
  • MuZero couples inner MCTS policy iteration with an outer process that learns its model, reward, value function, and policy.The search supplies an improved policy for acting and for training the outer policy; tree values use bootstrapped n-step returns.Source: §3 · p. 3
  • With only a small action subset available, the framework assumes actions are sampled from a proposal distribution and a process exists to build a locally improved policy.The questions include selecting samples, forming local improvements, using them to train or act with a global policy, and evaluating them for planning.Source: §3 · p. 3

4. Sample-based Policy Iteration

The section formulates policy improvement and evaluation over sampled actions as estimates of expectations under the fully improved policy. It then applies this framework to improve, train, act with, and evaluate policies using samples.

4. Sample-based Policy Iteration

  • The framework treats sampled actions as local information for estimating expectations under the fully improved policy, enabling principled policy improvement and evaluation.This addresses the gap between knowing an improved policy on sampled actions and reasoning about it across the full action space.Source: §4 · pp. 3–4

4.1. Operator view of Policy Improvement

  • Policy improvement is described as an improvement operator followed by projection into the space of realizable policies.Examples express policy-gradient, PPO, MPO, and AWR improvements as operators based on action values or advantages.Source: §4.1 · p. 4

4.2. Action-Independent Policy Improvement Operator

  • The paper states that the policy-improvement operators described in the preceding discussion are action-independent and gives MPO as an example.MPO’s operator uses the policy, action-value function, and a temperature parameter, normalized by a state-dependent factor.Source: §4.2 · p. 4

4.3. Sample-Based Action-Independent Policy Improvement Operator

  • The sample-based operator uses an empirical distribution over sampled actions and an empirical-to-proposal probability ratio; with MPO and β = π, it weights actions by exp(Q(s, a)/τ).The ratio corrects for the difference between the sampled-action distribution and the proposal distribution.Source: §4.3 · p. 4

4.4. Computing an expectation with respect to Iπ

  • As the sample count K increases, estimates of expectations under the improved policy become approximately normal around the true expectation, and the sample-based operator converges in distribution to the true operator.Figure 1 depicts sampled actions, the empirical distribution, and convergence toward the improved policy.Source: §4.4 · p. 4, Figure 1 · p. 5

4.5. Sample-based Policy Evaluation and Improvement

  • The estimated improved-policy expectation trains the global policy by cross-entropy, supports acting through resampling, and enables 1-step or n-step value backpropagation for planning.The value update can use X = r + γV′ to estimate the improved policy’s value in a search tree.Source: §4.5 · p. 5

5. Sampled MuZero

Sampled MuZero adapts MuZero’s search and learning loop to sampled actions, modifying its PUCT prior to account for sampling. Root search visits provide the improved behavior policy, while value learning follows MuZero’s n-step-return approach.

5. Sampled MuZero

  • Sampled MuZero applies MuZero’s sampling procedure to plan and learn in domains with arbitrarily complex action spaces.Source: §5 · p. 5

5.1. Inner Policy Evaluation and Improvement

  • In MCTS, retain n-step value averaging but replace PUCT’s π with adjusted sampled prior ˆπβ; this approximates sample-based policy improvement and avoids the naive correction’s potential instability.Source: §5.1 · pp. 5–6

5.2. Outer Policy Improvement

  • Normalized root visit counts form the sample-based improved policy and already account for how root actions were sampled; MuZero’s projection minimizes KL divergence to the representable policy.Source: §5.2 · p. 6

5.3. Outer Policy Evaluation

  • The agent samples behavior from the improved policy and trains its value function on n-step returns from trajectories generated by that policy.Source: §5.3 · p. 6

5.4. Search Tree Node Expansion

  • At node expansion, Sampled MuZero draws K ≪ N actions from β and returns each action with its policy and sampling probabilities.When simulations greatly exceed K, progressive widening could dynamically add actions along highly visited paths.Source: §5.4 · p. 6

5.5. Sampling distribution β

  • Sampling uses β = π, optionally temperature-modulated, with root noise in both β and π so low-prior moves can still be sampled and searched.Source: §5.5 · p. 6

6. Experiments

The experiments test whether Sampled MuZero works across varied action spaces and retains performance when sampling only a fraction of available actions. Results span Go and Atari, continuous-control tasks using state and pixel inputs, 56-dimensional humanoid control, and the Real-World RL Challenge.

6. Experiments

  • The experiments assess generality across discrete and continuous environments and robustness when only a small fraction of the action set is sampled.The comparison target is performance from algorithms that can consider the entire action set.Source: §6 · p. 6

6.1. Go

  • In 19x19 Go, K = 50 closely approached the full-action MuZero baseline despite sampling only a small part of the 362-action space.The experiments tested K ∈{15, 25, 50, 100}; performance improved as K increased.Source: §6.1 · p. 7, Figure 2 · p. 7

6.2. Atari

  • In Ms. Pacman, K = 3 samples rapidly approached the all-actions MuZero baseline, while K = 2 was insufficient for efficient learning.The Atari action space contains 18 actions.Source: §6.2 · p. 7, Figure 3 · p. 7

6.3. DeepMind Control Suite

  • Sampled MuZero performed well across DeepMind Control Suite tasks, with especially good results in hard and manipulator categories.It also learned from raw pixels, including control of a 21-dimensional humanoid; on 56-dimensional CMU humanoid tasks it outperformed prior results using more than an order of magnitude fewer environment interactions.Source: §6.3 · pp. 7, 9, Figure 4 · p. 8, Figure 5 · p. 8

6.4. Real-World RL Challenge Benchmark

  • Sampled MuZero significantly outperformed baseline algorithms at all three Real-World RL Challenge difficulties.The benchmark tests robustness to complications including delays, partial observability, and stochasticity.Source: §6.4 · p. 9

7. Conclusions

The paper presents a sample-based framework for learning and planning in complex action spaces, instantiated as Sampled MuZero. Its experiments show success across discrete and continuous benchmarks, with robust performance at small sample counts.

7. Conclusions

  • The unified framework accounts for sampled actions during policy improvement and evaluation, and the paper instantiates it in the model-based planning algorithm MuZero.The authors propose that the strategy could apply to other reinforcement-learning algorithms with action-independent policy improvement steps.Source: §7 · p. 9
  • Sampled MuZero succeeded across diverse discrete and continuous benchmarks and scaled gracefully to small sample counts; on RWRL, performance was strong across difficulties, especially on high-dimensional Humanoid tasks.The RWRL table covers easy, medium, and hard variations.Source: §7 · p. 9, Table 1 · p. 9

A.1. Gaussian policy parameterisation

  • Gaussian policy prediction performed similarly across almost all tasks but was harder to optimise than categorical prediction; entropy regularisation with coefficient 5e-3 helped.The categorical distribution required no added regularisation.Source: §A.1 · p. 12

A. DeepMind Control Suite and Real-World RL Experiments

  • DeepMind Control Suite curves compare Sampled MuZero with DMPO and D4PG, while Real-World RL curves cover easy, medium, and hard task variations.The plotted metric is mean episode return over millions of environment frames.Source: Figure 7 · p. 13, Figure 8 · p. 14

A.2. Sampled MuZero from Pixels

  • Sampled MuZero equalled or surpassed Dreamer on all 20 tasks, without action repeat, observation reconstruction, or hyperparameter retuning.The comparison used the 20-task, 5-million-environment-step setup from Hafner et al. (2019).Source: Table 2 · p. 15, §A.2 · p. 15

A.3. Ablation on the number of samples

  • On 21-dimensional humanoid.run, K = 3 samples sufficed for a good policy, and performance did not appear to improve beyond K = 10.The experiment compared K values of 3, 5, 10, 20, and 40 across six seeds.Source: §A.3 · p. 15

A.4. Reproducibility

The reproducibility experiments plot individual runs for state and raw-pixel inputs and for different action-sample counts. Small sample counts retained reproducible performance across runs.

A.4. Reproducibility

  • The hard and manipulator task plots show individual runs for state and raw-pixel inputs across three seeds.Figure 11 presents performance throughout training.Source: §A.4 · p. 15
  • The humanoid.run learning curves track performance throughout training across different numbers of action samples.Figure 10 reports six seeds per experiment.Source: Figure 10 · p. 15
  • Across six seeds on humanoid.run, performance remained reproducible even with small sample counts.The runs used K = 3, 5, 10, 20, or 40 action samples.Source: §A.4 · p. 15

A.5. Ablation on using ˆπβ vs π

The ablation compares using ˆπβ with π in Sampled MuZero’s PUCT formula at two sample counts, with and without root Q-value evaluation. The results favor ˆπβ, especially with few samples and no root Q evaluations.

A.5. Ablation on using ˆπβ vs π

  • The experiment tests the practical importance of using ˆπβ rather than π in the PUCT formula on humanoid.run.The authors expect the difference between the two forms to diminish as K increases.Source: §A.5 · p. 15
  • Using ˆπβ was much better than using π when the sample count was small and root Q values were not evaluated.The comparison included K = 5 and K = 20.Source: §A.5 · p. 15
  • Evaluating sampled-action Q values at the root attenuated the performance drop from using π, but ˆπβ still performed better.Root Q evaluation puts more emphasis on values and less on the prior in the PUCT formula.Source: §A.5 · p. 15

B. Go Experiments

The Go experiments adapt MuZero’s value training and search budget, then evaluate MuZero and Sampled MuZero checkpoints in matches. The supplied passages describe these settings but do not report comparative Go outcomes.

B. Go Experiments

  • For Go, value training used n-step bootstrapping with n = 25 and averaged target-network predictions from four consecutive states to account for alternating player perspectives.The authors observed reduced value overfitting and were able to train while generating less data.Source: B. Go Experiments · pp. 15–16
  • The Go experiments used a search budget of 400 simulations per move instead of 800 to reduce computation.Source: B. Go Experiments · p. 16
  • Checkpoint evaluations used 100 matches with 800 simulations per move, anchoring the Elo scale to a final MuZero baseline of 2000 Elo.Source: B. Go Experiments · p. 16

C. Atari Experiments

Atari checkpoints for MuZero and Sampled MuZero were evaluated during training, while the search section specifies the PUCT formula and experimental constants.

C. Atari Experiments

  • MuZero and Sampled MuZero checkpoints were evaluated during training by playing 100 games with a 50-simulation search budget per move.The evaluation used the same architecture, optimization, and hyperparameters as MuZero.Source: C. Atari Experiments · pp. 16–17

D. Search

  • Sampled MuZero used the full PUCT search formula, with c1 = 1.25 and c2 = 19652 in the experiments.When N(s) ≪ c2, the log in the exploration term is approximately zero, yielding the simplified form shown.Source: D. Search · p. 17
  • The displayed PUCT selection rule combines Q(s, a) with an exploration term when choosing an action.The supplied formula passage is incomplete, so the exploration term’s full expression is not specified here.Source: D. Search · p. 17

E. Sample-based Policy Improvement and Evaluation Proofs

The proofs characterize convergence and asymptotic normality for sample-based policy improvement and evaluation. They also identify an action-independent sample-based policy improvement operator associated with Iπ.

E. Sample-based Policy Improvement and Evaluation Proofs

F. The MuZero Policy Improvement Operator

  • The analysis identifies Îπβ as the action-independent sample-based policy improvement operator associated with Iπ.The visit-count distribution is obtained by searching with prior π̂β = β̂/β π, and Ẑβ normalizes Îπβ across actions.Source: F. The MuZero Policy Improvement… · p. 18

Related papers

Loading 2104.06303v1…