IMPALA: Importance Weighted Actor-Learner Architectures with the notation that c−1 = ρ−1 = 1 and Qt−2 s=0 cs = 1 for t = 0 and 1. Now the coefﬁcients (αt)t≥0 are non-negative in expectation. Indeed, since ¯ρ ≥¯c, we have Eµαt = E  ρt−1 −ct−1ρt  ≥Eµ  ct−1(1 −ρt)  ≥0, since Eµρt ≤Eµ  π(at|xt) µ(at|xt)  = 1. Thus V1(x) −V2(x) is a linear combination of the values V1 −V2 at other states, weighted by non-negative coefﬁcients whose sum is X t≥0 γtEµ " t−2 Y s=0 cs  [ρt−1 −ct−1ρt] # = X t≥0 γtEµ " t−2 Y s=0 cs  ρt−1 # − X t≥0 γtEµ " t−1 Y s=0 cs  ρt # = X t≥0 γtEµ " t−2 Y s=0 cs  ρt−1 # −γ−1  X t≥0 γtEµ " t−2 Y s=0 cs  ρt−1 # −1   = γ−1 −(γ−1 −1) X t≥0 γtEµ " t−2 Y s=0 cs  ρt−1 # | {z } ≥1+γEµρ0 ≤ 1 −(1 −γ)Eµρ0 ≤ 1 −(1 −γ)β < 1.


γ−1 −(γ−1 −1) X t≥0 γtEµ " t−2 Y s=0 cs  ρt−1 # | {z } ≥1+γEµρ0 ≤ 1 −(1 −γ)Eµρ0 ≤ 1 −(1 −γ)β < 1. We deduce that ∥RV1(x) −RV2(x)∥≤η∥V1 −V2∥∞, with η = γ−1 −(γ−1 −1) P t≥0 γtEµ h Qt−2 s=0 cs  ρt−1 i ≤ 1 −(1 −γ)β < 1, so R is a contraction mapping. Thus R possesses a unique ﬁxed point. Let us now prove that this ﬁxed point is V π¯ ρ. We have: Eµ  ρt rt + γV π¯ ρ(xt+1) −V π¯ ρ(xt)  xt i = X a µ(a|xt) min ¯ρ, π(a|xt) µ(a|xt) h r(xt, a) + γ X y p(y|xt, a)V π¯ ρ(y) −V π¯ ρ(xt) i = X a π¯ρ(a|xt) h r(xt, a) + γ X y p(y|xt, a)V π¯ ρ(y) −V π¯ ρ(xt) i | {z } =0 X b min ¯ρµ(b|xt), π(b|xt)  = 0, since this is the Bellman equation for V π¯ ρ. We deduce that RV π¯ ρ = V π¯ ρ, thus V π¯ ρ is the unique ﬁxed point of R. A.2. Online learning Theorem 2. Assume a tabular representation, i.e.


s the unique ﬁxed point of R. A.2. Online learning Theorem 2. Assume a tabular representation, i.e. the state and action spaces are ﬁnite. Consider a set of trajectories, with the kth trajectory x0, a0, r0, x1, a1, r1, . . . generated by following µ: at ∼µ(·|xt). For each state xs along this trajectory, update Vk+1(xs) = Vk(xs) + αk(xs) X t≥s γt−scs . . . ct−1  ρt rt + γVk(xt+1) −Vk(xt)  , (7) with ci = min ¯c, π(ai|xi) µ(ai|xi)  , ρi = min ¯ρ, π(ai|xi) µ(ai|xi)  , ¯ρ ≥¯c. Assume that (1) all states are visited inﬁnitely often, and (2) the stepsizes obey the usual Robbins-Munro conditions: for each state x, P k αk(x) = ∞, P k α2 k(x) < ∞. Then Vk →V π¯ ρ almost surely. The proof is a straightforward application of the convergence result for stochastic approximation algorithms to th


straightforward application of the convergence result for stochastic approximation algorithms to the ﬁxed point of a contraction operator, see e.g. Dayan & Sejnowski (1994); Bertsekas & Tsitsiklis (1996); Kushner & Yin (2003).


IMPALA: Importance Weighted Actor-Learner Architectures A.3. On the choice of qs in policy gradient The policy gradient update rule (4) makes use of the coefﬁcient qs = rs +γvs+1 as an estimate of Qπ¯ ρ(xs, as) built from the V-trace estimate vs+1 at the next state xs+1. The reason why we use qs instead of vs as target for our Q-value Qπ¯ ρ(xs, as) is to make sure our estimate of the Q-value is as unbiased as possible, and the ﬁrst requirement is that it is entirely unbiased in the case of perfect representation of the V-values. Indeed, assuming our value function is correctly estimated at all states, i.e. V = V π¯ ρ, then we have E[qs|xs, as] = Qπ¯ ρ(xs, as) (whereas we do not have this property for vt). Indeed, E[qs|xs, as] = rs + γE  V π¯ ρ(xs+1) + δs+1V π¯ ρ + γcs+1δs+2V π¯ ρ + . . .


perty for vt). Indeed, E[qs|xs, as] = rs + γE  V π¯ ρ(xs+1) + δs+1V π¯ ρ + γcs+1δs+2V π¯ ρ + . . .  = rs + γE  V π¯ ρ(xs+1)  = Qπ¯ ρ(xs, as) whereas E[vs|xs, as] = V π¯ ρ(xs) + ρs rs + γE  V π¯ ρ(xs+1)  −V π¯ ρ(xs)  + γcsδs+1V π¯ ρ + . . . = V π¯ ρ(xs) + ρs rs + γE  V π¯ ρ(xs+1)  −V π¯ ρ(xs)  = V π¯ ρ(xs)(1 −ρs) + ρsQπ¯ ρ(xs, as), which is different from Qπ¯ ρ(xs, as) when V π¯ ρ(xs) ̸= Qπ¯ ρ(xs, as).


IMPALA: Importance Weighted Actor-Learner Architectures B. Reference Scores Task t Human h Random r Experts IMPALA rooms collect good objects test 10.0 0.1 9.0 5.8 rooms exploit deferred effects test 85.7 8.5 15.6 11.0 rooms select nonmatching object 65.9 0.3 7.3 26.1 rooms watermaze 54.0 4.1 26.9 31.1 rooms keys doors puzzle 53.8 4.1 28.0 24.3 language select described object 389.5 -0.1 324.6 593.1 language select located object 280.7 1.9 189.0 301.7 language execute random task 254.1 -5.9 -49.9 66.8 language answer quantitative question 184.5 -0.3 219.4 264.0 lasertag one opponent large 12.7 -0.2 -0.2 0.3 lasertag three oponents large 18.6 -0.2 -0.1 4.1 lasertag one opponent small 18.6 -0.1 -0.1 2.5 lasertag three opponents small 31.5 -0.1 19.1 11.3 natlab fixed large map 36.9 2.2 34.7 1


1 -0.1 2.5 lasertag three opponents small 31.5 -0.1 19.1 11.3 natlab fixed large map 36.9 2.2 34.7 12.2 natlab varying map regrowth 24.4 3.0 20.7 15.9 natlab varying map randomized 42.4 7.3 36.1 29.0 skymaze irreversible path hard 100.0 0.1 13.6 30.0 skymaze irreversible path varied 100.0 14.4 45.1 53.6 pyschlab arbitrary visuomotor mapping 58.8 0.2 16.4 14.3 pyschlab continuous recognition 58.3 0.2 29.9 29.9 pyschlab sequential comparison 39.5 0.1 0.0 0.0 pyschlab visual search 78.5 0.1 0.0 0.0 explore object locations small 74.5 3.6 57.8 62.6 explore object locations large 65.7 4.7 37.0 51.1 explore obstructed goals small 206.0 6.8 135.2 188.8 explore obstructed goals large 119.5 2.6 39.5 71.0 explore goal locations small 267.5 7.7 209.4 252.5 explore goal locations large 194.5 3.1 83.1


71.0 explore goal locations small 267.5 7.7 209.4 252.5 explore goal locations large 194.5 3.1 83.1 125.3 explore object rewards few 77.7 2.1 39.8 43.2 explore object rewards many 106.7 2.4 58.7 62.6 Mean Capped Normalised Score: (P t min [1, (st −rt)/(ht −rt)]) /N 100% 0% 44.5% 49.4% Table B.1. DMLab-30 test scores.


IMPALA: Importance Weighted Actor-Learner Architectures B.1. Final training scores on DMLab-30 0 20 40 60 80 100 120 140 160 Human Normalised Score pyschlab_sequential_comparison pyschlab_visual_search lasertag_one_opponent_large lasertag_one_opponent_small language_execute_random_task lasertag_three_oponents_large rooms_exploit_deferred_effects_train pyschlab_arbitrary_visuomotor_mapping skymaze_irreversible_path_hard natlab_fixed_large_map rooms_keys_doors_puzzle skymaze_irreversible_path_varied lasertag_three_opponents_small pyschlab_continuous_recognition explore_object_rewards_few rooms_select_nonmatching_object rooms_watermaze explore_object_rewards_many explore_obstructed_goals_large explore_goal_locations_large natlab_varying_map_regrowth natlab_varying_map_randomized explore_objec


explore_goal_locations_large natlab_varying_map_regrowth natlab_varying_map_randomized explore_object_locations_large explore_object_locations_small explore_obstructed_goals_small rooms_collect_good_objects_train explore_goal_locations_small language_select_located_object language_answer_quantitative_question language_select_described_object A3C, deep IMPALA-Experts, deep IMPALA, deep, PBT Figure B.1. Human normalised scores across all DMLab-30 tasks.


IMPALA: Importance Weighted Actor-Learner Architectures C. Atari Scores ACKTR The Reactor IMPALA (deep, multi-task) IMPALA (shallow) IMPALA (deep) alien 3197.10 6482.10 2344.60 1536.05 15962.10 amidar 1059.40 833 136.82 497.62 1554.79 assault 10777.70 11013.50 2116.32 12086.86 19148.47 asterix 31583.00 36238.50 2609.00 29692.50 300732.00 asteroids 34171.60 2780.40 2011.05 3508.10 108590.05 atlantis 3433182.00 308258 460430.50 773355.50 849967.50 bank heist 1289.70 988.70 55.15 1200.35 1223.15 battle zone 8910.00 61220 7705.00 13015.00 20885.00 beam rider 13581.40 8566.50 698.36 8219.92 32463.47 berzerk 927.20 1641.40 647.80 888.30 1852.70 bowling 24.30 75.40 31.06 35.73 59.92 boxing 1.45 99.40 96.63 96.30 99.96 breakout 735.70 518.40 35.67 640.43 787.34 centipede 7125.28 3402.80 4916.84 55


40 96.63 96.30 99.96 breakout 735.70 518.40 35.67 640.43 787.34 centipede 7125.28 3402.80 4916.84 5528.13 11049.75 chopper command N/A 37568 5036.00 5012.00 28255.00 crazy climber 150444.00 194347 115384.00 136211.50 136950.00 defender N/A 113128 16667.50 58718.25 185203.00 demon attack 274176.70 100189 10095.20 107264.73 132826.98 double dunk -0.54 11.40 -1.92 -0.35 -0.33 enduro 0.00 2230.10 971.28 0.00 0.00 ﬁshing derby 33.73 23.20 35.27 32.08 44.85 freeway 0.00 31.40 21.41 0.00 0.00 frostbite N/A 8042.10 2744.15 269.65 317.75 gopher 47730.80 69135.10 913.50 1002.40 66782.30 gravitar N/A 1073.80 282.50 211.50 359.50 hero N/A 35542.20 18818.90 33853.15 33730.55 ice hockey -4.20 3.40 -13.55 -5.25 3.48 jamesbond 490.00 7869.20 284.00 440.00 601.50 kangaroo 3150.00 10484.50 8240.50 47.00 163


-5.25 3.48 jamesbond 490.00 7869.20 284.00 440.00 601.50 kangaroo 3150.00 10484.50 8240.50 47.00 1632.00 krull 9686.90 9930.80 10807.80 9247.60 8147.40 kung fu master 34954.00 59799.50 41905.00 42259.00 43375.50 montezuma revenge N/A 2643.50 0.00 0.00 0.00 ms pacman N/A 2724.30 3415.05 6501.71 7342.32 name this game N/A 9907.20 5719.30 6049.55 21537.20 phoenix 133433.70 40092.20 7486.50 33068.15 210996.45 pitfall -1.10 -3.50 -1.22 -11.14 -1.66 pong 20.90 20.70 8.58 20.40 20.98 private eye N/A 15177.10 0.00 92.42 98.50 qbert 23151.50 22956.50 10717.38 18901.25 351200.12 riverraid 17762.80 16608.30 2850.15 17401.90 29608.05 road runner 53446.00 71168 24435.50 37505.00 57121.00 robotank 16.50 68.50 9.94 2.30 12.96 seaquest 1776.00 8425.80 844.60 1716.90 1753.20 skiing N/A -10753.40 -8988.00 -


8.50 9.94 2.30 12.96 seaquest 1776.00 8425.80 844.60 1716.90 1753.20 skiing N/A -10753.40 -8988.00 -29975.00 -10180.38 solaris 2368.60 2760 1160.40 2368.40 2365.00 space invaders 19723.00 2448.60 199.65 1726.28 43595.78 star gunner 82920.00 70038 1855.50 69139.00 200625.00 surround N/A 6.70 -8.51 -8.13 7.56 tennis N/A 23.30 -8.12 -1.89 0.55 time pilot 22286.00 19401 3747.50 6617.50 48481.50 tutankham 314.30 272.60 105.22 267.82 292.11 up n down 436665.80 64354.20 82155.30 273058.10 332546.75 venture N/A 1597.50 1.00 0.00 0.00 video pinball 100496.60 469366 20125.14 228642.52 572898.27 wizard of wor 702.00 13170.50 2106.00 4203.00 9157.50 yars revenge 125169.00 102760 14739.41 80530.13 84231.14 zaxxon 17448.00 25215.50 6497.00 1148.50 32935.50 Table C.1.


9.00 102760 14739.41 80530.13 84231.14 zaxxon 17448.00 25215.50 6497.00 1148.50 32935.50 Table C.1. Atari scores after 200M steps environment steps of training. Up to 30 no-ops at the beginning of each episode.


IMPALA: Importance Weighted Actor-Learner Architectures D. Parameters In this section, the speciﬁc parameter settings that are used throughout our experiments are given in detail. Hyperparameter Range Distribution Entropy regularisation [5e-5, 1e-2] Log uniform Learning rate [5e-6, 5e-3] Log uniform RMSProp epsilon (ε) regularisation parameter [1e-1, 1e-3, 1e-5, 1e-7] Categorical Table D.1. The ranges used in sampling hyperparameters across all experiments that used a sweep and for the initial hyperparameters for PBT. Sweep size and population size are 24. Note, the loss is summed across the batch and time dimensions. Action Native DeepMind Lab Action Forward [ 0, 0, 0, 1, 0, 0, 0] Backward [ 0, 0, 0, -1, 0, 0, 0] Strafe Left [ 0, 0, -1, 0, 0, 0, 0] Strafe Right [ 0, 0, 1, 0, 0, 0, 0] Look


[ 0, 0, 0, -1, 0, 0, 0] Strafe Left [ 0, 0, -1, 0, 0, 0, 0] Strafe Right [ 0, 0, 1, 0, 0, 0, 0] Look Left [-20, 0, 0, 0, 0, 0, 0] Look Right [ 20, 0, 0, 0, 0, 0, 0] Forward + Look Left [-20, 0, 0, 1, 0, 0, 0] Forward + Look Right [ 20, 0, 0, 1, 0, 0, 0] Fire [ 0, 0, 0, 0, 1, 0, 0] Table D.2. Action set used in all tasks from the DeepMind Lab environment, including the DMLab-30 experiments. D.1. Fixed Model Hyperparameters In this section, we list all the hyperparameters that were kept ﬁxed across all experiments in the paper which are mostly concerned with observations speciﬁcations and optimisation. We ﬁrst show below the reward pre-processing function that is used across all experiments using DeepMind Lab, followed by all ﬁxed numerical values.


tion that is used across all experiments using DeepMind Lab, followed by all ﬁxed numerical values. −10 −5 0 5 10 Reward −1 0 1 2 3 4 5 Clipped Reward Figure D.1. Optimistic Asymmetric Clipping - 0.3 · min(tanh(reward), 0) + 5.0 · max(tanh(reward), 0)


IMPALA: Importance Weighted Actor-Learner Architectures Parameter Value Image Width 96 Image Height 72 Action Repetitions 4 Unroll Length (n) 100 Reward Clipping - Single tasks [-1, 1] - DMLab-30, including experts See Figure D.1 Discount (γ) 0.99 Baseline loss scaling 0.5 RMSProp momentum 0.0 Experience Replay (in Section 5.2.2 ) - Capacity 10,000 trajectories - Sampling Uniform - Removal First-in-ﬁrst-out Table D.3. Fixed model hyperparameters across all DeepMind Lab experiments.


IMPALA: Importance Weighted Actor-Learner Architectures E. V-trace Analysis E.1. Controlled Updates Here we show how different algorithms (On-Policy, No-correction, ε-correction, V-trace) behave under varying levels of policy-lag between the actors and the learner. 0.2 0.4 0.6 0.8 Environment Frames 1e9 0 10 20 30 40 50 60 Return 0 1 10 100 500 ϵ-correction 0.2 0.4 0.6 0.8 Environment Frames 1e9 0 1 10 100 500 No-correction 0.2 0.4 0.6 0.8 Environment Frames 1e9 0 1 10 100 500 V-trace rooms_watermaze 0.2 0.4 0.6 0.8 Environment Frames 1e9 0 5 10 15 20 25 30 Return 0 1 10 100 500 ϵ-correction 0.2 0.4 0.6 0.8 Environment Frames 1e9 0 1 10 100 500 No-correction 0.2 0.4 0.6 0.8 Environment Frames 1e9 1 0 10 100 500 V-trace rooms_keys_doors_puzzle 0.2 0.4 0.6 0.8 Environment Frames 1e9 5 0 5 10


s 1e9 1 0 10 100 500 V-trace rooms_keys_doors_puzzle 0.2 0.4 0.6 0.8 Environment Frames 1e9 5 0 5 10 15 20 25 30 35 Return 0 1 10 100 500 ϵ-correction 0.2 0.4 0.6 0.8 Environment Frames 1e9 0 1 10 100 500 No-correction 0.2 0.4 0.6 0.8 Environment Frames 1e9 10 1 100 0 500 V-trace lasertag_three_opponents_small 0.2 0.4 0.6 0.8 Environment Frames 1e9 0 50 100 150 200 250 Return 0 1 100 10 500 ϵ-correction 0.2 0.4 0.6 0.8 Environment Frames 1e9 0 1 10 100 500 No-correction 0.2 0.4 0.6 0.8 Environment Frames 1e9 10 1 0 100 500 V-trace explore_goal_locations_small 0.2 0.4 0.6 0.8 Environment Frames 1e9 0 5 10 15 20 25 30 35 40 45 Return 0 1 10 100 500 ϵ-correction 0.2 0.4 0.6 0.8 Environment Frames 1e9 0 1 10 100 500 No-correction 0.2 0.4 0.6 0.8 Environment Frames 1e9 0 1 10 100 500 V-trace se


es 1e9 0 1 10 100 500 No-correction 0.2 0.4 0.6 0.8 Environment Frames 1e9 0 1 10 100 500 V-trace seekavoid_arena_01 Figure E.1. As the policy-lag (the number of update steps the actor policy is behind learner policy) increases, learning with V-trace is more robust compared to ε-correction and pure on-policy learning.


IMPALA: Importance Weighted Actor-Learner Architectures E.2. V-trace Stability Analysis 1 5 9 13 17 21 24 Hyperparameter Combination 0 10 20 30 40 50 60 Final Return rooms_watermaze 1 5 9 13 17 21 24 Hyperparameter Combination 0 5 10 15 20 25 30 35 40 rooms_keys_doors_puzzle 1 5 9 13 17 21 24 Hyperparameter Combination −5 0 5 10 15 20 25 30 35 lasertag_three_opponents_small 1 5 9 13 17 21 24 Hyperparameter Combination 0 50 100 150 200 250 300 explore_goal_locations_small 1 5 9 13 17 21 24 Hyperparameter Combination 0 10 20 30 40 50 seekavoid_arena_01 V−trace − min(ρ, 1) 1 Step Importance Sampling − min(ρ, 1) ε−correction No-correction Figure E.2. Stability across hyper parameter combinations for different off-policy correction variants using replay.


lity across hyper parameter combinations for different off-policy correction variants using replay. V-trace is much more stable across a wide range of parameter combinations compared to ε-correction and pure on-policy learning. E.3. Estimating the State Action Value for Policy Gradient We investigated different ways of estimating the state action value function used to estimate advantages for the policy gradient calculation. The variant presented in the main section of the paper uses the V-trace corrected value function vs+1 to estimate qs = rs + γvs+1. Another possibility is to use the actor-critic baseline V (xs+1) to estimate qs = rs + γV (xs+1). Note that the latter variant does not use any information from the current policy rollout to estimate the policy gradient and relies on an acc


any information from the current policy rollout to estimate the policy gradient and relies on an accurate estimate of the value function. We found the latter variant to perform worse both when comparing the top 3 runs and an average over all runs of the hyperparameter sweep as can be see in ﬁgures E.3 and E.4. 0.0 0.2 0.4 0.6 0.8 1.0 Environment Frames 1e9 5 10 15 20 25 30 35 40 45 50 Return rooms_watermaze 0.0 0.2 0.4 0.6 0.8 1.0 Environment Frames 1e9 5 10 15 20 25 30 rooms_keys_doors_puzzle 0.0 0.2 0.4 0.6 0.8 1.0 Environment Frames 1e9 −5 0 5 10 15 20 25 30 35 lasertag_three_opponents_small 0.0 0.2 0.4 0.6 0.8 1.0 Environment Frames 1e9 0 50 100 150 200 250 explore_goal_locations_small 0.0 0.2 0.4 0.6 0.8 1.0 Environment Frames 1e9 5 10 15 20 25 30 35 40 45 seekavoid_arena_01 qs =rs +γ


.0 0.2 0.4 0.6 0.8 1.0 Environment Frames 1e9 5 10 15 20 25 30 35 40 45 seekavoid_arena_01 qs =rs +γ⋅vs+1 qs =rs +γ⋅V(xs+1) Figure E.3. Variants for estimation of state action value function - average over top 3 runs. 0.0 0.2 0.4 0.6 0.8 1.0 Environment Frames 1e9 5 10 15 20 25 30 35 Return rooms_watermaze 0.0 0.2 0.4 0.6 0.8 1.0 Environment Frames 1e9 4 6 8 10 12 14 16 18 20 22 rooms_keys_doors_puzzle 0.0 0.2 0.4 0.6 0.8 1.0 Environment Frames 1e9 −2 0 2 4 6 8 10 12 lasertag_three_opponents_small 0.0 0.2 0.4 0.6 0.8 1.0 Environment Frames 1e9 0 20 40 60 80 100 120 140 160 explore_goal_locations_small 0.0 0.2 0.4 0.6 0.8 1.0 Environment Frames 1e9 0 5 10 15 20 25 30 35 seekavoid_arena_01 qs =rs +γ⋅vs+1 qs =rs +γ⋅V(xs+1) Figure E.4.


nt Frames 1e9 0 5 10 15 20 25 30 35 seekavoid_arena_01 qs =rs +γ⋅vs+1 qs =rs +γ⋅V(xs+1) Figure E.4. Variants for estimation of state action value function - average over all runs. F. Population Based Training For Population Based Training we used a “burn-in” period of 20 million frames where no evolution is done. This is to stabilise the process and to avoid very rapid initial adaptation which hinders diversity. After collecting 5,000 episode rewards in total, the mean capped human normalised score is calculated and a random instance in the population is selected. If the score of the selected instance is more than an absolute 5% higher, then the selected instance weights and parameters are copied.


e is more than an absolute 5% higher, then the selected instance weights and parameters are copied. No matter if a copy happened or not, each parameter (RMSProp epsilon, learning rate and entropy cost) is permuted with 33% probability by multiplying with either 1.2 or 1/1.2. This is different from Jaderberg et al. (2017) in that our multiplication is unbiased where they use a multiplication of 1.2 or .8. We found that diversity is increased when the parameters are permuted even if no copy happened. We reconstruct the learning curves of the PBT runs in Figure 5 by backtracking through the ancestry of copied checkpoints for selected instances.


IMPALA: Importance Weighted Actor-Learner Architectures 0.0 0.2 0.4 0.6 0.8 1.0 Environment Frames 1e10 0.0000 0.0001 0.0002 0.0003 0.0004 0.0005 0.0006 0.0007 Learning Rate IMPALA - PBT - 8 GPUs IMPALA - PBT - 1 GPU Figure F.1. Learning rate schedule that is discovered by the PBT Jaderberg et al. (2017) method compared against the linear annealing schedule of the best run from the parameter sweep (red line). G. Atari Experiments All agents trained on Atari are equipped only with a feed forward network and pre-process frames in the same way as described in Mnih et al. (2016). When training experts agents, we use the same hyperparameters for each game for both IMPALA and A3C. These hyperparameters are the result of tuning A3C with a shallow network on the following games: breakout, pong, sp


eters are the result of tuning A3C with a shallow network on the following games: breakout, pong, space invaders, seaquest, beam rider, qbert. Following related work, experts use game-speciﬁc action sets. The multi-task agent was equipped with a feed forward residual network (see Figure 3 ). The learning rate, entropy regularisation, RMSProp ε and gradient clipping threshold were adapted through population based training. To be able to use the same policy layer on all Atari games in the multi-task setting we train the multi-task agent on the full Atari action set consisting of 18 actions. Agents were trained using the following set of hyperparameters:


IMPALA: Importance Weighted Actor-Learner Architectures Parameter Value Image Width 84 Image Height 84 Grayscaling Yes Action Repetitions 4 Max-pool over last N action repeat frames 2 Frame Stacking 4 End of episode when life lost Yes Reward Clipping [-1, 1] Unroll Length (n) 20 Batch size 32 Discount (γ) 0.99 Baseline loss scaling 0.5 Entropy Regularizer 0.01 RMSProp momentum 0.0 RMSProp ε 0.01 Learning rate 0.0006 Clip global gradient norm 40.0 Learning rate schedule Anneal linearly to 0 From beginning to end of training. Population based training (only multi-task agent) - Population size 24 - Start parameters Same as DMLab-30 sweep - Fitness Mean capped human normalised scores (P l min [1, (st −rt)/(ht −rt)]) /N - Adapted parameters Gradient clipping threshold Entropy regularisation Lea


, (st −rt)/(ht −rt)]) /N - Adapted parameters Gradient clipping threshold Entropy regularisation Learning rate RMSProp ε Table G.1. Hyperparameters for Atari experiments. References Bertsekas, D. P. and Tsitsiklis, J. N. Neuro-Dynamic Programming. Athena Scientiﬁc, 1996. Dayan, P. and Sejnowski, T. J. TD(λ) converges with probability 1. Machine Learning, 14(1):295–301, 1994. doi: 10.1023/A:1022657612745. Jaderberg, M., Dalibard, V., Osindero, S., Czarnecki, W. M., Donahue, J., Razavi, A., Vinyals, O., Green, T., Dunning, I., Simonyan, K., Fernando, C., and Kavukcuoglu, K. Population based training of neural networks. CoRR, abs/1711.09846, 2017. Kushner, H. and Yin, G. Stochastic Approximation and Recursive Algorithms and Applications. Stochastic Modelling and Applied Probability.


roximation and Recursive Algorithms and Applications. Stochastic Modelling and Applied Probability. Springer New York, 2003. ISBN 9780387008943. Mnih, V., Badia, A. P., Mirza, M., Graves, A., Lillicrap, T. P., Harley, T., Silver, D., and Kavukcuoglu, K. Asynchronous methods for deep reinforcement learning. ICML, 2016.


Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model Julian Schrittwieser,1∗Ioannis Antonoglou,1,2∗Thomas Hubert,1∗ Karen Simonyan,1 Laurent Sifre,1 Simon Schmitt,1 Arthur Guez,1 Edward Lockhart,1 Demis Hassabis,1 Thore Graepel,1,2 Timothy Lillicrap,1 David Silver1,2∗ 1DeepMind, 6 Pancras Square, London N1C 4AG. 2University College London, Gower Street, London WC1E 6BT. ∗These authors contributed equally to this work. Abstract Constructing agents with planning capabilities has long been one of the main challenges in the pursuit of artiﬁcial intelligence. Tree-based planning methods have enjoyed huge success in challenging domains, such as chess and Go, where a perfect simulator is available.


huge success in challenging domains, such as chess and Go, where a perfect simulator is available. However, in real-world problems the dynamics governing the environment are often complex and unknown. In this work we present the MuZero algorithm which, by combining a tree-based search with a learned model, achieves superhuman performance in a range of challenging and visually complex domains, without any knowledge of their underlying dynamics. MuZero learns a model that, when applied iteratively, predicts the quantities most directly relevant to planning: the reward, the action-selection policy, and the value function. When evaluated on 57 different Atari games - the canonical video game environment for testing AI techniques, in which model-based planning approaches have historically stru


vironment for testing AI techniques, in which model-based planning approaches have historically struggled - our new algorithm achieved a new state of the art. When evaluated on Go, chess and shogi, without any knowledge of the game rules, MuZero matched the superhuman performance of the AlphaZero algorithm that was supplied with the game rules. 1 Introduction Planning algorithms based on lookahead search have achieved remarkable successes in artiﬁcial intelligence. Human world champions have been defeated in classic games such as checkers [34], chess [5], Go [38] and poker [3, 26], and planning algorithms have had real-world impact in applications from logistics [47] to chemical synthesis [37].


gorithms have had real-world impact in applications from logistics [47] to chemical synthesis [37]. However, these planning algorithms all rely on knowledge of the environment’s dynamics, such as the rules of the game or an accurate simulator, preventing their direct application to real-world domains like robotics, industrial control, or intelligent assistants. Model-based reinforcement learning (RL) [42] aims to address this issue by ﬁrst learning a model of the environment’s dynamics, and then planning with respect to the learned model. Typically, these models have either focused on reconstructing the true environmental state [8, 16, 24], or the sequence of full observations [14, 20]. However, prior work [4, 14, 20] remains far from the state of the art in visually rich domains, such as


ver, prior work [4, 14, 20] remains far from the state of the art in visually rich domains, such as Atari 2600 games [2]. Instead, the most successful methods are based on model-free RL [9, 21, 18] – i.e. they estimate the optimal policy and/or value function directly from interactions with the environment. However, model-free algorithms are in turn far from the state of the art in domains that require precise and sophisticated lookahead, such as chess and Go. 1 arXiv:1911.08265v2 [cs.LG] 21 Feb 2020


In this paper, we introduce MuZero, a new approach to model-based RL that achieves state-of-the-art performance in Atari 2600, a visually complex set of domains, while maintaining superhuman performance in precision planning tasks such as chess, shogi and Go. MuZero builds upon AlphaZero’s [39] powerful search and search-based policy iteration algorithms, but incorporates a learned model into the training procedure. MuZero also extends AlphaZero to a broader set of environments including single agent domains and non-zero rewards at intermediate time-steps. The main idea of the algorithm (summarized in Figure 1) is to predict those aspects of the future that are directly relevant for planning. The model receives the observation (e.g.


cts of the future that are directly relevant for planning. The model receives the observation (e.g. an image of the Go board or the Atari screen) as an input and transforms it into a hidden state. The hidden state is then updated iteratively by a recurrent process that receives the previous hidden state and a hypothetical next action. At every one of these steps the model predicts the policy (e.g. the move to play), value function (e.g. the predicted winner), and immediate reward (e.g. the points scored by playing a move). The model is trained end-to-end, with the sole objective of accurately estimating these three important quantities, so as to match the improved estimates of policy and value generated by search as well as the observed reward.


atch the improved estimates of policy and value generated by search as well as the observed reward. There is no direct constraint or requirement for the hidden state to capture all information necessary to reconstruct the original observation, drastically reducing the amount of information the model has to maintain and predict; nor is there any requirement for the hidden state to match the unknown, true state of the environment; nor any other constraints on the semantics of state. Instead, the hidden states are free to represent state in whatever way is relevant to predicting current and future values and policies. Intuitively, the agent can invent, internally, the rules or dynamics that lead to most accurate planning.


ively, the agent can invent, internally, the rules or dynamics that lead to most accurate planning. 2 Prior Work Reinforcement learning may be subdivided into two principal categories: model-based, and model-free [42]. Model-based RL constructs, as an intermediate step, a model of the environment. Classically, this model is represented by a Markov-decision process (MDP) [31] consisting of two components: a state transition model, predicting the next state, and a reward model, predicting the expected reward during that transition. The model is typically conditioned on the selected action, or a temporally abstract behavior such as an option [43]. Once a model has been constructed, it is straightforward to apply MDP planning algorithms, such as value iteration [31] or Monte-Carlo tree search


htforward to apply MDP planning algorithms, such as value iteration [31] or Monte-Carlo tree search (MCTS) [7], to compute the optimal value or optimal policy for the MDP. In large or partially observed environments, the algorithm must ﬁrst construct the state representation that the model should predict. This tripartite separation between representation learning, model learning, and planning is potentially problematic since the agent is not able to optimize its representation or model for the purpose of effective planning, so that, for example modeling errors may compound during planning. A common approach to model-based RL focuses on directly modeling the observation stream at the pixellevel.


n approach to model-based RL focuses on directly modeling the observation stream at the pixellevel. It has been hypothesized that deep, stochastic models may mitigate the problems of compounding error [14, 20]. However, planning at pixel-level granularity is not computationally tractable in large scale problems. Other methods build a latent state-space model that is sufﬁcient to reconstruct the observation stream at pixel level [48, 49], or to predict its future latent states [13, 11], which facilitates more efﬁcient planning but still focuses the majority of the model capacity on potentially irrelevant detail. None of these prior methods has constructed a model that facilitates effective planning in visually complex domains such as Atari; results lag behind well-tuned, model-free methods,


anning in visually complex domains such as Atari; results lag behind well-tuned, model-free methods, even in terms of data efﬁciency [45]. A quite different approach to model-based RL has recently been developed, focused end-to-end on predicting the value function [41]. The main idea of these methods is to construct an abstract MDP model such that planning in the abstract MDP is equivalent to planning in the real environment. This equivalence is achieved by ensuring value equivalence, i.e. that, starting from the same real state, the cumulative reward of a trajectory through the abstract MDP matches the cumulative reward of a trajectory in the real environment. The predictron [41] ﬁrst introduced value equivalent models for predicting value (without actions).


The predictron [41] ﬁrst introduced value equivalent models for predicting value (without actions). Although the underlying model still takes the form of an MDP, there is no requirement for its transition model to match real states in the environment. Instead the MDP model is viewed as a hidden layer of a deep neural network. The unrolled MDP is trained such that the expected cumulative sum of rewards matches the expected value with respect to the real environment, e.g. by temporal-difference learning. 2


Figure 1: Planning, acting, and training with a learned model. (A) How MuZero uses its model to plan. The model consists of three connected components for representation, dynamics and prediction. Given a previous hidden state sk−1 and a candidate action ak, the dynamics function g produces an immediate reward rk and a new hidden state sk. The policy pk and value function vk are computed from the hidden state sk by a prediction function f. The initial hidden state s0 is obtained by passing the past observations (e.g. the Go board or Atari screen) into a representation function h. (B) How MuZero acts in the environment. A Monte-Carlo Tree Search is performed at each timestep t, as described in A.


s in the environment. A Monte-Carlo Tree Search is performed at each timestep t, as described in A. An action at+1 is sampled from the search policy πt, which is proportional to the visit count for each action from the root node. The environment receives the action and generates a new observation ot+1 and reward ut+1. At the end of the episode the trajectory data is stored into a replay buffer. (C) How MuZero trains its model. A trajectory is sampled from the replay buffer. For the initial step, the representation function h receives as input the past observations o1, ..., ot from the selected trajectory. The model is subsequently unrolled recurrently for K steps. At each step k, the dynamics function g receives as input the hidden state sk−1 from the previous step and the real action at+k


s function g receives as input the hidden state sk−1 from the previous step and the real action at+k. The parameters of the representation, dynamics and prediction functions are jointly trained, end-to-end by backpropagation-through-time, to predict three quantities: the policy pk ≈πt+k, value function vk ≈zt+k, and reward rt+k ≈ut+k, where zt+k is a sample return: either the ﬁnal reward (board games) or n-step return (Atari). Value equivalent models were subsequently extended to optimising value (with actions). TreeQN [10] learns an abstract MDP model, such that a tree search over that model (represented by a tree-structured neural network) approximates the optimal value function. Value iteration networks [44] learn a local MDP model, such that value iteration over that model (represented


ration networks [44] learn a local MDP model, such that value iteration over that model (represented by a convolutional neural network) approximates the optimal value function. Value prediction networks [28] are perhaps the closest precursor to MuZero: they learn an MDP model grounded in real actions; the unrolled MDP is trained such that the cumulative sum of rewards, conditioned on the actual sequence of actions generated by a simple lookahead search, matches the real environment. Unlike MuZero there is no policy prediction, and the search only utilizes value prediction. 3 MuZero Algorithm We now describe the MuZero algorithm in more detail. Predictions are made at each time-step t, for each of k = 1...K steps, by a model µθ, with parameters θ, conditioned on past observations o1, ..., o


ch of k = 1...K steps, by a model µθ, with parameters θ, conditioned on past observations o1, ..., ot and future actions at+1, ..., at+k. The model predicts three future quantities: the policy pk t ≈π(at+k+1|o1, ..., ot, at+1, ..., at+k), the value function vk t ≈E [ut+k+1 + γut+k+2 + ...|o1, ..., ot, at+1, ..., at+k], and the immediate reward rk t ≈ut+k, 3


where u. is the true, observed reward, π is the policy used to select real actions, and γ is the discount function of the environment. Internally, at each time-step t (subscripts t suppressed for simplicity), the model is represented by the combination of a representation function, a dynamics function, and a prediction function. The dynamics function, rk, sk = gθ(sk−1, ak), is a recurrent process that computes, at each hypothetical step k, an immediate reward rk and an internal state sk. It mirrors the structure of an MDP model that computes the expected reward and state transition for a given state and action [31]. However, unlike traditional approaches to model-based RL [42], this internal state sk has no semantics of environment state attached to it – it is simply the hidden state of th


state sk has no semantics of environment state attached to it – it is simply the hidden state of the overall model, and its sole purpose is to accurately predict relevant, future quantities: policies, values, and rewards. In this paper, the dynamics function is represented deterministically; the extension to stochastic transitions is left for future work. The policy and value functions are computed from the internal state sk by the prediction function, pk, vk = fθ(sk), akin to the joint policy and value network of AlphaZero. The “root” state s0 is initialized using a representation function that encodes past observations, s0 = hθ(o1, ..., ot); again this has no special semantics beyond its support for future predictions.


0 = hθ(o1, ..., ot); again this has no special semantics beyond its support for future predictions. Given such a model, it is possible to search over hypothetical future trajectories a1, ..., ak given past observations o1, ..., ot. For example, a naive search could simply select the k step action sequence that maximizes the value function. More generally, we may apply any MDP planning algorithm to the internal rewards and state space induced by the dynamics function. Speciﬁcally, we use an MCTS algorithm similar to AlphaZero’s search, generalized to allow for single agent domains and intermediate rewards (see Methods). At each internal node, it makes use of the policy, value and reward estimates produced by the current model parameters θ.


it makes use of the policy, value and reward estimates produced by the current model parameters θ. The MCTS algorithm outputs a recommended policy πt and estimated value νt. An action at+1 ∼πt is then selected. All parameters of the model are trained jointly to accurately match the policy, value, and reward, for every hypothetical step k, to corresponding target values observed after k actual time-steps have elapsed. Similarly to AlphaZero, the improved policy targets are generated by an MCTS search; the ﬁrst objective is to minimise the error between predicted policy pk t and search policy πt+k. Also like AlphaZero, the improved value targets are generated by playing the game or MDP. However, unlike AlphaZero, we allow for long episodes with discounting and intermediate rewards by bootst


er, unlike AlphaZero, we allow for long episodes with discounting and intermediate rewards by bootstrapping n steps into the future from the search value, zt = ut+1 +γut+2 +...+ γn−1ut+n + γnνt+n. Final outcomes {lose, draw, win} in board games are treated as rewards ut ∈{−1, 0, +1} occuring at the ﬁnal step of the episode. Speciﬁcally, the second objective is to minimize the error between the predicted value vk t and the value target, zt+k 1. The reward targets are simply the observed rewards; the third objective is therefore to minimize the error between the predicted reward rk t and the observed reward ut+k. Finally, an L2 regularization term is also added, leading to the overall loss: lt(θ) = K X k=0 lr(ut+k, rk t ) + lv(zt+k, vk t ) + lp(πt+k, pk t ) + c||θ||2 (1) where lr, lv, and lp


(θ) = K X k=0 lr(ut+k, rk t ) + lv(zt+k, vk t ) + lp(πt+k, pk t ) + c||θ||2 (1) where lr, lv, and lp are loss functions for reward, value and policy respectively. Supplementary Figure S2 summarizes the equations governing how the MuZero algorithm plans, acts, and learns. 4 Results We applied the MuZero algorithm to the classic board games Go, chess and shogi 2, as benchmarks for challenging planning problems, and to all 57 games in the Atari Learning Environment [2], as benchmarks for visually complex RL domains. In each case we trained MuZero for K = 5 hypothetical steps. Training proceeded for 1 million mini-batches of size 2048 in board games and of size 1024 in Atari. During both training and evaluation, MuZero used 800 simulations for each search in board games, and 50 simulations for


g and evaluation, MuZero used 800 simulations for each search in board games, and 50 simulations for each search in Atari. The representation 1For chess, Go and shogi, the same squared error loss as AlphaZero is used for rewards and values. A cross-entropy loss was found to be more stable than a squared error when encountering rewards and values of variable scale in Atari. Cross-entropy was used for the policy loss in both cases. 2Imperfect information games such as Poker are not directly addressed by our method. 4


Chess Shogi Go Atari rmblkans opopopop 0Z0Z0Z0Z Z0Z0Z0Z0 0Z0Z0Z0Z Z0Z0Z0Z0 POPOPOPO SNAQJBMR 香 桂 銀 金 玉 金 銀 桂 香 飛 角 歩 歩 歩 歩 歩 歩 歩 歩 歩 歩歩歩歩歩歩歩歩歩 角 飛 香桂銀金玉金銀桂香 Figure 2: Evaluation of MuZero throughout training in chess, shogi, Go and Atari. The x-axis shows millions of training steps. For chess, shogi and Go, the y-axis shows Elo rating, established by playing games against AlphaZero using 800 simulations per move for both players. MuZero’s Elo is indicated by the blue line, AlphaZero’s Elo by the horizontal orange line. For Atari, mean (full line) and median (dashed line) human normalized scores across all 57 games are shown on the y-axis. The scores for R2D2 [21], (the previous state of the art in this domain, based on model-free RL) are indicated by the horizontal orange lines.


te of the art in this domain, based on model-free RL) are indicated by the horizontal orange lines. Performance in Atari was evaluated using 50 simulations every fourth time-step, and then repeating the chosen action four times, as in prior work [25]. function uses the same convolutional [23] and residual [15] architecture as AlphaZero, but with 16 residual blocks instead of 20. The dynamics function uses the same architecture as the representation function and the prediction function uses the same architecture as AlphaZero. All networks use 256 hidden planes (see Methods for further details). Figure 2 shows the performance throughout training in each game. In Go, MuZero slightly exceeded the performance of AlphaZero, despite using less computation per node in the search tree (16 residual


e performance of AlphaZero, despite using less computation per node in the search tree (16 residual blocks per evaluation in MuZero compared to 20 blocks in AlphaZero). This suggests that MuZero may be caching its computation in the search tree and using each additional application of the dynamics model to gain a deeper understanding of the position. In Atari, MuZero achieved a new state of the art for both mean and median normalized score across the 57 games of the Arcade Learning Environment, outperforming the previous state-of-the-art method R2D2 [21] (a model-free approach) in 42 out of 57 games, and outperforming the previous best model-based approach SimPLe [20] in all games (see Table S1). We also evaluated a second version of MuZero that was optimised for greater sample efﬁciency.


S1). We also evaluated a second version of MuZero that was optimised for greater sample efﬁciency. Speciﬁcally, it reanalyzes old trajectories by re-running the MCTS using the latest network parameters to provide fresh targets (see Appendix H). When applied to 57 Atari games, using 200 million frames of experience per game, MuZero Reanalyze achieved 731% median normalized score, compared to 192%, 231% and 431% for previous state-ofthe-art model-free approaches IMPALA [9], Rainbow [17] and LASER [36] respectively. To understand the role of the model in MuZero we also ran several experiments, focusing on the board game of Go and the Atari game of Ms. Pacman. First, we tested the scalability of planning (Figure 3A), in the canonical planning problem of Go.


First, we tested the scalability of planning (Figure 3A), in the canonical planning problem of Go. We compared the performance of search in AlphaZero, using a perfect model, to the performance of search in MuZero, using a 5


Agent Median Mean Env. Frames Training Time Training Steps Ape-X [18] 434.1% 1695.6% 22.8B 5 days 8.64M R2D2 [21] 1920.6% 4024.9% 37.5B 5 days 2.16M MuZero 2041.1% 4999.2% 20.0B 12 hours 1M IMPALA [9] 191.8% 957.6% 200M – – Rainbow [17] 231.1% – 200M 10 days – UNREALa [19] 250%a 880%a 250M – – LASER [36] 431% – 200M – – MuZero Reanalyze 731.1% 2168.9% 200M 12 hours 1M Table 1: Comparison of MuZero against previous agents in Atari. We compare separately against agents trained in large (top) and small (bottom) data settings; all agents other than MuZero used model-free RL techniques. Mean and median scores are given, compared to human testers. The best results are highlighted in bold. MuZero sets a new state of the art in both settings. aHyper-parameters were tuned per game. learned model.


sets a new state of the art in both settings. aHyper-parameters were tuned per game. learned model. Speciﬁcally, the fully trained AlphaZero or MuZero was evaluated by comparing MCTS with different thinking times. MuZero matched the performance of a perfect model, even when doing much larger searches (up to 10s thinking time) than those from which the model was trained (around 0.1s thinking time, see also Figure S3A). We also investigated the scalability of planning across all Atari games (see Figure 3B). We compared MCTS with different numbers of simulations, using the fully trained MuZero. The improvements due to planning are much less marked than in Go, perhaps because of greater model inaccuracy; performance improved slightly with search time, but plateaued at around 100 simulations.


naccuracy; performance improved slightly with search time, but plateaued at around 100 simulations. Even with a single simulation – i.e. when selecting moves solely according to the policy network – MuZero performed well, suggesting that, by the end of training, the raw policy has learned to internalise the beneﬁts of search (see also Figure S3B). Next, we tested our model-based learning algorithm against a comparable model-free learning algorithm (see Figure 3C). We replaced the training objective of MuZero (Equation 1) with a model-free Q-learning objective (as used by R2D2), and the dual value and policy heads with a single head representing the Q-function Q(·|st). Subsequently, we trained and evaluated the new model without using any search. When evaluated on Ms.


ubsequently, we trained and evaluated the new model without using any search. When evaluated on Ms. Pacman, our model-free algorithm achieved identical results to R2D2, but learned signiﬁcantly slower than MuZero and converged to a much lower ﬁnal score. We conjecture that the search-based policy improvement step of MuZero provides a stronger learning signal than the high bias, high variance targets used by Q-learning. To better understand the nature of MuZero’s learning algorithm, we measured how MuZero’s training scales with respect to the amount of search it uses during training. Figure 3D shows the performance in Ms. Pacman, using an MCTS of different simulation counts per move throughout training.


rformance in Ms. Pacman, using an MCTS of different simulation counts per move throughout training. Surprisingly, and in contrast to previous work [1], even with only 6 simulations per move – fewer than the number of actions – MuZero learned an effective policy and improved rapidly. With more simulations performance jumped signiﬁcantly higher. For analysis of the policy improvement during each individual iteration, see also Figure S3 C and D. 5 Conclusions Many of the breakthroughs in artiﬁcial intelligence have been based on either high-performance planning [5, 38, 39] or model-free reinforcement learning methods [25, 29, 46]. In this paper we have introduced a method that combines the beneﬁts of both approaches.


5, 29, 46]. In this paper we have introduced a method that combines the beneﬁts of both approaches. Our algorithm, MuZero, has both matched the superhuman performance of high-performance planning algorithms in their favored domains – logically complex board games such as chess and Go – and outperformed state-of-the-art model-free RL algorithms in their favored domains – visually complex Atari games. Crucially, our method does not require any knowledge of the game rules or environment dynamics, potentially paving the way towards the application of powerful learning and planning methods to a host of realworld domains for which there exists no perfect simulator. 6


Figure 3: Evaluations of MuZero on Go (A), all 57 Atari Games (B) and Ms. Pacman (C-D). (A) Scaling with search time per move in Go, comparing the learned model with the ground truth simulator. Both networks were trained at 800 simulations per search, equivalent to 0.1 seconds per search. Remarkably, the learned model is able to scale well to up to two orders of magnitude longer searches than seen during training. (B) Scaling of ﬁnal human normalized mean score in Atari with the number of simulations per search. The network was trained at 50 simulations per search. Dark line indicates mean score, shaded regions indicate 25th to 75th and 5th to 95th percentiles. The learned model’s performance increases up to 100 simulations per search.


th to 95th percentiles. The learned model’s performance increases up to 100 simulations per search. Beyond, even when scaling to much longer searches than during training, the learned model’s performance remains stable and only decreases slightly. This contrasts with the much better scaling in Go (A), presumably due to greater model inaccuracy in Atari than Go. (C) Comparison of MCTS based training with Q-learning in the MuZero framework on Ms. Pacman, keeping network size and amount of training constant. The state of the art Q-Learning algorithm R2D2 is shown as a baseline. Our Q-Learning implementation reaches the same ﬁnal score as R2D2, but improves slower and results in much lower ﬁnal performance compared to MCTS based training.


D2, but improves slower and results in much lower ﬁnal performance compared to MCTS based training. (D) Different networks trained at different numbers of simulations per move, but all evaluated at 50 simulations per move. Networks trained with more simulations per move improve faster, consistent with ablation (B), where the policy improvement is larger when using more simulations per move. Surprisingly, MuZero can learn effectively even when training with less simulations per move than are enough to cover all 8 possible actions in Ms. Pacman. 7


6 Acknowledgments Lorrayne Bennett, Oliver Smith and Chris Apps for organizational assistance; Koray Kavukcuoglu for reviewing the paper; Thomas Anthony, Matthew Lai, Nenad Tomasev, Ulrich Paquet, Sumedh Ghaisas for many fruitful discussions; and the rest of the DeepMind team for their support. References [1] Kamyar Azizzadenesheli, Brandon Yang, Weitang Liu, Emma Brunskill, Zachary C. Lipton, and Animashree Anandkumar. Surprising negative results for generative adversarial tree search. CoRR, abs/1806.05780, 2018. [2] Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artiﬁcial Intelligence Research, 47:253–279, 2013. [3] Noam Brown and Tuomas Sandholm.


. Journal of Artiﬁcial Intelligence Research, 47:253–279, 2013. [3] Noam Brown and Tuomas Sandholm. Superhuman ai for heads-up no-limit poker: Libratus beats top professionals. Science, 359(6374):418–424, 2018. [4] Lars Buesing, Theophane Weber, Sebastien Racaniere, SM Eslami, Danilo Rezende, David P Reichert, Fabio Viola, Frederic Besse, Karol Gregor, Demis Hassabis, et al. Learning and querying fast generative models for reinforcement learning. arXiv preprint arXiv:1802.03006, 2018. [5] Murray Campbell, A. Joseph Hoane, Jr., and Feng-hsiung Hsu. Deep blue. Artif. Intell., 134(1-2):57–83, January 2002. [6] R. Coulom. Whole-history rating: A Bayesian rating system for players of time-varying strength. In International Conference on Computers and Games, pages 113–124, 2008.


of time-varying strength. In International Conference on Computers and Games, pages 113–124, 2008. [7] R´emi Coulom. Efﬁcient selectivity and backup operators in monte-carlo tree search. In International conference on computers and games, pages 72–83. Springer, 2006. [8] MP. Deisenroth and CE. Rasmussen. Pilco: A model-based and data-efﬁcient approach to policy search. In Proceedings of the 28th International Conference on Machine Learning, ICML 2011, pages 465–472. Omnipress, 2011. [9] Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Volodymir Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, Iain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures.


, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In Proceedings of the International Conference on Machine Learning (ICML), 2018. [10] Gregory Farquhar, Tim Rocktaeschel, Maximilian Igl, and Shimon Whiteson. TreeQN and ATreec: Differentiable tree planning for deep reinforcement learning. In International Conference on Learning Representations, 2018. [11] Carles Gelada, Saurabh Kumar, Jacob Buckman, Oﬁr Nachum, and Marc G. Bellemare. DeepMDP: Learning continuous latent space models for representation learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 2170–2179, Long Beach, California, USA, 09–15


97 of Proceedings of Machine Learning Research, pages 2170–2179, Long Beach, California, USA, 09–15 Jun 2019. PMLR. [12] Cloud tpu. https://cloud.google.com/tpu/. Accessed: 2019. [13] David Ha and J¨urgen Schmidhuber. Recurrent world models facilitate policy evolution. In Proceedings of the 32Nd International Conference on Neural Information Processing Systems, NIPS’18, pages 2455–2467, USA, 2018. Curran Associates Inc. 8


[14] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. arXiv preprint arXiv:1811.04551, 2018. [15] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In 14th European Conference on Computer Vision, pages 630–645, 2016. [16] Nicolas Heess, Greg Wayne, David Silver, Timothy Lillicrap, Yuval Tassa, and Tom Erez. Learning continuous control policies by stochastic value gradients. In Proceedings of the 28th International Conference on Neural Information Processing Systems - Volume 2, NIPS’15, pages 2944–2952, Cambridge, MA, USA, 2015. MIT Press.


ation Processing Systems - Volume 2, NIPS’15, pages 2944–2952, Cambridge, MA, USA, 2015. MIT Press. [17] Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Thirty-Second AAAI Conference on Artiﬁcial Intelligence, 2018. [18] Dan Horgan, John Quan, David Budden, Gabriel Barth-Maron, Matteo Hessel, Hado van Hasselt, and David Silver. Distributed prioritized experience replay. In International Conference on Learning Representations, 2018. [19] Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks.


ibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint arXiv:1611.05397, 2016. [20] Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model-based reinforcement learning for atari. arXiv preprint arXiv:1903.00374, 2019. [21] Steven Kapturowski, Georg Ostrovski, Will Dabney, John Quan, and Remi Munos. Recurrent experience replay in distributed reinforcement learning. In International Conference on Learning Representations, 2019. [22] Levente Kocsis and Csaba Szepesv´ari. Bandit based monte-carlo planning. In European conference on machine learning, pages 282–293. Springer, 2006.


ed monte-carlo planning. In European conference on machine learning, pages 282–293. Springer, 2006. [23] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classiﬁcation with deep convolutional neural networks. In Advances in neural information processing systems, pages 1097–1105, 2012. [24] Sergey Levine and Pieter Abbeel. Learning neural network policies with guided policy search under unknown dynamics. In Z. Ghahramani, M. Welling, C. Cortes, N. D. Lawrence, and K. Q. Weinberger, editors, Advances in Neural Information Processing Systems 27, pages 1071–1079. Curran Associates, Inc., 2014. [25] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al.


ess, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, 518(7540):529, 2015. [26] Matej Moravˇc´ık, Martin Schmid, Neil Burch, Viliam Lis`y, Dustin Morrill, Nolan Bard, Trevor Davis, Kevin Waugh, Michael Johanson, and Michael Bowling. Deepstack: Expert-level artiﬁcial intelligence in heads-up no-limit poker. Science, 356(6337):508–513, 2017. [27] Arun Nair, Praveen Srinivasan, Sam Blackwell, Cagdas Alcicek, Rory Fearon, Alessandro De Maria, Vedavyas Panneershelvam, Mustafa Suleyman, Charles Beattie, Stig Petersen, Shane Legg, Volodymyr Mnih, Koray Kavukcuoglu, and David Silver. Massively parallel methods for deep reinforcement learning. CoRR, abs/1507.04296, 2015.


vid Silver. Massively parallel methods for deep reinforcement learning. CoRR, abs/1507.04296, 2015. [28] Junhyuk Oh, Satinder Singh, and Honglak Lee. Value prediction network. In Advances in Neural Information Processing Systems, pages 6118–6128, 2017. [29] OpenAI. Openai ﬁve. https://blog.openai.com/openai-five/, 2018. 9


[30] Tobias Pohlen, Bilal Piot, Todd Hester, Mohammad Gheshlaghi Azar, Dan Horgan, David Budden, Gabriel Barth-Maron, Hado van Hasselt, John Quan, Mel Veˇcer´ık, et al. Observe and look further: Achieving consistent performance on atari. arXiv preprint arXiv:1805.11593, 2018. [31] Martin L. Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming. John Wiley & Sons, Inc., New York, NY, USA, 1st edition, 1994. [32] Christopher D Rosin. Multi-armed bandits with episode context. Annals of Mathematics and Artiﬁcial Intelligence, 61(3):203–230, 2011. [33] Maarten PD Schadd, Mark HM Winands, H Jaap Van Den Herik, Guillaume MJ-B Chaslot, and Jos WHM Uiterwijk. Single-player monte-carlo tree search. In International Conference on Computers and Games, pages 1–12.


gle-player monte-carlo tree search. In International Conference on Computers and Games, pages 1–12. Springer, 2008. [34] Jonathan Schaeffer, Joseph Culberson, Norman Treloar, Brent Knight, Paul Lu, and Duane Szafron. A world championship caliber checkers program. Artiﬁcial Intelligence, 53(2-3):273–289, 1992. [35] Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay. In International Conference on Learning Representations, Puerto Rico, 2016. [36] Simon Schmitt, Matteo Hessel, and Karen Simonyan. Off-policy actor-critic with shared experience replay. arXiv preprint arXiv:1909.11583, 2019. [37] Marwin HS Segler, Mike Preuss, and Mark P Waller. Planning chemical syntheses with deep neural networks and symbolic ai. Nature, 555(7698):604, 2018.


Planning chemical syntheses with deep neural networks and symbolic ai. Nature, 555(7698):604, 2018. [38] David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George van den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of Go with deep neural networks and tree search. Nature, 529(7587):484–489, January 2016. [39] David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al.


lou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(6419):1140–1144, 2018. [40] David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, Yutian Chen, Timothy Lillicrap, Fan Hui, Laurent Sifre, George van den Driessche, Thore Graepel, and Demis Hassabis. Mastering the game of go without human knowledge. Nature, 550:354–359, October 2017. [41] David Silver, Hado van Hasselt, Matteo Hessel, Tom Schaul, Arthur Guez, Tim Harley, Gabriel DulacArnold, David Reichert, Neil Rabinowitz, Andre Barreto, et al.


rthur Guez, Tim Harley, Gabriel DulacArnold, David Reichert, Neil Rabinowitz, Andre Barreto, et al. The predictron: End-to-end learning and planning. In Proceedings of the 34th International Conference on Machine Learning-Volume 70, pages 3191–3199. JMLR. org, 2017. [42] Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. The MIT Press, second edition, 2018. [43] Richard S Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning. Artiﬁcial intelligence, 112(1-2):181–211, 1999. [44] Aviv Tamar, Yi Wu, Garrett Thomas, Sergey Levine, and Pieter Abbeel. Value iteration networks. In Advances in Neural Information Processing Systems, pages 2154–2162, 2016.


ue iteration networks. In Advances in Neural Information Processing Systems, pages 2154–2162, 2016. [45] Hado van Hasselt, Matteo Hessel, and John Aslanides. When to use parametric models in reinforcement learning? arXiv preprint arXiv:1906.05243, 2019. 10


[46] Oriol Vinyals, Igor Babuschkin, Wojciech M Czarnecki, Micha¨el Mathieu, Andrew Dudzik, Junyoung Chung, David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in StarCraft II using multi-agent reinforcement learning. Nature, pages 1–5, 2019. [47] I Vlahavas and I Refanidis. Planning and scheduling. EETN, Greece, Tech. Rep, 2013. [48] Niklas Wahlstr¨om, Thomas B. Sch¨on, and Marc Peter Deisenroth. From pixels to torques: Policy learning with deep dynamical models. CoRR, abs/1502.02251, 2015. [49] Manuel Watter, Jost Tobias Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images.


n Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. In Proceedings of the 28th International Conference on Neural Information Processing Systems - Volume 2, NIPS’15, pages 2746–2754, Cambridge, MA, USA, 2015. MIT Press. Supplementary Materials • Pseudocode description of the MuZero algorithm. • Data for Figures 2, 3, S2, S3, S4 and Tables 1, S1, S2 in JSON format. Supplementary materials can be accessed from the ancillary ﬁle section of the arXiv submission. Appendix A Comparison to AlphaZero MuZero is designed for a more general setting than AlphaGo Zero [40] and AlphaZero [39]. In AlphaGo Zero and AlphaZero the planning process makes use of two separate components: a simulator implements the rules of the game, which are used to update the s


wo separate components: a simulator implements the rules of the game, which are used to update the state of the game while traversing the search tree; and a neural network jointly predicts the corresponding policy and value of a board position produced by the simulator (see Figure 1 A). Speciﬁcally, AlphaGo Zero and AlphaZero use knowledge of the rules of the game in three places: (1) state transitions in the search tree, (2) actions available at each node of the search tree, (3) episode termination within the search tree. In MuZero, all of these have been replaced with the use of a single implicit model learned by a neural network (see Figure 1 B): 1) State transitions. AlphaZero had access to a perfect simulator of the true dynamics process.


B): 1) State transitions. AlphaZero had access to a perfect simulator of the true dynamics process. In contrast, MuZero employs a learned dynamics model within its search. Under this model, each node in the tree is represented by a corresponding hidden state; by providing a hidden state sk−1 and an action ak to the model the search algorithm can transition to a new node sk = g(sk−1, ak). 2) Actions available. AlphaZero used the set of legal actions obtained from the simulator to mask the prior produced by the network everywhere in the search tree. MuZero only masks legal actions at the root of the search tree where the environment can be queried, but does not perform any masking within the search tree.


tree where the environment can be queried, but does not perform any masking within the search tree. This is possible because the network rapidly learns not to predict actions that never occur in the trajectories it is trained on. 3) Terminal nodes. AlphaZero stopped the search at tree nodes representing terminal states and used the terminal value provided by the simulator instead of the value produced by the network. MuZero does not give special treatment to terminal nodes and always uses the value predicted by the network. Inside the tree, the search can proceed past a terminal node - in this case the network is expected to always predict the same value. This is achieved by treating terminal states as absorbing states during training.


t the same value. This is achieved by treating terminal states as absorbing states during training. In addition, MuZero is designed to operate in the general reinforcement learning setting: single-agent domains with discounted intermediate rewards of arbitrary magnitude. In contrast, AlphaGo Zero and AlphaZero were designed to operate in two-player games with undiscounted terminal rewards of ±1. 11


Appendix B Search We now describe the search algorithm used by MuZero. Our approach is based upon Monte-Carlo tree search with upper conﬁdence bounds, an approach to planning that converges asymptotically to the optimal policy in single agent domains and to the minimax value function in zero sum games [22]. Every node of the search tree is associated with an internal state s. For each action a from s there is an edge (s, a) that stores a set of statistics {N(s, a), Q(s, a), P(s, a), R(s, a), S(s, a)}, respectively representing visit counts N, mean value Q, policy P, reward R, and state transition S. Similar to AlphaZero, the search is divided into three stages, repeated for a number of simulations.


imilar to AlphaZero, the search is divided into three stages, repeated for a number of simulations. Selection: Each simulation starts from the internal root state s0, and ﬁnishes when the simulation reaches a leaf node sl. For each hypothetical time-step k = 1...l of the simulation, an action ak is selected according to the stored statistics for internal state sk−1, by maximizing over an upper conﬁdence bound [32][39], ak = arg max a  Q(s, a) + P(s, a) · pP b N(s, b) 1 + N(s, a)  c1 + log P b N(s, b) + c2 + 1 c2  (2) The constants c1 and c2 are used to control the inﬂuence of the prior P(s, a) relative to the value Q(s, a) as nodes are visited more often. In our experiments, c1 = 1.25 and c2 = 19652.


to the value Q(s, a) as nodes are visited more often. In our experiments, c1 = 1.25 and c2 = 19652. For k < l, the next state and reward are looked up in the state transition and reward table sk = S(sk−1, ak), rk = R(sk−1, ak). Expansion: At the ﬁnal time-step l of the simulation, the reward and state are computed by the dynamics function, rl, sl = gθ(sl−1, al), and stored in the corresponding tables, R(sl−1, al) = rl, S(sl−1, al) = sl. The policy and value are computed by the prediction function, pl, vl = fθ(sl). A new node, corresponding to state sl is added to the search tree. Each edge (sl, a) from the newly expanded node is initialized to {N(sl, a) = 0, Q(sl, a) = 0, P(sl, a) = pl}. Note that the search algorithm makes at most one call to the dynamics function and prediction function


e that the search algorithm makes at most one call to the dynamics function and prediction function respectively per simulation; the computational cost is of the same order as in AlphaZero. Backup: At the end of the simulation, the statistics along the trajectory are updated. The backup is generalized to the case where the environment can emit intermediate rewards, have a discount γ different from 1, and the value estimates are unbounded 3. For k = l...0, we form an l −k-step estimate of the cumulative discounted reward, bootstrapping from the value function vl, Gk = l−1−k X τ=0 γτrk+1+τ + γl−kvl (3) For k = l...1, we update the statistics for each edge (sk−1, ak) in the simulation path as follows, Q(sk−1, ak) := N(sk−1, ak) · Q(sk−1, ak) + Gk N(sk−1, ak) + 1 N(sk−1, ak) := N(sk−1, ak) + 1


ollows, Q(sk−1, ak) := N(sk−1, ak) · Q(sk−1, ak) + Gk N(sk−1, ak) + 1 N(sk−1, ak) := N(sk−1, ak) + 1 (4) In two-player zero sum games the value functions are assumed to be bounded within the [0, 1] interval. This choice allows us to combine value estimates with probabilities using the pUCT rule (Eqn 2). However, since in many environments the value is unbounded, it is necessary to adjust the pUCT rule. A simple solution would be to use the maximum score that can be observed in the environment to either re-scale the value or set the pUCT constants appropriately [33]. However, both solutions are game speciﬁc and require adding prior knowledge to the MuZero algorithm. To avoid this, MuZero computes normalized Q value estimates Q ∈[0, 1] by using the minimum-maximum values observed in the sear


utes normalized Q value estimates Q ∈[0, 1] by using the minimum-maximum values observed in the search tree up to that point. When a node is reached during the selection stage, the algorithm computes the normalized Q values of its edges to be used in the pUCT rule using the equation below: Q(sk−1, ak) = Q(sk−1, ak) −mins,a∈T ree Q(s, a) maxs,a∈T ree Q(s, a) −mins,a∈T ree Q(s, a) (5) 3In board games the discount is assumed to be 1 and there are no intermediate rewards. 12


Appendix C Hyperparameters For simplicity we preferentially use the same architectural choices and hyperparameters as in previous work. Speciﬁcally, we started with the network architecture and search choices of AlphaZero [39]. For board games, we use the same UCB constants, dirichlet exploration noise and the same 800 simulations per search as in AlphaZero. Due to the much smaller branching factor and simpler policies in Atari, we only used 50 simulations per search to speed up experiments. As shown in Figure 3B, the algorithm is not very sensitive to this choice. We also use the same discount (0.997) and value transformation (see Network Architecture section) as R2D2 [21]. For parameter values not mentioned in the text, please refer to the pseudocode.


tion) as R2D2 [21]. For parameter values not mentioned in the text, please refer to the pseudocode. Appendix D Data Generation To generate training data, the latest checkpoint of the network (updated every 1000 training steps) is used to play games with MCTS. In the board games Go, chess and shogi the search is run for 800 simulations per move to pick an action; in Atari due to the much smaller action space 50 simulations per move are sufﬁcient. For board games, games are sent to the training job as soon as they ﬁnish. Due to the much larger length of Atari games (up to 30 minutes or 108,000 frames), intermediate sequences are sent every 200 moves. In board games, the training job keeps an in-memory replay buffer of the most recent 1 million games received; in Atari, where the visual obser


n-memory replay buffer of the most recent 1 million games received; in Atari, where the visual observations are larger, the most recent 125 thousand sequences of length 200 are kept. During the generation of experience in the board game domains, the same exploration scheme as the one described in AlphaZero [39] is used. Using a variation of this scheme, in the Atari domain actions are sampled from the visit count distribution throughout the duration of each game, instead of just the ﬁrst k moves. Moreover, the visit count distribution is parametrized using a temperature parameter T: pα = N(α)1/T P b N(b)1/T (6) T is decayed as a function of the number of training steps of the network. Speciﬁcally, for the ﬁrst 500k training steps a temperature of 1 is used, for the next 250k steps a temper


cally, for the ﬁrst 500k training steps a temperature of 1 is used, for the next 250k steps a temperature of 0.5 and for the remaining 250k a temperature of 0.25. This ensures that the action selection becomes greedier as training progresses. Appendix E Network Input Representation Function The history over board states used as input to the representation function for Go, chess and shogi is represented similarly to AlphaZero [39]. In Go and shogi we encode the last 8 board states as in AlphaZero; in chess we increased the history to the last 100 board states to allow correct prediction of draws. For Atari, the input of the representation function includes the last 32 RGB frames at resolution 96x96 along with the last 32 actions that led to each of those frames.


32 RGB frames at resolution 96x96 along with the last 32 actions that led to each of those frames. We encode the historical actions because unlike board games, an action in Atari does not necessarily have a visible effect on the observation. RGB frames are encoded as one plane per color, rescaled to the range [0, 1], for red, green and blue respectively. We perform no other normalization, whitening or other preprocessing of the RGB input. Historical actions are encoded as simple bias planes, scaled as a/18 (there are 18 total actions in Atari). Dynamics Function The input to the dynamics function is the hidden state produced by the representation function or previous application of the dynamics function, concatenated with a representation of the action for the transition.


tion of the dynamics function, concatenated with a representation of the action for the transition. Actions are encoded spatially in planes of the same resolution as the hidden state. In Atari, this resolution is 6x6 (see description of downsampling in Network Architecture section), in board games this is the same as the board size (19x19 for Go, 8x8 for chess, 9x9 for shogi). 13


In Go, a normal action (playing a stone on the board) is encoded as an all zero plane, with a single one in the position of the played stone. A pass is encoded as an all zero plane. In chess, 8 planes are used to encode the action. The ﬁrst one-hot plane encodes which position the piece was moved from. The next two planes encode which position the piece was moved to: a one-hot plane to encode the target position, if on the board, and a second binary plane to indicate whether the target was valid (on the board) or not. This is necessary because for simplicity our policy action space enumerates a superset of all possible actions, not all of which are legal, and we use the same action space for policy prediction and to encode the dynamics function input.


, and we use the same action space for policy prediction and to encode the dynamics function input. The remaining ﬁve binary planes are used to indicate the type of promotion, if any (queen, knight, bishop, rook, none). The encoding for shogi is similar, with a total of 11 planes. We use the ﬁrst 8 planes to indicate where the piece moved from - either a board position (ﬁrst one-hot plane) or the drop of one of the seven types of prisoner (remaining 7 binary planes). The next two planes are used to encode the target as in chess. The remaining binary plane indicates whether the move was a promotion or not. In Atari, an action is encoded as a one hot vector which is tiled appropriately into planes.


r not. In Atari, an action is encoded as a one hot vector which is tiled appropriately into planes. Appendix F Network Architecture The prediction function pk, vk = fθ(sk) uses the same architecture as AlphaZero: one or two convolutional layers that preserve the resolution but reduce the number of planes, followed by a fully connected layer to the size of the output. For value and reward prediction in Atari we follow [30] in scaling targets using an invertible transform h(x) = sign(x)( p |x| + 1 −1 + ϵx), where ϵ = 0.001 in all our experiments. We then apply a transformation φ to the scalar reward and value targets in order to obtain equivalent categorical representations. We use a discrete support set of size 601 with one support for every integer between −300 and 300.


We use a discrete support set of size 601 with one support for every integer between −300 and 300. Under this transformation, each scalar is represented as the linear combination of its two adjacent supports, such that the original value can be recovered by x = xlow ∗plow + xhigh ∗phigh. As an example, a target of 3.7 would be represented as a weight of 0.3 on the support for 3 and a weight of 0.7 on the support for 4. The value and reward outputs of the network are also modeled using a softmax output of size 601. During inference the actual value and rewards are obtained by ﬁrst computing their expected value under their respective softmax distribution and subsequently by inverting the scaling transformation.


der their respective softmax distribution and subsequently by inverting the scaling transformation. Scaling and transformation of the value and reward happens transparently on the network side and is not visible to the rest of the algorithm. Both the representation and dynamics function use the same architecture as AlphaZero, but with 16 instead of 20 residual blocks [15]. We use 3x3 kernels and 256 hidden planes for each convolution. For Atari, where observations have large spatial resolution, the representation function starts with a sequence of convolutions with stride 2 to reduce the spatial resolution. Speciﬁcally, starting with an input observation of resolution 96x96 and 128 planes (32 history frames of 3 color channels each, concatenated with the corresponding 32 actions broadcast


2 history frames of 3 color channels each, concatenated with the corresponding 32 actions broadcast to planes), we downsample as follows: • 1 convolution with stride 2 and 128 output planes, output resolution 48x48. • 2 residual blocks with 128 planes • 1 convolution with stride 2 and 256 output planes, output resolution 24x24. • 3 residual blocks with 256 planes. • Average pooling with stride 2, output resolution 12x12. • 3 residual blocks with 256 planes. • Average pooling with stride 2, output resolution 6x6. The kernel size is 3x3 for all operations. For the dynamics function (which always operates at the downsampled resolution of 6x6), the action is ﬁrst encoded as an image, then stacked with the hidden state of the previous step along the plane dimension. 14


Appendix G Training During training, the MuZero network is unrolled for K hypothetical steps and aligned to sequences sampled from the trajectories generated by the MCTS actors. Sequences are selected by sampling a state from any game in the replay buffer, then unrolling for K steps from that state. In Atari, samples are drawn according to prioritized replay [35], with priority P(i) = pα i P k pα k , where pi = |νi −zi|, ν is the search value and z the observed n-step return. To correct for sampling bias introduced by the prioritized sampling, we scale the loss using the importance sampling ratio wi = ( 1 N · 1 P (i))β. In all our experiments, we set α = β = 1. For board games, states are sampled uniformly.


1 P (i))β. In all our experiments, we set α = β = 1. For board games, states are sampled uniformly. Each observation ot along the sequence also has a corresponding MCTS policy πt, estimated value νt and environment reward ut. At each unrolled step k the network has a loss to the value, policy and reward target for that step, summed to produce the total loss for the MuZero network (see Equation 1). Note that, in board games without intermediate rewards, we omit the reward prediction loss. For board games, we bootstrap directly to the end of the game, equivalent to predicting the ﬁnal outcome; for Atari we bootstrap for n = 10 steps into the future. To maintain roughly similar magnitude of gradient across different unroll steps, we scale the gradient in two separate locations: • We scale the


dient across different unroll steps, we scale the gradient in two separate locations: • We scale the loss of each head by 1 K , where K is the number of unroll steps. This ensures that the total gradient has similar magnitude irrespective of how many steps we unroll for. • We also scale the gradient at the start of the dynamics function by 1 2. This ensures that the total gradient applied to the dynamics function stays constant. In the experiments reported in this paper, we always unroll for K = 5 steps. For a detailed illustration, see Figure 1. To improve the learning process and bound the activations, we also scale the hidden state to the same range as the action input ([0, 1]): sscaled = s−min(s) max(s)−min(s). All experiments were run using third generation Google Cloud TPUs [12].


d = s−min(s) max(s)−min(s). All experiments were run using third generation Google Cloud TPUs [12]. For each board game, we used 16 TPUs for training and 1000 TPUs for selfplay. For each game in Atari, we used 8 TPUs for training and 32 TPUs for selfplay. The much smaller proportion of TPUs used for acting in Atari is due to the smaller number of simulations per move (50 instead of 800) and the smaller size of the dynamics function compared to the representation function. Appendix H Reanalyze To improve the sample efﬁciency of MuZero we introduced a second variant of the algorithm, MuZero Reanalyze. MuZero Reanalyze revisits its past time-steps and re-executes its search using the latest model parameters, potentially resulting in a better quality policy than the original search.


latest model parameters, potentially resulting in a better quality policy than the original search. This fresh policy is used as the policy target for 80% of updates during MuZero training. Furthermore, a target network [25] ·, v−= fθ−(s0), based on recent parameters θ−, is used to provide a fresher, stable n-step bootstrapped target for the value function, zt = ut+1 + γut+2 + ... + γn−1ut+n + γnv− t+n. In addition, several other hyperparameters were adjusted – primarily to increase sample reuse and avoid overﬁtting of the value function. Speciﬁcally, 2.0 samples were drawn per state, instead of 0.1; the value target was weighted down to 0.25 compared to weights of 1.0 for policy and reward targets; and the n-step return was reduced to n = 5 steps instead of n = 10 steps.


olicy and reward targets; and the n-step return was reduced to n = 5 steps instead of n = 10 steps. Appendix I Evaluation We evaluated the relative strength of MuZero (Figure 2) in board games by measuring the Elo rating of each player. We estimate the probability that player a will defeat player b by a logistic function p(a defeats b) = (1+10(celo(e(b)−e(a))))−1, and estimate the ratings e(·) by Bayesian logistic regression, computed by the BayesElo program [6] using the standard constant celo = 1/400. Elo ratings were computed from the results of a 800 simulations per move tournament between iterations of MuZero during training, and also a baseline player: either Stockﬁsh, Elmo or AlphaZero respectively. Baseline 15


Figure S1: Repeatability of MuZero in Atari for ﬁve games. Total reward is shown on the y-axis, millions of training steps on the x-axis. Dark line indicates median score across 10 separate training runs, light lines indicate individual training runs, and the shaded region indicates 25th to 75th percentile. players used an equivalent search time of 100ms per move. The Elo rating of the baseline players was anchored to publicly available values [39]. In Atari, we computed mean reward over 1000 episodes per game, limited to the standard 30 minutes or 108,000 frames per episode [27], using 50 simulations per move unless indicated otherwise. In order to mitigate the effects of the deterministic nature of the Atari simulator we employed two different evaluation strategies: 30 noop random starts


nature of the Atari simulator we employed two different evaluation strategies: 30 noop random starts and human starts. For the former, at the beginning of each episode a random number of between 0 and 30 noop actions are applied to the simulator before handing control to the agent. For the latter, start positions are sampled from human expert play to initialize the Atari simulator before handing the control to the agent [27]. 16


Game Random Human SimPLe [20] Ape-X [18] R2D2 [21] MuZero MuZero normalized alien 227.75 7,127.80 616.90 40,805.00 229,496.90 741,812.63 10,747.5 % amidar 5.77 1,719.53 74.30 8,659.00 29,321.40 28,634.39 1,670.5 % assault 222.39 742.00 527.20 24,559.00 108,197.00 143,972.03 27,664.9 % asterix 210.00 8,503.33 1,128.30 313,305.00 999,153.30 998,425.00 12,036.4 % asteroids 719.10 47,388.67 793.60 155,495.00 357,867.70 678,558.64 1,452.4 % atlantis 12,850.00 29,028.13 20,992.50 944,498.00 1,620,764.00 1,674,767.20 10,272.6 % bank heist 14.20 753.13 34.20 1,716.00 24,235.90 1,278.98 171.2 % battle zone 2,360.00 37,187.50 4,031.20 98,895.00 751,880.00 848,623.00 2,429.9 % beam rider 363.88 16,926.53 621.60 63,305.00 188,257.40 454,993.53 2,744.9 % berzerk 123.65 2,630.42 57,197.00 53,318.70 85,9


53 621.60 63,305.00 188,257.40 454,993.53 2,744.9 % berzerk 123.65 2,630.42 57,197.00 53,318.70 85,932.60 3,423.1 % bowling 23.11 160.73 30.00 18.00 219.50 260.13 172.2 % boxing 0.05 12.06 7.80 100.00 98.50 100.00 832.2 % breakout 1.72 30.47 16.40 801.00 837.70 864.00 2,999.2 % centipede 2,090.87 12,017.04 12,974.00 599,140.30 1,159,049.27 11,655.6 % chopper command 811.00 7,387.80 979.40 721,851.00 986,652.00 991,039.70 15,056.4 % crazy climber 10,780.50 35,829.41 62,583.60 320,426.00 366,690.70 458,315.40 1,786.6 % defender 2,874.50 18,688.89 411,944.00 665,792.00 839,642.95 5,291.2 % demon attack 152.07 1,971.00 208.10 133,086.00 140,002.30 143,964.26 7,906.4 % double dunk -18.55 -16.40 24.00 23.70 23.94 1,976.3 % enduro 0.00 860.53 2,177.00 2,372.70 2,382.44 276.9 % ﬁshing derby -91.71


4.00 23.70 23.94 1,976.3 % enduro 0.00 860.53 2,177.00 2,372.70 2,382.44 276.9 % ﬁshing derby -91.71 -38.80 -90.70 44.00 85.80 91.16 345.6 % freeway 0.01 29.60 16.70 34.00 32.50 33.03 111.6 % frostbite 65.20 4,334.67 236.90 9,329.00 315,456.40 631,378.53 14,786.7 % gopher 257.60 2,412.50 596.80 120,501.00 124,776.30 130,345.58 6,036.8 % gravitar 173.00 3,351.43 173.40 1,599.00 15,680.70 6,682.70 204.8 % hero 1,026.97 30,826.38 2,656.60 31,656.00 39,537.10 49,244.11 161.8 % ice hockey -11.15 0.88 -11.60 33.00 79.30 67.04 650.0 % jamesbond 29.00 302.80 100.50 21,323.00 25,354.00 41,063.25 14,986.9 % kangaroo 52.00 3,035.00 51.20 1,416.00 14,130.70 16,763.60 560.2 % krull 1,598.05 2,665.53 2,204.80 11,741.00 218,448.10 269,358.27 25,083.4 % kung fu master 258.50 22,736.25 14,862.50 97,830.00


4.80 11,741.00 218,448.10 269,358.27 25,083.4 % kung fu master 258.50 22,736.25 14,862.50 97,830.00 233,413.30 204,824.00 910.1 % montezuma revenge 0.00 4,753.33 2,500.00 2,061.30 0.00 0.0 % ms pacman 307.30 6,951.60 1,480.00 11,255.00 42,281.70 243,401.10 3,658.7 % name this game 2,292.35 8,049.00 2,420.70 25,783.00 58,182.70 157,177.85 2,690.5 % phoenix 761.40 7,242.60 224,491.00 864,020.00 955,137.84 14,725.3 % pitfall -229.44 6,463.69 -1.00 0.00 0.00 3.4 % pong -20.71 14.59 12.80 21.00 21.00 21.00 118.2 % private eye 24.94 69,571.27 35.00 50.00 5,322.70 15,299.98 22.0 % qbert 163.88 13,455.00 1,288.80 302,391.00 408,850.00 72,276.00 542.6 % riverraid 1,338.50 17,118.00 1,957.80 63,864.00 45,632.10 323,417.18 2,041.1 % road runner 11.50 7,845.00 5,640.60 222,235.00 599,246.70 613,411.80


45,632.10 323,417.18 2,041.1 % road runner 11.50 7,845.00 5,640.60 222,235.00 599,246.70 613,411.80 7,830.5 % robotank 2.16 11.94 74.00 100.40 131.13 1,318.7 % seaquest 68.40 42,054.71 683.30 392,952.00 999,996.70 999,976.52 2,381.5 % skiing -17,098.09 -4,336.93 -10,790.00 -30,021.70 -29,968.36 -100.9 % solaris 1,236.30 12,326.67 2,893.00 3,787.20 56.62 -10.6 % space invaders 148.03 1,668.67 54,681.00 43,223.40 74,335.30 4,878.7 % star gunner 664.00 10,250.00 434,343.00 717,344.00 549,271.70 5,723.0 % surround -9.99 6.53 7.00 9.90 9.99 120.9 % tennis -23.84 -8.27 24.00 -0.10 0.00 153.1 % time pilot 3,568.00 5,229.10 87,085.00 445,377.30 476,763.90 28,486.9 % tutankham 11.43 167.59 273.00 395.30 491.48 307.4 % up n down 533.40 11,693.23 3,350.30 401,884.00 589,226.90 715,545.61 6,407.0 % v


5.30 491.48 307.4 % up n down 533.40 11,693.23 3,350.30 401,884.00 589,226.90 715,545.61 6,407.0 % venture 0.00 1,187.50 1,813.00 1,970.70 0.40 0.0 % video pinball 0.00 17,667.90 565,163.00 999,383.20 981,791.88 5,556.9 % wizard of wor 563.50 4,756.52 46,204.00 144,362.70 197,126.00 4,687.9 % yars revenge 3,092.91 54,576.93 5,664.30 148,595.00 995,048.40 553,311.46 1,068.7 % zaxxon 32.50 9,173.30 42,286.00 224,910.70 725,853.90 7,940.5 % # best 0 5 0 5 13 37 Table S1: Evaluation of MuZero in Atari for individual games with 30 random no-op starts. Best result for each game highlighted in bold. Each episode is limited to a maximum of 30 minutes of game time (108k frames). SimPLe was only evaluated on 36 of the 57 games, unavailable results are indicated with ‘-’.


ames). SimPLe was only evaluated on 36 of the 57 games, unavailable results are indicated with ‘-’. Human normalized score is calculated as snormalized = sagent−srandom shuman−srandom . 17


Game Random Human Ape-X [18] MuZero MuZero normalized alien 128.30 6,371.30 17,732.00 713,387.37 11,424.9 % amidar 11.79 1,540.43 1,047.00 26,638.80 1,741.9 % assault 166.95 628.89 24,405.00 143,900.58 31,115.2 % asterix 164.50 7,536.00 283,180.00 985,801.95 13,370.9 % asteroids 877.10 36,517.30 117,303.00 606,971.12 1,700.6 % atlantis 13,463.00 26,575.00 918,715.00 1,653,202.50 12,505.6 % bank heist 21.70 644.50 1,201.00 962.11 151.0 % battle zone 3,560.00 33,030.00 92,275.00 791,387.00 2,673.3 % beam rider 254.56 14,961.02 72,234.00 419,460.76 2,850.5 % berzerk 196.10 2,237.50 55,599.00 87,308.60 4,267.3 % bowling 35.16 146.46 30.00 194.03 142.7 % boxing -1.46 9.61 81.00 56.60 524.5 % breakout 1.77 27.86 757.00 849.59 3,249.6 % centipede 1,925.45 10,321.89 5,712.00 1,138,294.60 13,533.9


kout 1.77 27.86 757.00 849.59 3,249.6 % centipede 1,925.45 10,321.89 5,712.00 1,138,294.60 13,533.9 % chopper command 644.00 8,930.00 576,602.00 932,370.10 11,244.6 % crazy climber 9,337.00 32,667.00 263,954.00 412,213.90 1,726.9 % defender 1,965.50 14,296.00 399,865.00 823,636.00 6,663.7 % demon attack 208.25 3,442.85 133,002.00 143,858.05 4,441.0 % double dunk -15.97 -14.37 22.00 23.12 2,443.1 % enduro -81.84 740.17 2,042.00 2,264.20 285.4 % ﬁshing derby -77.11 5.09 22.00 57.45 163.7 % freeway 0.17 25.61 29.00 28.38 110.9 % frostbite 90.80 4,202.80 6,512.00 613,944.04 14,928.3 % gopher 250.00 2,311.00 121,168.00 129,218.68 6,257.6 % gravitar 245.50 3,116.00 662.00 3,390.65 109.6 % hero 1,580.30 25,839.40 26,345.00 44,129.55 175.4 % ice hockey -9.67 0.53 24.00 52.40 608.5 % jamesbond 33.5


80.30 25,839.40 26,345.00 44,129.55 175.4 % ice hockey -9.67 0.53 24.00 52.40 608.5 % jamesbond 33.50 368.50 18,992.00 39,107.20 11,663.8 % kangaroo 100.00 2,739.00 578.00 13,210.50 496.8 % krull 1,151.90 2,109.10 8,592.00 257,706.70 26,802.6 % kung fu master 304.00 20,786.80 72,068.00 174,623.60 851.1 % montezuma revenge 25.00 4,182.00 1,079.00 57.10 0.8 % ms pacman 197.80 15,375.05 6,135.00 230,650.24 1,518.4 % name this game 1,747.80 6,796.00 23,830.00 152,723.62 2,990.7 % phoenix 1,134.40 6,686.20 188,789.00 943,255.07 16,969.6 % pitfall -348.80 5,998.91 -273.00 -801.10 -7.1 % pong -17.95 15.46 19.00 19.20 111.2 % private eye 662.78 64,169.07 865.00 5,067.59 6.9 % qbert 159.38 12,085.00 380,152.00 39,302.10 328.2 % riverraid 588.30 14,382.20 49,983.00 315,353.33 2,281.9 % road runner 2


380,152.00 39,302.10 328.2 % riverraid 588.30 14,382.20 49,983.00 315,353.33 2,281.9 % road runner 200.00 6,878.00 127,112.00 580,445.00 8,688.9 % robotank 2.42 8.94 69.00 128.80 1,938.3 % seaquest 215.50 40,425.80 377,180.00 997,601.01 2,480.4 % skiing -15,287.35 -3,686.58 -11,359.00 -29,400.75 -121.7 % solaris 2,047.20 11,032.60 3,116.00 2,108.08 0.7 % space invaders 182.55 1,464.90 50,699.00 57,450.41 4,465.9 % star gunner 697.00 9,528.00 432,958.00 539,342.70 6,099.5 % surround -9.72 5.37 6.00 8.46 120.5 % tennis -21.43 -6.69 23.00 -2.30 129.8 % time pilot 3,273.00 5,650.00 71,543.00 405,829.30 16,935.5 % tutankham 12.74 138.30 128.00 351.76 270.0 % up n down 707.20 9,896.10 347,912.00 607,807.85 6,606.9 % venture 18.00 1,039.00 936.00 21.10 0.3 % video pinball 0.00 15,641.09 873,989.0


07,807.85 6,606.9 % venture 18.00 1,039.00 936.00 21.10 0.3 % video pinball 0.00 15,641.09 873,989.00 970,881.10 6,207.2 % wizard of wor 804.00 4,556.00 46,897.00 196,279.20 5,209.9 % yars revenge 1,476.88 47,135.17 131,701.00 888,633.84 1,943.0 % zaxxon 475.00 8,443.00 37,672.00 592,238.70 7,426.8 % # best 0 6 5 46 Table S2: Evaluation of MuZero in Atari for individual games from human start positions. Best result for each game highlighted in bold. Each episode is limited to a maximum of 30 minutes of game time (108k frames). 18


Model s0 = hθ(o1, ..., ot) rk, sk = gθ(sk−1, ak) pk, vk = fθ(sk)   pk, vk, rk = µθ(o1, ..., ot, a1, ..., ak) Search νt, πt = MCTS(s0 t, µθ) at ∼πt Learning Rule pk t , vk t , rk t = µθ(o1, ..., ot, at+1, ..., at+k) zt =  uT for games ut+1 + γut+2 + ... + γn−1ut+n + γnνt+n for general MDPs lt(θ) = K X k=0 lr(ut+k, rk t ) + lv(zt+k, vk t ) + lp(πt+k, pk t ) + c||θ||2 Losses lr(u, r) =  0 for games φ(u)T log r for general MDPs lv(z, q) =  (z −q)2 for games φ(z)T log q for general MDPs lp(π, p) = πT log p Figure S2: Equations summarising the MuZero algorithm. Here, φ(x) refers to the representation of a real number x through a linear combination of its adjacent integers, as described in the Network Architecture section. 19


Figure S3: Details of MuZero evaluations (A-B) and policy improvement ablations (C-D). (A-B) Distribution of evaluation depth in the search tree for the learned model for the evaluations in Figure 3A-B. The network was trained over 5 hypothetical steps, as indicated by the red line. Dark blue line indicates median depth from the root, dark shaded region shows 25th to 75th percentile, light shaded region shows 5th to 95th percentile. (C) Policy improvement in Ms. Pacman - a single network was trained at 50 simulations per search and is evaluated at different numbers of simulations per search, including playing according to the argmax of the raw policy network. The policy improvement effect of the search over the raw policy network is clearly visible throughout training.


mprovement effect of the search over the raw policy network is clearly visible throughout training. This consistent gap between the performance with and without search highlights the policy improvement that MuZero exploits, by continually updating towards the improved policy, to efﬁciently progress towards the optimal policy. (D) Policy improvement in Go - a single network was trained at 800 simulations per search and is evaluated at different numbers of simulations per search. In Go, the playing strength improvement from longer searches is much larger than in Ms. Pacman and persists throughout training, consistent with previous results in [40]. This suggests, as might intuitively be expected, that the beneﬁt of models is greatest in precision planning domains. 20


Figure S4: Learning curves of MuZero in Atari for individual games. Total reward is shown on the y-axis, millions of training steps on the x-axis. Line indicates mean score across 1000 evaluation games, shaded region indicates standard deviation. 21


Decision Transformer: Reinforcement Learning via Sequence Modeling Lili Chen∗,1, Kevin Lu∗,1, Aravind Rajeswaran2, Kimin Lee1, Aditya Grover2, Michael Laskin1, Pieter Abbeel1, Aravind Srinivas†,1, Igor Mordatch†,3 ∗equal contribution †equal advising 1UC Berkeley 2Facebook AI Research 3Google Brain {lilichen, kzl}@berkeley.edu Abstract We introduce a framework that abstracts Reinforcement Learning (RL) as a sequence modeling problem. This allows us to draw upon the simplicity and scalability of the Transformer architecture, and associated advances in language modeling such as GPT-x and BERT. In particular, we present Decision Transformer, an architecture that casts the problem of RL as conditional sequence modeling.


ecision Transformer, an architecture that casts the problem of RL as conditional sequence modeling. Unlike prior approaches to RL that ﬁt value functions or compute policy gradients, Decision Transformer simply outputs the optimal actions by leveraging a causally masked Transformer. By conditioning an autoregressive model on the desired return (reward), past states, and actions, our Decision Transformer model can generate future actions that achieve the desired return. Despite its simplicity, Decision Transformer matches or exceeds the performance of state-of-the-art model-free ofﬂine RL baselines on Atari, OpenAI Gym, and Key-to-Door tasks. R s a a s a a causal transformer emb. + pos. enc. linear decoder . . . 21 return state action . . .


R s a a s a a causal transformer emb. + pos. enc. linear decoder . . . 21 return state action . . . ^ R ^ Figure 1: Decision Transformer architecture1. States, actions, and returns are fed into modalityspeciﬁc linear embeddings and a positional episodic timestep encoding is added. Tokens are fed into a GPT architecture which predicts actions autoregressively using a causal self-attention mask. 1Our code is available at: https://sites.google.com/berkeley.edu/decision-transformer arXiv:2106.01345v2 [cs.LG] 24 Jun 2021


Contents 1 Introduction 3 2 Preliminaries 4 2.1 Ofﬂine reinforcement learning . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4 2.2 Transformers . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4 3 Method 4 4 Evaluations on Ofﬂine RL Benchmarks 6 4.1 Atari . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 6 4.2 OpenAI Gym . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 7 5 Discussion 8 5.1 Does Decision Transformer perform behavior cloning on a subset of the data? . . . 8 5.2 How well does Decision Transformer model the distribution of returns? . . . . . . 8 5.3 What is the beneﬁt of using a longer context length? . . . . . . . . . . . . . . . . .