Dream and Search to Control: Latent Space Planning for Continuous Control
Authors: Anurag Koul
Varun Kumar Vijay
Alan Fern
Somdeb Majumdar
Venue:  (2021)
Keywords: Reinforcement Learning
Model Based RL
Continuous Control
Search
Planning
MCTS

Abstract: Learning and planning with latent space dynamics has been shown to be useful for sample efficiency in model-based reinforcement learning (MBRL) for discrete and continuous control tasks. In particular, recent work, for discrete action spaces, demonstrated the effectiveness of latent-space planning via Monte-Carlo Tree Search (MCTS) for bootstrapping MBRL during learning and at test time. However, the potential gains from latent-space tree search have not yet been demonstrated for environments with continuous action spaces.  In this work, we propose and explore an MBRL approach for continuous action spaces based on tree-based planning over learned latent dynamics.  We show that it is possible to demonstrate the types of bootstrapping benefits as previously shown for discrete spaces. In particular, the approach achieves improved sample efficiency and performance on a majority of challenging continuous-control benchmarks compared to the state-of-the-art. 

Under review as a conference paper at ICLR 2021
DREAM AND SEARCH TO CONTROL: LATENT SPACE PLANNING FOR CONTINUOUS CONTROL
Anonymous authors Paper under double-blind review
ABSTRACT
Learning and planning with latent space dynamics has been shown be useful for sample efficiency in model-based reinforcement learning (MBRL) for dis- crete and continuous control tasks. In particular, recent work, for discrete ac- tion spaces, demonstrated the effectiveness of latent-space planning via Monte- Carlo Tree Search (MCTS) for bootstrapping MBRL during learning and at test time. However, the potential gains from latent-space tree search have not yet been demonstrated for environments with continuous action spaces. this work, we propose and explore an MBRI. approach for continuous action spaces based on tree-based planning over learned latent dynamics. We show that it is possible to demonstrate the types of bootstrapping benefits as previously shown for discrete spaces. In particular, the approach achieves improved sample efficiency and per- formance majority of challenging continuous-contro benchmarks compared to the state-of-the-arr.
INTRODUCTION
Deep reinforcement learning (RL) has been effective at solving sequential decision-making prob- lems with varying levels of difficulty. The solutions generally fall into one of two categories: model-free and model-based methods. Model-free methods (Haarnoja et al. 2018 Silver et al.] 2014 Lillicrap et al., 2015] Fujimoto et 2018 22018/Sccullii
[Schulman et al. 2017) directly learn a policy or action-values. However, they are usually considered to be cample-inefficien
Model-base based meth-
ods (Lee et al. 2019 Gregor et al. 2019 Zhang et al. 2018 Ha & Schmidhuber] 2018) learn the environment's s dynamies. Common model- based approaches involve sampling trajectories from the learned dynamics to train using RL applying planning algorithm directly on the Iearned dynam- ics (Ha & Schmidhuber 2018 Halner et al. 20]9a
However, learning multi-step dynamics in the raw observation space is challenging. This is primarily because it involves the reconstruction of nigh-dimenssonal features (e.g., pixels) in order to roll-out trajectories, which is error proo-ppoon process. Instead, recent work has focused on learning latent space models (Hafner et al.) 2019a[b). This has been shown to improve robustness and sample efficiency by eliminating the need for high-dimensional reconstruction during inference.
Learning on latent dynamics: Once the dynamics have been learned, a classic approach to sample trajectories from to learn a policy using RL. This approach is usually motivated by sample efficiency. Dreamer (Hafner et al. 2019a) took this approach and demonstrated state-of-the-art performance on continuous control by performing gradient-based RL on learned latent dynamics. Another approach is to perform look-ahead search where the dynamics are used for a multi-step rollout determine an optimal action. This could be accompanied by a value estimate and/or policy that produces state-action mappings to narrow the search space or reduce the search's 's depth. MuZero (Schrittwieser al. 2019) took this approach and applied tree-based search on latent dynamics however, it restricted to discrete action spaces only. The role of look-ahead search using learned latent dynamics has not been explored sufficiently for continuous action spaces.
Our contribution: In this paper, extend the idea of performing look-ahead search using learned latent dynamics to continuous action spaces. Our high level approach is shown in Fig I We build on top of Dreamer and modify how actions are sampled during online planning. Instead of sampling actions from the current policy, we search over a set of actions sampled from a mix distributions. For our search mechanism, implement MCTS but also investigate simple rollout algorithm

---

Under review as a conference paper at ICLR 2021
Environment interaction
action
transition data
Replay Buffer
sample
rollouts
Learn Model
imagine
rollouts
Agent
Learn Policy
& Value
obs, reward
updated model/agent
Figure 1: Overall Approach. Phase 1: At each time-step, determine the actions be stepped via look-ahead search over latent dynamics and write the resulting transition data in a replay buffer. Phase 2: Sample fixed-horizon rollouts from replay buffer. Phase 3: Optimize latent model to reconstruct rewards and observations. Phase 4: Sample imaginary rollouts from the latent model to train agent.
that trades off performance for compute complexity. We observed that look-ahead search results in optimal actions early on. This, in turn, leads to faster convergence of model estimates and of optimal policies. However,these benefits come the cost of computational time since deeper and iterative look-ahead search takes more time.
2 RELATED WORK
Model-Based RL involves learning a dynamical model and utilizing it to sample trajectories to train policy or performing look-ahead search for policy improvement. World Model (Ha & Schmidhu- 2018) learns a latent representation of observations and learns dynamical model over this latent representation. It then utilizes the learned model to optimize a linear controller via evolution. The latent representations help to plan in low dimensional abstract space. PlaNet (Hafner et al. (019b) does end-to-end training of the latent observation encoding and dynamics and does online planning using CEM. Dreamer, the current state-of-the-art (Hafner et al.] (209a) builds on top of PlaNet and uses analytic gradients to efficiently learn long-horizon behaviors for visual control purely by latent imagination. We built our work top of Dreamer.
Planning using look-ahead search: Given dynamics model, one can search over its state-action space to determine optimal action. If one has perfect model but sub-optimal action-valus function, then in fact, a deeper look-ahead search will usually yield better policies. This look-ahead search could be done both during online decision making and during policy improvement. AlphaGo (Silver et al., [017]] [ Lowrey et al. 2018) combined MCTS with a value estimator to plan and explore with known ground-trute dynamics. MuZero (Schrittwieser al. 2019) learned latent dynamic model and performed MCTS over it. They conditioned the latent representation on value and multi- step reward prediction rather than observation reconstruction. These works are limited to discrete action space due to large action space in continuous action space tasks.
variation on MCTS is Progressive Widening (Coulom 2006 Chaslot al, 2008 Couêtoux al., 2011) where the the child-actions of node are expanded based on its visitation count. This has been exploited continuous action spaces by adding sampled actions from a proposal distribution to create a discrete set of child-actions. AOC (Moerland et al.) 2018) utilized this approach to apply MCTS in continuous action spaces over true dynamics.However. they showed their results only for the Pendulum-vO task in OpenAI Gym. A common thread across all these prior works is the availability of known ground-trure -rrth dynamics. Another approach that can be used to reduce the size of the lhek-ahead search space is to utilize hierarchical optimistic optimization (HOO), which splits the action space and gradually increases the resolution actions (Mansley et al. 2011 Bubeck et al. 2009).
2

---

Under review as a conference paper at ICLR 2021
3 DREAMER: LEARNING LATENT DYNAMICS AND POLICY
MBRL approaches usually involve collecting data from the environment, dynamics learning, and behavior learning by utilizing data collected from environment-agent interaction and/or trajectories generated via learned dynamics. We adopt Dreamer's s dynamics and behavior learning approach. And, unlike Dreamer, we perform decision-tin -immme planning via look-ahead search during exploration to determine action, instead of just using learned behavior policy. In this section, we review Dreamer's dynamics learning component in subsection (3.1) and the behavior learning component in subsection (3.2 After that, we describe investigation of look-ahead search for decision-time planning in section
3.1 DYNAMICS LEARNING
This comprises of a representation model that encodes the history of observations and actions to create a Markovian belief of the current state. This belief state and an action are consumed by transition model that predicts the model's future latent state. The transition model is conditioned on observation model and reward model. The observation model predicts the observation corresponding to the current latent state and the reward model predicts the corresponding reward. The observation model is used only to condition the transition model and does not play role during inference. These models are summarized n Eq. where (s,, a, Or, Ft)) represent latent state, action, observation, and reward at discrete-time step I.
Representation model:
(e(s11st-1, (o(0))st)
Observation model:
(1)
Reward model:
(r((())))
Transition model:
qo 8t 0-11)
The transition model is implemented recurrent state-space model (RSSM) (Hafner et al.
2019b).
The representation, observation model and reward model are dense networks. Observation features are given the representation model, which are reconstructed by the observation model. This is
unlike Dreamer which trains the model over raw pixels. The combined parameter vector in Eq.
D is updated by stochastic backpropagation (Kingma Welling, 2013] Rezende al.] 2014).
optimizes 8 using joint reconstruction loss (JRec) comprising of an observation reconstruction loss (Jo), reward reconstruction loss (JR) and KL regularizer (Jn), as described in Eq. (2)
JAEC = Ep E (J+ (Th)TT)) + const
To = Inq((r(s))
(2)
JA = Ing(res)s
JB - -3KL ( (11))
3.2 BEHAVIOR LEARNING
Given a transition dynamics, the action and value models are learnt from the imagined trajectories [+H (Sr; Or, )(+H of finite-horizon (H) over latent state space. Dreamer utilizes these trajectories for learning behaviour via an actor-critic approach. The action model implements the policy and aims to predict actions that solve the imagination environment. The value model estimates the expected imagined rewards that the action model achieves from each state Sy, where the discrete-time index during imagination.
Action model:
a, ~ q0(a)|s7) (x(s)) = Eg(( D 7. -T-
Value model:
As in Eq. 33. action and value models implemented dense neural network with parameters and th respectively. As shown in Eq. [ the action model outputs a canh-transformed Gaus- sian (Haarnoja al. 2018). This allows for re-parameterized sampling (Kingma & Welling 2013
3

---

Under review as a conference paper at ICLR 2021
Rezende et al. 2014) that views sampled actions deterninistically dependent on the neural net- work output, allowing Dreamer to backpropagate analytic gradients through the sampling operation.
a, = tanh((u + (5+))) E ~~ -NoNoalm. 1).
(4)
Dreamer uses VA (Bq. 8 Suttonn& Barto, 2018), an exponentially weighted average of the estimates for different k-step returns for the imagination rollouts. The k varies from .... H.
It helps to balance bias and variance of target for value estimates by acting as an intermediary between the 1-step and the Monte-Carlo return. The objective for the action model qs (a, (s,) is to predict actions that result in state trajectories with high value estimates as shown in Eq. 5 It uses analytic gradients through the learned dynamics to maximize the value estimates. The objective for the value model (w((s), in turn, is to regress the value estimates as shown in Eq. 6 while learning behaviors.
Also, Dreamer freezes the world model
H+H
max P9g-90
(5)
min E Rag Σ
(o)(T)) va(++))
(6)
h-1
(*++)) = Ego, 1o
7"th n- Ta
ra + - (m0( with 4= min - kk kk H),
(7)
n=t
H-1
Va(++) = (1
(8)
n=1
OUR APPROACH: DECISION-TIME PLANNING IN LATENT SPACES
Planning could be used to improve policy in two ways (uttonn & Barto 2018 1) Background planning involves optimizing a policy or value estimates with targets received from search or dy- optimal action to be stepped into the environment via search after encountering each new state. The simplest case would be to determine action based on current policy or value estimates. An alternate approach would be to perform look-ahead search with the dynamics for determining an improved action.
namic programming. 2) Decision-time planning involves determining
In this work, we only focus decision-time planning. We hypothesize that when we have imperfect valio estimates during training, we could improve on that by look-ahead search in updated dynamics, leading to better action selection. These better actions lead to less noisy data in the replay buffer which leads to a better model and, in turn, a better policy. Under an oracle model, this is guaranteed to give better policy than the current proposal policy. We implement two variations of decision-time planning: rollout and MCTS.
4.1 ROLLOUT
Rollout is based Monte-Carlo Control and involves rolling trajectories from given state and action pair for a certain depth (or termination) in order to determine an action-value estimate. We sample actions based on the current policy as well some other distribution (such uni- form/random) for better exploration from a given state. For each of the sampled actions, we perform rollout using current stochastic policy and estimate V(A) in Eq. (8). A greedy action that max- imizes the action- value estimates is selected and some noise is added to the greedy action during exploration. The computation time needed by rollout algorithm depends on the number of actions that have to be evaluated and the depth of the simulated trajectories. Upon action selection and en- vironment interaction, we get a new observation based on which we re-estimate our belief state and re-plan for the next time-step.
Rollout is related to random shooting (Nagabandi et al. 2018) which generates multiple action se- quences, computes a reward for each imagined trajectory and chooses the first action of the trajectory with the highest cumulative reward. It is also similar to Cross- Pross-Entropy method (CEM), which is an iterative procedure of sampling action sequences and fitting new distribution top-k action quences with the highest rewards. The main distinction of Rollout with these is that the actions are already sampled from a proposal distribution which helps narrow action search space.
4

---

Under review as a conference paper at ICLR 2021
ot at-1
81-1
Vi
(1) ~ de
(a) Dreamer
(b) Rollout
(c) MCTS
Figure 2: Decision-time planning comparison: For a current state Dreamer simply samples an action using the current policy. b) Rollout samples n actions and unrolls the dynamics for each action transition for depth using mode action of current stochastic policy. MCTS iteratively extends a search tree via trade off over exploration and exploitation. In both (b) and (c), rollouts used to update action-value estimates of the root edges (actions) and Value estimates of leaf nodes are used for bootstrapping. The blue line indicates path of backpropagted return for single rollout/simulation
4.2 MCTS
Unlike the Rollout algorithm which follows current policy for rollout, MCTS iteratively extends search tree starting from root node representing the current environment state. In this section, we describe how adapt MCTS for continuous action spaces.
Every node of the search tree associated with an internal state 8. For each action a from 8 there is edge (s, that stores set of statistics ((s,, a), Q(s, a), ass,a), representing visit counts N, mean value Q, policy P, reward Rt and state transition S. Starting from the root node, MCTS executes the following three steps for given computational budget:
S(s, a), respectively
Selection:
Starting from a root node 8 we traverse through the tree until we reach a leaf node s At each time-step of traversal, an action a is selected at each node. This is done by computing upper confidence bound (UCB) score (Rosin [001]] Silver et al.) 2018) for each child action(edge) using the stored internal statistics. The selected action is the one that maximizes this UCB.
Continuous action. In continuous action spaces, we cannot enumerate all actions since there are infinitely many actions in continuous space. alternative could be have fixed number of actions for each node, where each action is sampled from distribution. This fixed number could be pre-defined or determined via hyper-parameter search. This makes it similar to discrete-action MCTS.
Another alternative is to use progressive widening, where we slowly grow the number of child actions of state in the tree as function of the total number of visits to that state {n(s) EB children(S) (s, b)}. Thus, if children(s) < (()) then, we add a new action to the parent node and expand that new action node for simulation. Otherwise we simply select one of the actions from the eurrent set of actions using upper confidence bound. This was also done by (Moerland et al. 2018).
V2,
Σg N2s 1)
+
at arg max Q( P(s, 1 (s)
c1 log
C2
The constants C1 and C2 used to control the influence of the prior P(s, a) relative the value Qss, a) as nodes are visited more often.
Expansion: Upon reaching un-expanded edge, compute the leaf state s' = po(ss-11,0 a') and it's reward o ro(' and a node corresponding to is added to the tree. In the case of fixed actions set for each node, sample actions and add them edges to the node. Whereas, in the case of
5

---

Under review as a conference paper at ICLR 2021
progressive widening, edges are added as they will be introduced upon re- re-selection during next simulation.
Backup: At the end of the simulation, the statistics along the trajectory updated. The backup is generalized to the case where the environment can emit intermediate rewards, have a discount different from and the value estimates unbounded. For k I...0, form an i-k step estimate of the cumulative discounted reward, bootstrapping from the value function for as shown in Eq. 10
1-1-k
G G==
Σ Tk+14 + gl
(10)
For I...,, update the statisties for each edge (st-1, ) in the simulation path
Q(sk-1,at) N(sh-1 aaa) 1(st-((((ss)) ,a*) Gk N(s , = -N a aa)) (1I)
N(s--1,at) +1
During action selection, we use normalized Q value estimates Q [0, 1] in the pUCT rule by using the eq. 12 as done by MuZero. This normalization is done by the menimum-maximan values observed so far in the tree.
Q(s a²) ming Q(s,a)
Q(s ,aa)=
(12)
maTeee Q(8, a) aeTree Q(s,a)
5 EXPERIMENTS
In our experiments, we aim to determine the importance of search over latent space during explo- ration and, if that helps in policy improvement and better sample efficiency for control tasks, based on our hypothesis in section (4).
Tasks. We evaluated 20 control tasks from the DeepMind Control suite (Tassa al. 2020). which are implemented in Mujoco. Each task's episode involves 1000 steps, and each step reward is in the range [0, 1] We use an action repeat of 2 across all tasks, as done by Dreamer. All episodes have random initial states. In all the considered tasks, we use the default feature representation of the system state, consisting of information such as joint positions and velocities, as well as additional sensor values and target positions where provided by the environment.
Implementation. We build our work on top of the existing Dreamer architecture and optimization with the change in way actions are selected during exploration. For the rollout case, sample 100 actions from the proposal distribution and 50 actions from a uniform random distribution and performed
perform rollout for 10 steps in latent space. These rollouts giving us the freedom of choosing a large number of actions. The choice of large action count was motivated from (Van de Wiele al.) 2020)). In the case of MCTS, we performed 50 simulations, and all child actions were sampled from the proposal policy distribution. This choice for the number of simulations was adopted from MuZero. We restricted ourselves from hyper-parameter search over range of simulations counts due to computational budget. In Dreamer, we added exploration noise from Normal(O, 0.3) the mode action of the current stochastic policy during exploration. Furthermore, for rollout and mets, the same exploration noise was added to all sampled actions as well as selected greedy action during exploration. Also, for MCTS, we used uniform prior for each edge and set = 1.25 C2 19652, Cpuct = 0.05, Cpw 1 and
batch operation
0.5
Performance . We treat Dreamer as our baseline, which already improves over model-free methods like DDPG and A3C in terms of sample efficiency shared in (Hafner et al. 2019a). In the original paper, Dreamer operated directly on the pixel space. In this work, we adapt Dreamer to operate on the feature space provided by the environment. We summarize training curves for 2 x 10 steps in the figure (5) and show the mean performance across different environment interaction modes in figure 3ap. Due to computational budget, we limited our experiments only 2 million environment

---

Under review as a conference paper at ICLR 2021
steps to have comparable comparison with MCTS experiments which require much longer training time as shown in figure (3b) due to synchronous nature training.
In general, observe that by incorporating look-ahead search during exploration, sample efficiency is improves. Furthermore, we also observe that the variance learned policies' behavior is smaller than the baseline. It further improves look-ahead search. Also, In some cases, such as Hopper, we do observe that both search methods are slightly sample-ineffccient compared the baseline. One potential reason is that in the initial phase of training the reward and value estimates' error may be high, which accumulates over look-ahead search. This could lead to bad updated estimates of aation-valles during decision-timg time planning, thereby it diverges from good actions. In both the cases of Hopper, search improves with training and eventually surpasses the baseline.
replace the rollout search with the more complex MCTS
1000
300
M
episode reward
800
tmm(hhuur)
600
200
400
dreamer
100
200
dreamer rollout
dreamer
mcts
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
env. steps
1e6
env. steps
1e6
(a) Data Consumption
(b) Time Consumption
Figure 3: Mean performance across 20 Benchmark tasks of dm-control suite. Each task run for 3 random seeds. b) Time-taken by each explore method. MCTS is significantly slower than rollout due the increased complexity of its search.
Notably, in most cases, we also observe that the performance of Rollout's matches that of MCTS. This implies that complex and time- consuming look-ahead search may not be ways required in the considered dm-control tasks, after a proposal policy has been learned. This is interesting trade-off MCTS takes computationally much longer than the rollouts which are performed batch operation.
1000
800
episode reward
600
400
While the above version of MCTS was im-
200
fixed mots: progressive
plemented with progressive widening, we also
investigated classic approach of simply hav-
ing fixed children for each node. We consid- ered 20 children for each node. We evalu- ated this implementation for environments and show the mean performance across these tasks fig. 4b. The environments consid- ered were Acrobat Swing Up. Cheetah Run, Reacher Hard, QuadrupedRum, Finger Spin, Hopper Hop. Walker Stand and Walk. Both the variants were given budget of 50 simulations. across both the methods.
0.0
0.5
1.0
1.5
env. steps
1e6
Figure 4: Comparison of mets varitions progres- sive widening and fixed children. We share mean scores across 8 environments trained on 3 random seeds.
We observe performance to be almost equivalent

---

Under review as a conference paper at ICLR 2021
6 SUMMARY & FUTURE WORK
In this work, we investigated the role of decision-time planning with the tree-based search over a learned latent dynamics for continuous control tasks during exploration. We show this leads to better sample efficiency in many tasks compared to the state-of-the-art Dreamer. The work can benefit from more complex tree-based search or better action selection expansion in tree. Also, we observe that the search component significantly increases the training time. As part of future work, we will explore the inclusion of distributional training for reducing computation time while still maintaining sample efficiency. Also, utilizing the output of the search as target estimates during optimization is another potential future thread. Finally, we will extend our investigation over more challenging tasks in terms of action-space, observation space, and complex goals.
Acrobot Swingup
Ball In Cupcatch
Cartpole Balance
Cartpolebalance Sparse
1000
1000
400
750
800
500
600
200
500
250
250
200
0
1e6
1e6
1e6
1eB
Cartpole Swingup
Cartpoleswingup Sparse
Cheetah Run
Finger Spin
800
750
800
750
600
500
600
400
250
250
200
0
1e6
1e8
1e6
1e6
Fingertum Easy
Fingertum Hard
Hopper Hop
Hopper
1000
1000
750
300
750
750
500
500
200
500
250
250
100
250
0
1e6
1e6
1e6
1eB
Pendulum Swingup
Quadruped Run
Quadruped Walk
Reacher Easy
1000
750
750
750
500
500
500
500
250
250
260
250
0
2
1e6
Reacher Hard
Walker Run
Walker Stand
Walker Walk
1000
800
1000
600
800
800
500
600
400
400
250
400
dreamer mets
200
200
dreamer rollout
200
dreamer
1e6
1eG
1e6
1e6
Figure 5: Comparison of Dreamer with different search modes during exploration. Each curve represents the mean value over 3 runs with different seeds.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Sébastien Bubeck, Gilles Stoltz, Csaba Szepesvári, and Rémi Munos. Online optimization in x- armed bandits. In Advances Neural Information Processing Systems, pp- 201-208, 2009.
Guillaume M JB Chaslot, Mark HM Winands, Н JAAP VAN DEN HERIK, Jos WHM Uiterwijk, and Bruno Bouzy. Progressive strategies for monte-carlo tree search. New Mathematics and Natural Computation, 4(03):343-357, 2008. 2008.
Adrien Couêtoux, Jean-Baptiste Hoock, Nataliya Sokolovska, Olivier Teytaud, and Nicolas Bon- nard. Continuous upper confidence trees. In International Conference on Learning and Intelligent Optimization, pp. 433-445, Springer, 2011.
Rémi Coulom. Efficient selectivity and backup operators monte-carlo tree search. In International conference on computers and games, pp. 72-83. Springer, 2006.
Scott Fujimoto, Herke Van Hoof, and David Meger. Addressing function approximation error in actor-critic methods. arXiv preprint arXiv: /802.09477 2018.
Karol Gregor, Danilo Jimenez Rezende, Frederic Besse, Yan Wu, Hamza Merzic, and Aaron van den Oord. Shaping belief states with generative environment models for rl. In Advances in Neural Information Processing Systems, pp. 13475-13487, 2019.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 1803. 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with stochastic actor. arXiv preprint 2018.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603. 2019a.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pp. 2555-2566 2019b.
Diederik P Kingma and Max Welling. 13/2.61 14, 2013.
Auto-encoding variational bayes. arXiv preprint
Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with latent variable model. arXiv preprint 2019.
Timothy Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv:1509.022 02971, 2015.
Kendall Lowrey, Aravind Rajeswaran, Sham Kakade, Emanuel Todorov, and Igor Mordatch. Plan online, learn offline: Efficient learning and exploration via model-based control. arXiv preprint arXiv:1811.01848. 2018
Christopher Mansley, Ari Weinstein, and Michael Littman. Sample-based planning for contin- uous action markov decision processes. In ICAPS, 2011.
Thomas M Moerland, Joost Broekens, Aske Plaat, and Catholijn M Jonker. A0c: Alpha zero in continuous action space. arXiv preprint arXiv: 09613, 2018.
A. Nagabandi, G. Kahn, R. S. Fearing, and S. Levine. Neural network dynamics for model-based deep reinforcement learning with model-free fine-tuning. In 2018 IEEE International Conference on Robotics and Automation (ICRA), pp. 7559-7566. 2018.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. arXiv preprint arXiv: /001://400.08.2 2014.

---

Under review as a conference paper at ICLR 2021
Christopher D Rosin. Multi-armed bandits with episode context. Annals of Mathematics and Artifi- cial Intelligence, 61(3):203-22 2011.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, al. Mastering atari, go, chess and shogi by planning with a learned model. arXiv preprint urXiv:/91..08220 1.08265, 2019.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 2017.
David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. 2014.
David Silver, Julian Schrittwieser, Karen Simonyan, loannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676):354-3333 2017.
David Silver, Thomas Hubert, Julian Schrittwieser, loannis Antonoglou, Matthew Lai, Arthur Guez, Mare 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): 140- 1144, 2018
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Yuval Tassa, Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, and Nicolas Heess. dm.ontrol: Software and tasks for continuous control, 2020.
Tom Van de Wiele, David Warde-Farley, Andriy Mnih, and Volodymyr Mnih. Q-learning in enor- mous action spaces via amortized approximate maximization. arXiv preprint arXiv:2001.08116, 2020.
Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew Johnson, and Sergey Levine. Solar: Deep structured representations for model-based reinforcement learning. (2018). arXiv preprint 1s.LG/1808.09105, 2018.
10


Reviews:
[{"review_id": "qgmwgsQO1BZ", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "##########################################################################\n\nSummary:\n\nThe paper is developed on top of the Dreamer architecture, i.e. learning a latent space dynamics based on the image inputs to train policies. The difference is that, instead of using the already trained policy, this paper used MPC or MCTS to sample actions during the exploration phase to reduce the bias.  The authors demonstrated that their approach led to overall improved sample efficiency and final policy performance across many MuJoCo benchmark tests.\n\n##########################################################################\n\nReasons for score: \n\nOverall,  the methodology is sound and the writing is clear. The contribution, however, seemed minor since it is a small modification to the original Dreamer framework, and the improvement in the performance is not significant. Thus, my rating for this paper is weak acceptance. \n\n##########################################################################\n\nPros:\n\n(1) The writing is clear and easy to understand.\n\n(2) Comprehensive studies on many experiments to study the effectiveness of their method, unlike many learning papers which only selected a small subset of validation tasks. This gives us a full picture of the strength and weakness of the proposed approach. \n\nCons:\n\n(1) As mentioned before, the theoretic contribution of the paper is small. It modifies the SoTA with a minor tweak, and the results are not that significant.\n\n(2) In this paper, the observations used are from the original state spaces of the environments. On the contrary, Dreamer assumes inputs in the image space and that is the reason a latent space was introduced. It remains to see if the claim still holds if the studies are executed in the pixel space. \n\n##########################################################################\n\nConclusion:\n\nPlease add ablation studies in the image space as well and see if the conclusion still holds. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "lSNBJ0i5Zx", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes to integrate planning into Dreamer. The main idea is to apply a planning module on top of Dreamer to improve the quality of action selection. The planning via MCTS is on the learnt latent dynamics and the policy learnt by Dreamer. One of the challenges addressed in the paper is to perform planning on continuous action spaces. The proposed method is evaluated on 20 control tasks from the DeepMind Control suite, and compared against the original Dreamer algorithm, and a baseline planning method that does only rollout simulations. \n\n\nThe problem of using planning to enhance action selection for model-based RL is interesting and worth studying. However, the proposed idea is quite incremental. It might be worth trying, however, the experiment results are not exhaustive to evaluate the real benefits of the proposal. \n\nIn overall, the application of MCTS and the baseline Rollout on top of the deamer's learnt latent dynamics is quite straightforward. As the domain is continuous, therefore there is a bit challenge on the search tree's representation. Most techniques used in the paper is quite standard from existing works. In addition, there would be helpful if there are more ablation studies to look at the effect of the way MCTS is setup, .e.g. the amount of the fixed actions at branching, the number of simulations, etc. Those settings would affect how deep the policy tree is built, which roughly similar to the setting of horizon $H$ in Dreamer. The trade-off between this horizon length with the estimation error in model learning was well ablated in the Dreamer paper. It would be helpful to see the same ablation here. Given that MCTS can do planning under uncertainty (POMDP), i.e. on inaccurate model estimation, it would be great if the proposed idea discusses on this possibility and could address the problem of performance degradation with a large long look-ahead horizon.\n\nThe experiment results are not very convincing. There are many unfinished experiments. The proposed idea does not always outperform the baseline. A complex MCTS planning while consuming expensive computation, but performs worse than the baseline Rollout, and sometimes worse than the original Dreamer. More ablation studies might also be needed to make fair comparisons, i.e. while MCTS requires more planning time, could more computation budget be allocated for the baseline Rollout (more simulations or with larger tree settings) and Dreamer (more batch updates for action and value models)?\n\n\nThe notations, i.e. transitions, policy, etc., in Section 3 and 4 should be made consistent.", "extra_scores": {}, "extra_text": {}}, {"review_id": "fq8nPYujPA8", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors extend the Dreamer algorithm to use a different policy optimization mechanism, either a form of Monte-Carlo control or MCTS applied online in the continuous latent space at each decision step. This paper combines existing algorithms and compares two variants (“Rollout” vs “MCTS”) on continuous control domains.\n\nExperimental results in 20 control tasks suggest that the proposed approaches and Dreamer perform similarly on most tasks. On some tasks, the “Rollout” and “MCTS” action-selection approach seems to have an advantage, but there is no clear trend showing that MCTS provides an advantage over the Monte-Carlo control in these cases. The number of simulations and rollouts were fixed for these experiments, so it prevents more nuanced interpretation of these results. Overall it’s not clear that the alternative action-selection mechanism for Dreamer proposed in the paper has any clear advantage overall. In the cases where there is an advantage (e.g. Hopper Hop, Quadruped Run), we don’t know whether the gains come from having a stronger policy improvement or from the modified behavior policy (which may help exploration), it would be interesting to investigate these questions in more detail.\n\nOverall, the paper is clearly written and combines existing ideas in a sensible way, but it’s not clear what the take-away or potential impact of the paper is given there isn’t a particularly strong finding that comes out of this paper at the moment. Perhaps the authors could clarify their main takeaway for further discussion. \n\nAdditional questions:\n\n* Is the policy prior for search updated based on the search policy (as in MuZero)?\n* What network architecture was used for the experiments?\n\nMinor things:\n\n\n* Missing references for Progressive Widening technique in MCTS.\n* It might be worth clarifying in the text that the tree policy is not actually a proper UCB, but is an approximation (PUCT which incorporates a prior policy).\n* Return notation is inconsistent. G is used to denote random returns in Eq 10, but V is used in Eq 5.\n* Missing q in last term of Eq2?\n* The quality of the figures (resolution/format) should be improved\n* is improves -> is improved\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "OStEZYh92h", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "summary:\nThis paper extends Dreamer, a model-based RL algorithm trained through latent imagination, by additionally performing decision-time planning in the learned latent-space dynamics for action selection. Most of the components follow those of Dreamer: from the experiences collected by the agent, it learns a (latent-space) world model that comprises representation model, observation model, reward model, and transition model, which are trained by minimizing reconstruction loss with a KL regularizer. The value and action models are also trained as same as Dreamer. It computes value estimates for the imagined trajectories and performs gradient ascent using the reparameterized gradient. Still, unlike Dreamer, the proposed method does not work on raw pixels but uses low-dimensional features such as joint positions and velocities. Finally, additional online planning, either a simple rollout or an MCTS, is performed to select an action at each interaction with the environment. Experimental results demonstrate that the proposed shows a better sample efficiency in several domains.\n\n\n\npros:\n- This work shows that performing a search in a model-based RL can potentially benefit.\n\n- Experimental results show that the proposed agent improves the sample-efficiency of the Dreamer baseline in several domains.\n\n\nconcerns:\n- The main weakness of this work is the novelty. The difference to Dreamer is the additional adoption of online planning (or decision-time planning), but this additional component itself is not new to model-based RL.\n\n- Planning in the learned latent-space dynamics is not well-motivated in the situation where the agent takes 'low-dimensional features' as an observation directly. Why does the agent have to learn the complex latent dynamics even when the low-dimensional features are accessible? Under this circumstance, it seems to be more natural to learn and plan on the dynamics in the raw (low-dimensional) observation space.\n\n- It is not convincing that the proposed agent significantly outperforms the baseline. It is unclear what exactly the shaded area denotes in the plot (standard deviation? or standard error?), but the shaded areas of Dreamer and Dreamer+MCTS are being overlapped in Figure 3a.\n\n- The analysis of why search is beneficial during exploration seems to be not substantial. More ablation experiments could have strengthened the paper. The raised hypothesis, that 'search can be beneficial when action-value estimation is imperfect', is not supported by evident ablation study. If the inaccurate action-value estimation is the problem, why (compounding) model error during a search should be less problematic?\n\n\n\ncomments and questions:\n- In order to improve sample efficiency, structured exploration may be important. What aspect of the search of the proposed method can be helpful for better exploration (or other factors contributing to performance)?\n\n- It seems that the dynamics of the continuous latent variables are stochastic. Since the number of latent states is infinite, there may be a need for special treatment to make MCTS tractable, but it is not clearly described. How was the stochastic transition handled in MCTS? (e.g. double progressive-widening is used?)\n\n- Above Eq. (9): there is no definition of $\\alpha$.\n\n- The values of J in Eq. (2) should be maximized, not minimized, thus it seems to be awkward to be called a loss for each J. $q$ is omitted in the KL term.\n\n- Experiments could have been more thorough. In Figure 5, not all the experiments were conducted until the 2 * 10^6 timesteps (e.g. dreamer+mcts in Quadruped Run, Reacher Hard, ...). It would be great to see ablation studies that show the effect of the planning horizon, search budget, and so on, which may be helpful to understand why the proposed method could perform better than the baseline. We can also explore which role of search is more significant between the search for exploration of collecting training samples and search at evaluation time.\n\n- Other model-based RL algorithms that operate in the default feature representation could be a good baseline for comparison, e.g. MBPO (Janner et al. NeurIPS 2019).\n\n- In figures 3-5, what does the shaded area represent? (standard deviation? standard error?)", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes an extension to the Dreamer agent in which planning (either via MCTS or rollouts) is used to select actions, rather than sampling from the policy prior. The results show small improvements over the baseline Dreamer agent.

Pros:
- Important study on incorporating decision-time planning into Dyna-based agents
- Evaluation on many control tasks rather than just a few

Cons:
- Lack of ablations and detailed analysis
- Claims aren't backed up by quantitative results

The reviewers generally felt that the approach taken in the paper lacked novelty. I agree that the approach is somewhat incremental (in fact I think it is also an instance of [1]). While both incremental changes and reimplementations of older methods with newer techniques can indeed be valuable, the current paper falls short in terms of the evaluation. As pointed out by several reviewers, there is no in-depth analysis explaining the design choices in which rollouts or MCTS are most likely to help (e.g. search budget, exploration parameters, etc.). As these parameters can play a large role in performance, I think it is important to characterize their effect on the agent---otherwise, I do not think there is a clear learning regarding how to translate these results to other domains and tasks. Additionally, and perhaps even more seriously, there are a number of claims made in the paper about the proposed method being more data efficient or higher performance. But, it is not clear visually that these improvements are statistically significant, and no quantitative tests have been run (and if the authors want to make a claim about data efficiency, I'd especially encourage them to report a metric like cumulative regret). Finally, while the incomplete runs are not a reason for rejection on their own, they do add to my overall sense that the paper is incomplete in its current form.

Given the above reasons, I do not feel this paper is ready for publication at ICLR. I'd encourage the authors to perform more careful ablations of the effect of incorporating search into the agent, and to back up their claims with more rigorous quantitative results.

One small point: the authors wrote in the rebuttal that "we are not aware of any work which investigates look-ahead search-based planning for continuous control with learned dynamics". Grill et al. [2] uses MCTS with learned dynamics in a modification of MuZero, though only applies it in one continuous control task (Cheetah Run).

1. Silver, D., Sutton, R. S., & Müller, M. (2008). Sample-based learning and search with permanent and transient memories. ICML.
2. Grill, J. B., Altché, F., Tang, Y., Hubert, T., Valko, M., Antonoglou, I., & Munos, R. (2020). Monte-Carlo tree search as regularized policy optimization. ICML.
Decision: Reject