E VALUATING MODEL - BASED PLANNING AND PLANNER AMORTIZATION FOR CONTINUOUS CONTROL
Arunkumar Byravan ∗† Leonard Hasenclever ∗† Piotr Trochim ∗ Mehdi Mirza ∗
Alessandro Davide Ialongo ‡§ Yuval Tassa ∗ Jost Tobias Springenberg ∗ Abbas Abdolmaleki ∗
Nicolas Heess ∗
Josh Merel ‡¶
Martin Riedmiller ∗
A BSTRACT
There is a widespread intuition that model-based control methods should be able to surpass the data efﬁciency of model-free approaches.
In this paper we attempt to evaluate this intuition on various challenging locomotion tasks.
We take a hybrid approach, combining model predictive control (MPC) with a learned model and model-free policy learning; the learned policy serves as a proposal for MPC.
We ﬁnd that well-tuned model-free agents are strong baselines even for high DoF control problems but MPC with learned proposal distributions and models (both trained on the ﬂy or transferred from related tasks) can signiﬁcantly improve performance and data efﬁciency in hard multi-task/multi-goal settings.
Finally, we show that it is possible to distil a model-based planner into a policy that amortizes the planning computation without any loss of performance.
Videos of agents performing different tasks can be seen on our website.
I NTRODUCTION
In recent years, model-free RL algorithms have improved and scaled to the point where it is feasible to learn adaptive behavior for high-dimensional systems in diverse settings (Schulman et al., 2017;
Heess et al., 2017).
Despite these improvements, there is a widespread intuition that model-based methods can further improve data efﬁciency.
This has led to recent advances that demonstrate improved learning efﬁciency by leveraging model learning during policy training (Lowrey et al., 2019; Nagabandi et al., 2020; Hubert et al., 2021).
However, there is also work which urges moderation in interpreting these results, by showing that well-tuned model-free baselines can compare favorably against some model-based approaches (Springenberg et al., 2020; van Hasselt et al., 2019).
In this paper, we focus on model-based RL with learned models and proposals.
1 We use modelpredictive control (MPC) to improve the quality of behavior generated from a policy that serves as a proposal for the planning procedure.
When acting, the MPC-based search procedure improves the quality of the data collected which serves as a sort of “active exploration”.
This data can then be used to improve the proposal, in effect consolidating the gains.
This hybrid approach interpolates between model-free RL, on one side, and trajectory optimization with a model on the other.
Fundamentally, the spectrum on which this hybrid approach is situated reﬂects a trade-off in terms of reusability / generality versus compute cost at deployment time.
Policies obtained by model-free
RL often generalize poorly outside of the situations they have been trained for, but are efﬁcient to execute (i.e., they are fully amortized).
Models offer potentially greater generalization, insofar as the model is accurate over a broad domain of states, but it can be computationally costly to derive actions from models.
Ultimately, the pure planning approach involving deriving actions from models is usually prohibitive and some form of additional knowledge is required to render the search space tractable for real-time settings.
Coming from the model-free RL perspective, the natural approach is to learn a policy that serves as a relatively assumption-free proposal for the planning process.
While in most model-based RL publications the ambition is to speed up policy learning for a single task, we believe this is not necessarily the most promising setting, as it is difﬁcult to dissociate the learning of a dynamics model (which is potentially task independent) from that of a value function (which is task speciﬁc).
A different intuition is that model-based RL might not accelerate learning on a particular problem, but rather enable efﬁcient behavior learning on new tasks with the same embodiment.
In such transfer settings, we might hope that a dynamics model will offer complementary generalization to the policy, and we can transfer the policy, model, or both.
In this work, we validate these intuitions about the relative value of the different components of this system.
For high-dimensional locomotion problems, we ﬁnd that even with a good learned model successful MPC requires good proposal distributions to succeed.
This effect is particularly pronounced in domains with multiple goals (or equivalently multiple different tasks).
Building on this insight we show that it is sometimes possible to leverage learned models with a limited search budget to boost exploration and learn policies more efﬁciently.
Policies that are trained by off-policy updates from data acquired through planning do not reliably perform well when used without planning in our experiments.
To overcome this problem, and enable planner amortization into a policy that does not require planning at test time, we propose training the policy with a combination of a behavioral cloning objective (on MPC data) and an off-policy update with MPO (Abdolmaleki et al., 2018b).
When transferring models and proposals to other tasks we ﬁnd only marginal improvements in data efﬁciency relative to a well-tuned model-free baseline.
Overall, our results suggest that while MPC with learned models can lead to more data efﬁciency, and planners can be amortized effectively into compact policies, it is not a silver bullet and model-free methods are strong baselines.
B ACKGROUND & RELATED WORK
Model-predictive control refers to the use of model-based search or planning over a short horizon for selecting an action.
In order to make it computationally tractable, it is common to “seed” the search at a given timestep with either an action from a proposal policy or an action obtained by warmstarting from the result of the previous timestep.
In simulation, it has been demonstrated that a fairly generic MPC implementation can be effective for control of a relatively high DoF humanoid (Tassa et al., 2012) and that MPC with learned models (Chua et al., 2018; Wang & Ba, 2020; Nagabandi et al., 2020) can achieve data-efﬁcient solutions to simple control problems.
Real RC helicopter control has also been achieved using an MPC approach that made use of a learned model (Abbeel et al., 2006).
MPC approaches are gaining wider use in robotics (multicopters (Neunert et al., 2016;
Torrente et al., 2021), quadruped (Grandia et al., 2019; Sleiman et al., 2021), humanoids (Tassa et al., 2014; Kuindersma et al., 2016)), and dexterous manipulation (Nagabandi et al., 2020); but the computational speed of the planner is a bottleneck for hardware deployment.
There are different ways around this, with the core practical solution being to plan in lower dimensional reduced coordinate models.
Alternatively, POPLIN (Wang & Ba, 2020) explores learning proposals for MPC and planner distillation but is tested on simple tasks and does not leverage model-free RL for proposal learning, and LOOP (Sikchi et al., 2022) uses an off-policy actor to control exploration for MPC.
Amortized policies map states to actions quickly, but implicitly reﬂect considerable knowledge by being trained on diverse data.
Model-free RL produces policies which amortize the knowledge reﬂected in rollouts required to produce them.
Similarly, it is possible to produce diverse trajectories from a planner and distil them into a single policy (Levine & Koltun, 2013; Mordatch & Todorov, 2014; Mordatch et al., 2015; Pan et al., 2017; Xie et al., 2021).
Reinforcement learning approaches with MPC have become more popular recently, and our work ﬁts within this space.
As noted in the introduction, previous work often emphasizes the role of amortization through the learning of value functions and models.
TreePI (Springenberg et al., 2020),
MuZero (Schrittwieser et al., 2020; Hubert et al., 2021), SAVE (Hamrick et al., 2020), DPI (Sun et al., 2018) and MPC-Net (Carius et al., 2020) all perform versions of hybrid learning with model-based
MCTS or planning being used to gather data which is then used to train the model and value function to accelerate learning.
Other recently proposed algorithmic innovations blend MPC with learned value estimates to trade off model and value errors (Bhardwaj et al., 2021).
Here, we primarily consider learning dynamics models to enable transfer to new settings with different reward functions.
Other uses of models unlike hybrid MPC-RL schemes have also been explored in the literature; however, they are not the focus of this work.
Nevertheless we highlight two of these approaches: value gradients can be backpropagated through dynamics models to improve credit assignment (Nguyen &
Widrow, 1990; Heess et al., 2015; Amos et al., 2020) and it is possible to train policies on model rollouts to improve data efﬁciency (Sutton, 1991; Janner et al., 2019).
Recently, there has also been considerable effort towards exploring model-based ofﬂine RL methods in order to gain full value from ofﬂine datasets (Yu et al., 2020; Argenson & Dulac-Arnold, 2021; Kidambi et al., 2020).
T ASKS
In this paper we consider a number of locomotion tasks of varying complexity, simulated with the
MuJoCo (Todorov et al., 2012) physics simulator.
We consider two embodiments: an 8-DoF ant from dm_control (Tunyasuvunakool et al., 2020) and a model of the Robotis OP3 robot with 20 degrees of freedom.
For each embodiment, we consider three tasks: walking forward, walking backward and “go-to-target-pose” (GTTP), a challenging task that is the focus of our evaluation.
In all tasks, the agent receives egocentric proprioceptive observations and additional task observations.
In the walking forward and backward tasks the agent is rewarded for maximizing forward (or backward) velocity in the direction of a narrow corridor.
For the OP3 robot we also include a small pose regularization term.
The task is speciﬁed through a relative target direction observation.
The GTTP task, which builds on existing motion tracking infrastructure(Hasenclever et al., 2020), consists of either body on a plane, with a target pose in relative coordinates as a task-speciﬁc observation and proximity to the target pose rewarded.
When the agent is within a threshold distance of the target pose (i.e. it achieves the current target), it gets a sparse reward and a new target is sampled.
For the ant we use target poses from policies trained on a standard go-to-target task (Tunyasuvunakool et al., 2020).
For the OP3, we use poses from the CMU mocap database (cmu) (retargeted to the robot).
We use thousands of different target poses; the agent has to learn to transition between them.
Thus the GTTP task can be seen as either a single highly diverse task or as a multi-task setting with strongly related tasks and consistent dynamics.
We believe the GTTP task should be particularly amenable to model-based methods: it combines a high-dimensional control problem with a diverse goal distribution.
This makes it hard to solve quickly with model-free methods.
However, since all goals share the same dynamics, a dynamics model should help leverage the common structure.
See Sec.
A for an in-depth motivation and Sec.
F.5 for results on tasks from the DeepMind Control Suite (Tassa et al., 2018).
M ETHOD 4. M ODEL PREDICTIVE CONTROL (MPC) FOR BEHAVIOR GENERATION
We consider control obtained by model-based planning that reﬁnes initial action-samples from a proposal distribution.
MPC executes the ﬁrst action from the planned action sequence and iteratively re-plans after each timestep (c.f. actor loop in Alg. 1).
In this scheme, actions are sampled either from the planner ( PLANNER ) with probability p plan or the proposal π θ (with parameters θ ) with probability 1 − p plan .
Setting p plan = 0 results in using only samples from the proposal and vice-versa for p plan = 1 .
Setting p plan = 0 .
5 leads to interleaved execution of proposal and planner actions, like the stochastic mixing of student and expert actions in DA GGER (Ross et al., 2011).
The PLANNER subroutine takes in the current state s t , the proposal π θ (with parameters θ ), a model m φ (with parameters φ ) that predicts next state s t +1 given current state s t and action a t , the known reward function r ( s t , a t , s t +1 ) and, optionally, a learned state-value function V ψ ( s ) (with parameters ψ ) that predicts the return given state s .
While we consider primarily a Sequential Monte Carlo based planner (SMC, Alg. 2) (Piché et al., 2019; Gordon et al., 1993) and the Cross-Entropy Method (CEM,
Alg. 3) (Botev et al., 2013), other planners (Williams et al., 2017) could be used.
See Sec.
B.
Given: Randomly initialized proposal π θ , (pre-trained or random) model m φ , random critic Q ψ , optionally (pre-trained or random) task-agnostic proposal ρ ω . //
Modules to be learned
Given: Known reward r , planning probability p plan , replay buffer B , MPO loss weight α , BC loss weight β , learning rates & optimizers (ADAM) for the different modules. //
Known modules and parameters //
MPC loop – Asynchronously on the actors
while True do
Initialize ENV and observe state s 0 .
while episode is not terminated do // Choose between planner and proposal action depending on p plan // Use mixture of (pre-trained) task-agnostic proposal ( ρ ω ) and learned proposal ( π θ ) as proposal for proposal transfer exps (Sec. 5.4).
Use pre-trained model for model transfer exps (Sec. 5.3) .
a t ∼ (cid:26) PLANNER ( s t , π θ , m φ , r, V ψ ) if x ≤ p plan π θ ( s t ) otherwise.
where x ∼ U [0 , 1]
Step ENV ( s t , a t ) → ( s t +1 , r t ) and write transition to replay buffer B end while end while //
Asynchronously on the learner while True do
Sample batch B of trajectories, each of sequence length T from the replay buffer B Update action-value function Q ψ based on B using Retrace (Munos et al., 2016).
Update model m φ based on B using multi-step loss in Equation 20 (Sec. 4.2) Update proposal π θ based on B using Equation 4 (Sec. 4.3) Optionally, for from-scratch experiments, update task-agnostic proposal ρ ω using behavioural cloning (Equation 1) on transitions in B .
end while
The hybrid agent in Alg. 1 has several learnable components: the model m φ , the proposal π θ as well as the (optional) value function V ψ ; these are learned based on data generated by the MPC loop on the actor and we measure learning progress w.r.t environment steps collected.
Note that we assume the reward function is known since it is under the control of the practitioner in most robotics applications.
4. L EARNING DYNAMICS MODELS
Successfully applying MPC requires learning good dynamics models.
In this paper, we train predictive models m φ that take in the current state s t and action a t and predict the next state s t +
= m φ ( s t , a t ) .
This model has both learned and hand-designed components; the task-agnostic proprioceptive observations are predicted via “black-box” neural networks while any non-proprioceptive – taskspeciﬁc – observations (e.g. relative pose of the target) are calculated in closed form from the predicted proprioceptive observations together with a known kinematic model of the robot.
The learned components are parameterized as a set of deterministic feed-forward MLPs, one per observation group (e.g. joint positions), and predict the next observation from the current observations and action.
This model is trained end-to-end via a multi-step squared error (Eqn. 20) between an open loop sequence of predicted states and true states sampled from the replay buffer.
We use deterministic models as they are more amenable to multi-step losses and an ablation study (Sec.
F.4 of the supplement) as well as a recent benchmarking effort (Lutter et al., 2021) found minimal differences between using deterministic vs stochastic model ensembles for model-based planning.
Training happens from scratch, concurrently with policy learning.
More details are in Sec. C of the supplement.
4. L EARNING PROPOSAL DISTRIBUTIONS
As we will show, planning with a good dynamics model by itself is insufﬁcient for solving challenging continuous control problems.
That is, generating behavior according to the MPC loop in Alg. 1 with a goal-unaware proposal fails to solve the challenging GTTP tasks even when planning with the ground-truth model and a known reward function.
This is due to the fact that the action space is large, and the planning algorithm is limited by computational constraints and a ﬁnite horizon.
In order for planning to succeed, we ﬁnd it vital to provide planners with an action proposal distribution, which helps search by increasing the probability of sampling plausible, task-relevant actions.
In addition, we show that the learned proposal can itself be deployed without planning at test time; this can be particularly useful in compute constrained settings where planning on the robot may not be feasible.
But what constitutes a good proposal?
A natural answer to this question is that a good proposal produces actions, for each state, that lead to task success, and the planner selects the best among already good candidate actions.
A simple strategy to obtain such a proposal is to iteratively amortize the planning process.
That is, starting from any proposal, execute the planner for a given duration to return an improved action and ﬁt the returned action; this will sharpen the proposal and lead to the planner focusing on better actions in the next iteration.
This idea has been used in the RL community before and underlies modern search based algorithms such as AlphaZero (Schrittwieser et al., 2020).
In the MPC setting we can formalize the idea as follows: let π B ( a | s ) refer to the stochastic mixture of PLANNER ( s t , π θ , m φ , r, V ψ ) and π θ ( a | s ) with probability p plan (i.e., as actions are selected according to the MPC loop in Alg. 1).
And let D π B be a dataset of states and actions collected by executing this policy – note that in practice for data efﬁciency reasons we consider a dataset D ˜ π B , where ˜ π B is the average behavior distribution during training, which changes over time due to the proposal and model being learned.
Clearly, assuming the planning procedure produces improved actions, we expect that ˜ π B is at least as good as the average proposal policy ˜ π θ .
We can then improve our proposal by minimizing the forward KL divergence E s ∈D ˜
π B [ KL (˜ π B [Equation]
·| s ) | π θ ( ·| s ))] which leads to an improving average proposal over time, and is equivalent to maximizing:
[Equation]
i.e. the proposal is learned by amortizing the planner via behavioral cloning of the planner actions.
Unfortunately, in cases where the planner cannot ﬁnd an improvement on the proposal (as is the case for random proposals in our experiments), the above objective may lead to premature convergence at sub-optimal behavior (see e.g. Wang & Ba (2020)).
To ameliorate this issue, we can consider hybrid updates that both amortize the planner, but also directly favor actions that lead to an improvement in terms of cumulative return via an off-policy RL policy update.
We will use the MPO (Abdolmaleki et al., 2018b) algorithm, although other recent RL algorithms could be used instead.
More precisely, the MPO policy improvement step involves using a learned action-value function Q π θ ( s, a ) ≈
E π
θ [ (cid:80) t γ t r t ( s t , a t ) | s
=
s, a 0 = a ] , and optimizing the KL constrained RL objective:
[Equation]
It is known (see Abdolmaleki et al., 2018b) that the solution of this optimization is given in closed form as q ( a | s )
∝ π θ ( a | s ) exp( Q π θ ( s, a ) /η ) , where η is a dual variable optimized such that the KLconstraint on the policy is fulﬁlled.
We can ﬁt this improved policy by minimizing the KL divergence E s [Equation]
˜ π B [ KL ( q ( ·| s ) | π θ [Equation] ·| s ))] which is equivalent to maximizing the weighted log-likelihood:
[Equation]
where π θ (cid:48) is the last reference policy (ﬁxed for this optimization).
In practice MPO uses an additional trust-region constraint to stabilize learning.
We can combine the two objectives (Eqns. 1 and 3) for improving the proposal via a simple weighting to obtain the complete objective
[Equation]
The full agent showing the actor and learner loops is shown in Alg.
In our experiments we compare several variants corresponding to different choices of p plan , α and β .
See Sec.
D for further details and rationale on our choice of MPO as the learning algorithm.
R ESULTS 5. E VALUATING PRE - TRAINED MODELS AND PROPOSALS
In a ﬁrst set of experiments we study how MPC (Alg. 1, p plan = 1 ) performs on the locomotion tasks considered in the paper.
We evaluate performance with two different models, the ground truth
MuJoCo simulator and a pre-trained model that is trained on data from a successful agent on the corresponding task (see Sec. 4.2).
We also consider two different proposal distributions: a zero-mean
Gaussian proposal and a task-agnostic proposal that is pre-trained with a behavioral cloning objective on logged data from a successful agent (similar to the prior distribution in Galashov et al.
(2019), see Sec.
E.2).
This pre-trained proposal depends only on proprioceptive information but not task speciﬁc observations, and therefore is a reasonable prior capturing average behavior.
We assume a known reward function but do not use a learned value function.
The best results of a large hyper
[Table]
planner.
parameter sweep are shown in Table 1, together with performance of the two proposals and baseline performance of a successful agent whose data was used to train the pre-trained model and proposal (labelled ‘Near-Optimal performance’).
Results from an additional sweep comparing performance of planning w.r.t computational budget and proposal use can be seen in Sec.
F.1.
When planning with a zero-mean Gaussian proposal and the ground truth dynamics, both planners improve signiﬁcantly on the performance of the Gaussian proposal for the simpler Ant tasks but not for the harder OP3 tasks (this may be in part because the OP3 tasks terminate if the robot falls, which makes planning hard).
This result holds across a wide range of computational budgets (see subsection F.1).
Planning with a pre-trained model and the Gaussian proposal performs similar to the proposal itself on the harder OP3 tasks but there is a small improvement, albeit less than when planning with the ground truth dynamics, for the simpler Ant tasks.
Speciﬁcally on the Ant GTTP task, planning leads to better rewards but fails to consistently achieve target poses and there remains a large qualitative difference between the planner results and near-optimal performance.
When planning with the pre-trained task-agnostic proposal instead, the performance far exceeds the results from planning with the Gaussian proposal or executing the pre-trained proposal without any planning.
This highlights the need for a suitable proposal, especially for the high-dimensional OP tasks, further motivating our approach learn a proposal for MPC with model-free RL.
When planning with a pre-trained proposal, using the ground truth dynamics or a pre-trained model achieve similar performance, indicating that our pre-trained models are suitable for planning.
Lastly, both planners perform similarly across tasks, though on the harder GTTP tasks SMC slightly outperforms CEM.
We use SMC throughout the paper as it makes better use of the proposal unlike CEM which uses it only for plan initialization (see Sec.
F.1 for full results).
5. L EVERAGING MPC WITH A LEARNED MODEL AND PROPOSAL
Next, we evaluate different approaches to learning a proposal from scratch for MPC based on the approach discussed in Sec. 4.3.
In particular, we consider the following variants:
• MPO: MPO with a distributional critic similar to Hoffman et al.
Corresponds to α = 1 , β = 0 , p plan = 0 from Eqn. 4 and Alg. 1.
MPO hyper-parameters are tuned for each task (and listed in Sec. F.2 of the supplement).
• MPC+MPO : MPC to collect data.
MPO objective for learning ( α = 1 , β = 0 , p plan = 0 . 5 ).
• MPO+BC : Adding MPO and BC objectives ( α = 1 , β > 0 , p plan = 0 ), where β is tuned per task.
• MPC+MPO+BC : MPC to collect data.
Combined MPO+BC objective ( α = 1 , β > 0 , p plan = 0 . 5 ).
The model is trained from scratch for the MPC variants (see Sec. 4.2).
We choose p plan = 0 .
5 for all our MPC experiments as it worked better than p plan = 1 . 0 , and use 250 samples and a planning horizon of 10 for SMC (see Sec.
F.2 for ablations).
We tune the BC objective weight β per task.
Figure 2 (left & center columns) presents performance comparisons on the go to target pose tasks 2 , showing the reward and corresponding number of achieved target poses for each variant.
MPC+MPO outperforms MPO signiﬁcantly with regards to actor performance (see left column, solid lines).
However even though MPC+MPO led to higher reward trajectories earlier in the learning process, the proposal distribution by itself tended to perform only as well as the MPO baseline.
We speculate that this phenomenon is due to the increased off-policyness of the planner trajectories.
Adding the BC objective mitigates this issue, improving the performance of both MPO 3 and MPC+MPO .
The resulting variant MPC+MPO+BC signiﬁcantly outperforms all other baselines in terms of actor and proposal performance as well as learning speed.
Additionally, we tried bootstrapping with the learned critic for planning, but did not observe signiﬁcant improvements on the GTTP tasks (see Fig. 12).
We also tested our approach on simpler forward and backward walking tasks for both the Ant and
OP3 bodies.
Figure 2 (right column) shows the results from the backward walking tasks.
For the
Ant (top row), MPC+MPO signiﬁcantly outperforms MPO early on during training but reaches similar asymptotic performance while for the OP3 (bottom row) the difference between MPO and MPC+MPO is small throughout training.
Interestingly, on these simpler tasks the proposal for MPC+MPO matches the actor performance and adding the BC objective gives only minor performance gains.
We posit that a well-tuned implementation of the model-free MPO baseline achieves near-optimal performance on these tasks and provides a strong baseline that is hard to beat both in terms of data efﬁciency and performance even with the true reward function and bootstrapping with a learned critic.
We present further details as well as forward walking results showing similar trends in the supplement (Fig. 13).
Lastly, since the BC objective substantially improves results for both MPO and MPC+MPO we use the BC variants in all further experiments.
Videos of learned behaviors can be seen at our website.
Ablations: We additionally ran several experiments ablating our design choices.
The results of these experiments can be found in the supplement.
Figure 10 and Figure 11 show the results of varying p plan as well as the number of samples used by the planner and the planning horizon, respectively.
In
Figure 12 we show the effect of bootstrapping with the learned value function within MPC.
Figureand Figure 18 compare deterministic models with PETS-style stochastic models (Chua et al., 2018) in our setting.
Furthermore, in Figure 19 we report results using ensembles rather than a single model.
Finally, Figure 20 shows results on the walker and humanoid tasks in the DeepMind Control Suite.
5. M ODEL TRANSFER
In this section, we study how well models can transfer between tasks.
We ﬁrst explore to what extent a learned model can boost performance on a complex task.
As a control experiment we transfer learned models from the OP3 GTTP task to the same task.
We consider three different settings (in addition to the baseline MPO+BC ) where a) no model is transferred and one is trained from scratch on the target task ( MPC+MPO+BC ), b) the transferred model is ﬁnetuned on the target task ( MPC+MPO+BC+Finetune ) and c) the transferred model is kept frozen throughout training on the target task ( MPC+MPO+BC+Frozen ).
In this experiment (Figure 3 left column, solid lines)
we ﬁnd little improvements in learning speed or asymptotic performance when transferring a model vs learning from scratch.
Transferring a frozen model (which can fail to generalize to out of distribution data) performs slightly worse vs ﬁnetuning the transferred model.
This trend also holds in other transfer settings such as transferring models from the forward walking task to GTTP and transferring from the GTTP task to backward walking (Figure 3, center/right column, respectively, solid lines).
While this result is perhaps surprising, it does agree with our initial investigation of planning with learned models (Table 1) where we saw poor performance on all our tasks without a good proposal for the planner to leverage.
Transferring a good model still does not solve the initial exploration problem, especially on the harder tasks.
Videos of learned behaviors can be seen here.
5. P ROPOSAL ( AND MODEL ) TRANSFER
In addition to transferring learned models across tasks we also consider transferring learned proposals.
On each source task we train a proposal that is dependent only on proprioceptive information and lacks task-speciﬁc information (similar to the one used in subsection 5.1) and is thus transferable across tasks.
As a simple approach to transfer a pre-trained proposal we used it for action generation while keeping the learning objective unchanged.
Concretely, we use a mixture of the reloaded proposal from a source task and the learned proposal on the target task as our proposal distribution ( π θ ) for the MPC loop in algorithm 1.
The mixture weight of the reloaded proposal is annealed linearly from 1 to 0 in a ﬁxed number of learning steps (tuned per task); see Sec.
F.3 of the supplementary material.
Again, we ﬁrst transfer the proposal trained on the GTTP task to the same task (Figure 3, left column, dashed lines).
Transferring the proposal leads to faster learning for both MPO+BC and MPC+MPO+BC as well as to a smaller extent, better asymptotic performance.
Transferring the model and proposal together ( MPC+MPO+BC+Finetune , MPC+MPO+BC+Frozen ) does not lead to any additional improvements, further strengthening our intuitions from the model transfer experiments.
Next we transfer proposals between different source and target tasks, which yielded some nuanced insights into the need for compatibility of the tasks.
Figure 3 (center column, dashed lines) shows the results when transferring a proposal from the forward walking task to the GTTP task and Figure(right column, dashed lines) presents proposal transfer results from the GTTP task to backward walking.
Interestingly, proposal transfer hurts both MPC+MPO+BC and MPO+BC in both these cases, especially for the high-dimensional OP3, leading to slower learning and lower ﬁnal performance on both the target tasks.
These results provide complementary insights regarding proposal transfer: there should be good overlap between the data distributions of the source and target tasks for proposal transfer to succeed.
This is not the case in these transfer experiments; forward walking has a very narrow goal distribution compared to GTTP making the resulting proposal far too peaked, and while a proposal from the GTTP task would be quite broad (see website for some trajectory rollouts) it is highly unlikely to capture the behavior of walking backwards.
Overall, these results provide encouragement that combining the right proposal, potentially trained on a diverse set of tasks, together with model-based planning can lead to efﬁcient and performant learning on downstream tasks.
For more transfer results and plots showing amortized policy performance see the supplement (Sec. F.3).
D ISCUSSION
Our initial experiments highlighted that a good dynamics model is not enough to solve challenging locomotion tasks with model predictive control when computation time is limited, especially in tasks with high task and control complexity.
In such settings a good proposal is necessary to guide the planner.
Motivated by this ﬁnding we study different approaches to learning proposals for MPC, considering variants that combine the learning objective from MPO, a model-free RL algorithm, together with a behavioral cloning objective for efﬁcient planner amortization.
We also evaluate transfer performance across tasks where either the proposal, learned model, or both are transferred.
Overall our results show that for the locomotion domains considered in this paper, MPC with a learned model and proposal can yield modest improvements in data efﬁciency relative to well-tuned model-free baselines.
We found that the gains are larger as task and control complexity increase.
On simpler walking tasks we saw very small improvements that are further diminished if an amortized policy is desired at test time.
On our most challenging task, the OP3 go to target pose task, we see both signiﬁcantly faster learning speed and improved asymptotic performance.
A common justiﬁcation for model-based approaches is the intuition that models can more easily transfer to related tasks since the dynamics of a body are largely task independent.
We attempted to validate this intuition but had difﬁculty achieving large gains in data efﬁciency even when transferring to the same task.
We speculate that this ﬁnding is related to the difﬁculty of planning with a limited search budget in a multi-goal/multi-task setting even with a near perfect model.
If the overall system is limited by the lack of a good proposal then model transfer by itself may have a negligible effect.
When transferring models and proposals to new simpler tasks we also did not ﬁnd substantial beneﬁts.
There are a number of other potential pitfalls in this setting in addition to the lack of a good proposal.
If a proposal leads to trajectories that are inconsistent with the state distribution on which the model was trained, MPC may add little value.
As we observed when transferring from the go-to-target-pose tasks to walking backwards, an unsuitable proposal may limit asymptotic performance.
Additionally, on tasks with fairly narrow goal distributions a well-tuned model-free method can perform just as well as a model-based agent.
This suggests tasks with multi-task/multi-goal settings provide a good test bed to showcase the strengths of model-based approaches and aid further research.
On the whole, the gains from MPC with learned models in our setting are meaningful, but not so dramatic as to be a silver bullet, a ﬁnding similar to Springenberg et al.
; Hamrick et al.
This paper focused on learning complex locomotion tasks from state features, using a structured dynamics model as well as focusing on MPC as the way to leverage the model.
There are a number of additional settings where models can be used differently and may aid transfer.
For example, in partially observed tasks with pixel observations, transferring models and representations may lead to improvements in data efﬁciency (Byravan et al., 2020; Hafner et al., 2018; 2020). A CKNOWLEDGMENTS
We would like to thank Jackie Kay, Alex X. Lee, Claudio Fantacci, Jessica Hamrick and Tom Erez as well as the rest of the team at Deepmind for insightful discussions and feedback.
The data used in this project was obtained from mocap.cs.cmu.edu.
The database was created with funding from NSF
EIA-0196217.
A T ASKS
We consider a number of locomotion tasks of varying complexity, simulated with the MuJoCo (Todorov et al., 2012) physics simulator.
We consider two embodiments: an 8-DoF ant from dm_control (Tunyasuvunakool et al., 2020) and a model of the Robotis OP3 robot with 20 degrees of freedom.
For each embodiment, we consider three tasks: walking forward, walking backward and “go-to-target-pose” (GTTP), a challenging task that is the focus of our evaluation.
In all tasks, the agent receives egocentric proprioceptive observations (joint angles, joint, linear and angular velocities as well as end effector positions).
In addition we provide the world z-axis in the robot’s frame of reference.
For the ant the proprioceptive observations are 37 dimensional.
For the OP3 robot the proprioceptive observations are 64 dimensional.
In addition to proprioceptive observations the agent also receives task speciﬁc observations that describe the task goal and differ by task.
A. W ALKING TASKS
In the walking forward and backward tasks the agent is rewarded for maximizing forward (or backward) velocity in the direction of a narrow corridor:
[Equation]
where v denotes the velocity of the robot and e target denotes a unit vector in the direction of desired movement (both in the frame of the agent).
For the OP3 robot we also include a small pose regularization term to regularize towards a walking pose.
The agent observes the target direction e target ∈ R 3 in it’s egocentric frame of reference as a task-speciﬁc observation for the walking tasks.
[Figure]
The GTTP task consists of either body on a plane, with a target pose in relative coordinates as a task-speciﬁc observation and proximity to the target pose rewarded.
When the agent is within a threshold distance of the target pose (i.e. it achieves the current target), it gets a sparse reward and a new target is sampled.
For the ant we use target poses from policies trained on a standard go-to-target task (Tunyasuvunakool et al., 2020).
For the OP3, we use poses from the CMU mocap database (cmu) (retargeted to the robot).
We use thousands of different target poses; the agent has to learn to transition between them.
Thus the GTTP task can be seen as either a single highly diverse task or as a multi-task setting with strongly related tasks and consistent dynamics.
We extend existing motion tracking infrastructure (Hasenclever et al., 2020) to build our task.
The agent agent receives task speciﬁc observations that specify the target pose relative to the agent.
Speciﬁcally, we use relative joint angles, as well as relative root position and relative positions and orientations of a number of different body parts, all expressed in the egocentric frame of the embodiment.
These task observations are 107 dimensional for the ant and 163 dimensional for the
OP3, respectively.
Task reward The task reward consists of two parts, a dense reward term that corresponds to how well the target pose is matched and a sparse reward that is added if the target pose is ‘achieved’:
[Equation]
where we use r threshold = 0 .
65 , N = 1 , S = 50 for the OP3 and r threshold = 0 .
8 , N = 10 , S = 50 for the ant.
The rewards are similarly robot speciﬁc and will be described below.
The rewards as well as the tolerance r threshold and N were tuned to give visually good behaviors with a large-scale on-policy agent.
For the dense reward we use rewards of the following form:
[Equation]
where d denotes the center of mass distance between the agent and the target, r pose is a reward that depends chieﬂy on the robots joint conﬁguration but less on the relative position to the target pose and m robot is a robot speciﬁc multiplier.
We use m robot = 4 for the OP3 and m robot
= 1 for the ant.
Ant pose reward For the ant we use the following pose reward:
[Equation]
where d root denotes the relative root distance.
r pos and r quat are terms penalizing deviations from the target pose in terms of the positions and orientations of the individual MuJoCo bodies that make up the Ant:
[Equation]
where p bodies and p refbodies are the body positions of the ant and the reference, respectively, and q bodies and p refbodies are the body positions of the ant and the reference, respectively,
OP3 pose reward For the OP3 we use the following pose reward (similar to Peng et al.
Hasenclever et al.
with terms penalizing deviations in terms of the center of mass, the joint angle velocities, the end effector positions and the joint orientations:
[Equation]
The ﬁrst term r com penalizes deviations from the centre of mass:
[Equation]
[Equation]
where p com and p refcom are the positions of the centre of mass of the simulated character and the mocap reference, respectively.
The second term r vel penalizes deviations from the reference joint angle velocities:
[Equation] [Equation]
where q vel and q refvel are the joint angle velocities of the simulated character and the mocap reference, respectively.
The third term r app penalizes deviations from the reference end effector positions:
[Equation] [Equation]
where p app and p refapp are the end effector positions of the agent and the target pose, respectively.
Finally, r quat penalizes deviations from the reference in terms of the quaternions describing the body orientations:
[Equation] [Equation]
where (cid:23) denotes quaternion differences and q bodies and q refbodies are the joint quaternions of the agent and the target pose, respectively.
Planner rewards The reward terms above cannot be computed solely from observations available to the agent.
In principle it would be possible to instead use reward terms that are only based on observations.
However, in practice, we use planner rewards that are correlated with but different from the task reward.
For the OP3 we use the following reward for planning:
[Equation]
[Equation]
where J is the number of joints.
For the Ant we use a similar reward with slightly different lengthscales for the planning reward:
[Equation]
[Equation]
[Equation]
Target pose sampling
At the beginning of an episode and whenever a target pose is achieved we sample a new target.
The target pose is sampled uniformly from a reference set of target poses.
For the ant this set is derived from expert policies that have been trained on a standard go-to-target task from dm_control (Tunyasuvunakool et al., 2020).
For the OP3 we use target poses from the CMU mocap database (cmu).
To determine the position of the new target relative to the agent we sample both a heading relative to current agent heading (sampled uniformly between − 120 ◦ and 120 ◦ ) and a random distance.
We also sample the orientation of the target randomly relative to the current agent orientation (heading shift sampled uniformly between − 120 ◦ and 120 ◦ ).
To make this challenging task a bit easier we use what we call an intra-episode curriculum to determine the distance of the target pose from the agent.
We linearly scale up the random distance with each achieved target (up to some maximum value).
Concretely, let U [ a, b ] be a uniform distribution and let M be the number of curriculum phases.
Then after having achieved m target poses, the distance of the next target pose is drawn from U [min( m/M, 1) a, min( m/M, 1) b ] .
We use
M = 4 and a = 0 .
5 , b = 5 meters for the ant and a = 0 .
05 , b = 1 .
meters for the OP3.
Motivating the go-to-target-pose task
The go-to-target-pose task can be motivated and understood from a few distinct perspectives.
Firstly, it arises naturally as a temporally abstract version of a motion tracking task, and indeed we build our task as an extension of existing motion tracking infrastructure (Hasenclever et al., 2020) that is available as part of dm_control (Tunyasuvunakool et al., 2020).
In a tracking task the target pose and reward change every timestep such that the agent is rewarded for producing behavior that tracks a time-varying reference motion.
In the GTTP task, the instruction and reward deﬁnitions can be similar to the tracking task; however, the GTTP task involves a ﬁxed target pose for a number of timesteps until the target is achieved.
Insofar as the policy that solves a tracking task is an inverse dynamics model ( s, s (cid:48)
→ a ), the policy that solves the GTTP task is a temporally abstract inverse model.
From another perspective, the GTTP task is the essential self-rearrangement task that can be performed by any avatar in an empty environment.
Batra et al.
(2020) argue that a conceptual uniﬁcation of goal-directed navigation tasks, visual servoing, and object manipulation tasks are that they all are rearrangement tasks that can be standardized as challenges for embodied control.
Rearrangement tasks vary along multiple dimensions, including temporal abstraction and degree of speciﬁcation.
A motion tracking task is essentially a self-rearrangement challenge that involves a target pose being fully speciﬁed at every timestep.
The GTTP task involves a target pose being speciﬁed with an unspeciﬁed offset and the policy must handle the temporal abstraction by closing in on the target pose.
While both of these tasks fully specify the desired pose, one can alternatively consider versions of the tasks where the desired pose is not fully speciﬁed, or even further where the desired pose is about external objects rather than the pose of the body.
To surmount the issues related to temporal abstraction, policies may, implicitly or explicitly, break problems down into subgoals.
Curiously, goal-directed policies (or inverse models) also may themselves serve as effective abstractions for achieving subgoals, playing the role of lower-level controllers insofar as they can abstract away the capabilities required to execute movements to a subgoal.
We believe the GTTP task should be particularly amenable to model-based methods: it combines a high-dimensional control problem with a diverse goal distribution.
This makes it hard to solve quickly with model-free methods.
However, since the dynamics are shared between all goal poses, a dynamics models should be beneﬁcial in leveraging the common structure.
B P LANNERS
We consider two sample based planners, primarily a Sequential Monte Carlo based non-iterative planner (SMC)
(Piché et al., 2019; Gordon et al., 1993) and the Cross-Entropy Method (CEM) (Botev et al., 2013).
Our SMC and CEM planners are shown in algorithm 2 and 3, respectively.
We brieﬂy describe the planners and discuss how they can be warm-started with a provided action proposal distribution.
SMC planner:
The SMC planner (Alg. 2) is a non-iterative sample-based planner that maintains a number of particles S corresponding to different model rollouts.
All particles are initialized to start at the initial state s 0 .
At each step of the rollout (upto a horizon of H ), an action is sampled from a proposal for each particle; this proposal can be either a ﬁxed distribution or the learned proposal.
These actions are then rolled out through the model to compute the next state, reward and optionally, a value estimate.
The particles are then resampled according to the weighted exponentiated reward/advantage with a temperature parameter τ .
The rollout is performed upto a horizon of H and the ﬁrst action from a randomly chosen particle (amongst the surviving particles at the ﬁnal timestep) is chosen as the action to be executed.
Unlike CEM, SMC samples actions from the proposal at each step within the rollout and is therefore more suited to warm-starting using a learned proposal; we choose it as our planner primarily for this reason.
1: Given: state s 0 , proposal π m
, reward r , planning horizon H , number of samples S , planner temperature and optionally value function V ψ 2: {
= s
= //
Initialise state 3:
Let x ( i ) 0:
t = { ( s ( i ) 0 , a ( i ) 0 ) , . . .
, ( s ( i ) t , a ( i ) t ) , s ( i ) t +1 } denote the i th particle up to time t 4: for t = 0 . . .
H do 5: { a ( i ) t ∼ π θ ( ·| s ( i ) t ) } S i = //
Sample actions 6: { s ( i ) t +
= m φ ( s ( i ) t , a ( i ) t ) } Si
= //
Take model step.
7: { r ( i ) t = r
( s ( i ) t , a ( i ) t , s ( i ) t +1 ) }
Si = //
Evaluate rewards.
8: calculate advantage
A ( i ) t based on rewards and value function V ψ
Resampling 9: w ( i ) t ∝ exp (cid:16)
A
( i ) t /τ (cid:17) //
Resampling weights proportional to exponentiated advantage 10: i 1 , i 2 , . . .
i S ∼ Categorical ( { w ( i ) t } Si =1 ) //
Draw resampling indices 11: x ( k ) 0: t ← x ( i k ) 0: t ∀ k = 1 . . .
S // resample trajectories 12: end for 13: i ∼ U{ 1 , . . .
, S } //
Sample action index 14: return a ( i )//
Return ﬁrst action from a random particle
CEM planner: The CEM planner (Alg. 3) is an iterative sample-based planner maintains a distribution over action sequences that is usually parameterized as a Gaussian with mean µ and standard deviation σ .
At the start of planning, the mean µ is initialized to an open-loop sequence of actions sampled from a proposal distribution (either ﬁxed or learned) and the standard deviation σ is initialized to σ init .
In each iteration, S trajectories are sampled from this distribution and rolled out through the model and their returns are computed.
The top E fraction of trajectories, ranked by the return, are retained and their mean µ elite and standard deviation σ elite are computed.
These are used to updated the mean µ and standard deviation σ via an exponential average (with weights α mean
and α std ).
This procedure is repeated for I iterations and after the ﬁnal iteration, the ﬁrst action from µ is executed in the environment.
Unlike SMC, CEM uses the proposal only for plan initialization at the ﬁrst iteration and uses the distribution N ( µ, σ ) for further trajectory samples.
C M ODELS
Successfully applying model predictive control requires learning good dynamics models.
In this paper, we train single-step feedforward dynamics models m φ that take in the current state s t and action a t and predict the next state s t +
= m φ ( s t , a t ) .
This dynamics model has both learned and hand-designed components; the task-agnostic proprioceptive observations are predicted via “black-box” neural networks while any non-proprioceptive – task-speciﬁc – observations (e.g. relative pose of the target) are calculated in closed form from the predicted proprioceptive observations in combination with a known kinematic model of the robot.
Fig.
5 (top) shows the network architecture of the learned components of the predictive dynamics model ( m φ ).
We use a deterministic, feed-forward MLP per proprioceptive observation i.e. one each for joint angles, joint velocities, linear velocities and so on.
This MLP takes as input all the proprioceptive observations and the action from the current timestep and predicts a delta change in the observation being predicted; this delta change is scaled by an embodiment speciﬁc timestep
Algorithm
Given: state s 0 , proposal π θ , model m φ , reward r , planning horizon H , number of samples S , elite fraction E , noise standard deviation σ init , number of iterations I , and optionally value function V ψ //
Rollout proposal distribution using the model to initialize the plan.
( s 0 , a 0 , s 1 , . .
, s H ) ← rollout_with_proposal ( m φ , π θ , H ) µ ← [ a 0 , a 1 , . .
. , a H ] // initial plan (H × action dimension) σ ← σ init for i = 1 . . .
I do for k = 1 . . .
S do p k ∼ N ( µ, σ ) //
Sample candidate actions.
// Evaluate candidate action sequences open loop according to the model and compute associated returns.
r k ← evaluate_actions ( m φ , p k , H, r, V ψ ) end for
Rank candidate sequences by reward and retain top E fraction.
Compute mean µ elite and per-dim standard deviation σ elite based on the retained elite sequences.
µ ← (1 − α mean ) µ + α mean µ elite // Update mean; α mean = 0 .σ ← (
− // Update standard deviation; α = 0 .end for return ﬁrst action in µ ( dt = 0 .
for Ant, dt = 0 .
03 for OP3) and added to the observation at the current timestep to get the prediction at the next timestep:
[Equation]
17) where ˆ s kt +1 denotes the k-th predicted observation at time t + 1 , s kt is the k-th observation at time t and m kφ is the k-th MLP.
The only exception to this is for joint angle predictions, where we make a delta prediction on top of the joint velocity predictions to get the predicted joint angles at the next timestep:
[Equation]
[Equation]
Here, q t and ˆ q t +1 denote the observed and predicted joint angles at t and t + 1 respectively, m qφ denotes the MLP that predicts the delta change in joint angles and ˆ˙ q t +1 is the predicted joint velocities at t
+ 1 .
In effect, we implement an Euler integration step for joint angles based on the predicted joint velocities ˆ˙ q t +1 ) and a correction term from the MLP m qφ ( s t , a t ) .
In the case of the go-to-target-pose task, we have an additional learned “forward kinematics” model that predicts the 3D positions and orientations (represented as quaternions) of different parts of the embodiment.
The architecture of this model is shown in Fig.
5 (bottom).
This network takes in only the joint angles of the robot and predicts the 3D positions and orientations of the robot’s bodies via a feed-forward, deterministic MLP; as before, there are two separate MLPs, one each for the position and orientation predictions.
C. T RAINING
We train our models alongside the policy and critic using data collected from the MPC loop in
Alg. 1 that is saved in the replay buffer.
We sample a batch of trajectories ( s 1: T , a 1: T − 1 , r 1: T ) from this replay buffer, where T is the trajectory length (set to 10 in all our experiments).
This batch of trajectories is also used for policy and critic learning.
Given a trajectory of length T , we perform an open-loop rollout with the model using the initial state s 1 and the entire sequence of actions a 1: T − 1 :
[Equation]
[Equation]
[Figure] [Figure]
where the ﬁrst prediction ˆ s 2 uses the true state s 1 and subsequent predictions use the predicted state from the previous timesteps.
We then compute a mean squared error between the predictions and the true states:
[Equation]
where ˆ s i and s i denote the predicted and true state respectively.
This objective J m is minimized for training the model parameters φ ; we use an Adam optimizer (Kingma & Ba, 2014) with a ﬁxed learning rate of 1e-4 for training, and use a sequence length of T = 10 .
We use a similar training procedure for the forward kinematics model (used only for the GTTP task) but instead use the joint angles from the true states s 1: T as inputs to the network as opposed to predictions.
We compute a squared error between the true and predicted 3D positions and a quaternion difference for the between the true and predicted orientations; these are summed together and averaged across time to get the ﬁnal objective for training.
As before, we minimize this objective using an Adam optimizer (Kingma & Ba, 2014) with a ﬁxed learning rate of 1e-4 and a sequence length of 10 for training the model parameters.
C. I NTEGRATING MODEL PREDICTIONS FOR PLANNING
As mentioned earlier, only proprioceptive observations are predicted by the learned models but the policy and critic depend on both proprioceptive and task-speciﬁc observations (see e.g. Fig. 6).
Consequently, we need a way to generate future task-speciﬁc observations given an initial state and action sequence for multi-step planning to work.
We do this by integrating task-speciﬁc observations in closed form based on the predictions from the learned dynamics (and kinematics) models.
As all our observations (proprio and task-speciﬁc) are speciﬁed in an egocentric frame of reference we do not need any external information for this integration other than the observations at the current timestep and the predictions from our learned models.
We brieﬂy describe the integration process for all our tasks next.
For the walking tasks, the only task-speciﬁc observation to be integrated is the relative direction of the target.
Since the angular velocity of the robot is a proprioceptive observation (and the timestep dt is known) we can generate a change in orientation by Euler integrating this velocity from the unit quaternion.
By rotating the relative target direction with this change in orientation we can obtain the predicted relative target direction at t + 1 .
For the GTTP tasks, there are several task-speciﬁc observations including the relative joint angle differences between the robot’s current conﬁguration and the target, the relative root position and the relative 3D positions and quaternions of the target bodies w.r.t to the current conﬁguration.
The integration of relative joint angles is straightforward: we can subtract the Euler integrated joint velocities to the relative joint angles at t to get the relative joint angles at t + 1 .
Similarly, we can integrate the relative root position by adding the translation in the egocentric frame of reference (obtained by Euler integration of the linear velocities).
The integration of the relative body positions and orientations is more involved, and uses a combination of the predicted translation, rotation (obtained by Euler integration of linear/angular velocities) and the observed & predicted positions and orientations of the robot (at t and t + 1 respectively); as before, this does not need any external information not available as part of the egocentric observations.
The procedure for predicting the full next state ˆ s t +1 given the current state s t and a t is as follows:
Predict all the proprioceptive observations via the learned dynamics model:
ˆ s kt +
= m φ ( s t , a t ); ∀ k . 2.
(Only for the GTTP task) Predict the body positions and orientations at t +1 via the predicted joint angles at t + 1 (from step 1 above).
Integrate the task-speciﬁc observations based on the predictions from the learned dynamics (and optionally, the kinematics) models.
This procedure is used within the PLANNER subroutine (see Alg. 1) for generating future states via multi-step rollouts.
D MPO
Maximum A-Posteriori Policy Optimization (MPO) (Abdolmaleki et al., 2018b) is a continuous control RL algorithm that performs an expectation minimization form of policy iteration.
There are two steps comprising policy evaluation and policy improvement .
The policy evaluation step receives as input a policy π k and evaluates and action-value function Q π k θ ( s, a ) by minimizing the squared TD error:
[Equation]
where ˆ θ denotes the parameters of a target network (Mnih et al., 2015) that are periodically updated from θ and d π denotes the distribution over s t , a t , s t +1 , a t +1 induced by π .
In practice we use a replay-buffer of samples in order to perform the policy evaluation step.
The second step, policy improvement consists of optimizing the objective ¯ J ( s, π ) =
E π
[ Q π k θ ( s, a )] for states s drawn from a state distribution µ ( s ) .
In practice the state distribution we draw samples from is the replay buffer.
By improving ¯ J in all states s , we improve our objective.
To do so, a two step procedure is performed.
First, we construct a non-parametric estimate q such that ¯ J ( s, q ) ≥ ¯ J ( s, π k ) .
This is done by maximizing ¯ J ( s, q ) while ensuring that the solution, locally, stays close to the current policy π k ; i.e. E µ ( s ) [ KL ( q ( . | s )
||
π k ( .
| s ))] < (cid:15) .
This optimization has a closed form solution given as
[Equation]
where η is a temperature parameter that can be computed by minimizing a convex dual function (Abdolmaleki et al., 2018a).
Second, we project this non-parametric representation back onto a parameterized policy by solving the optimization problem
π k +1 =
arg min π E µ ( s )
[ KL ( q ( a | s ) ||
π ( a | s )] , where π k +1 is the new and improved policy and where one typically employs additional regularization (Abdolmaleki et al., 2018b).
Note that this amounts to supervised learning with samples drawn from q ( a | s ) ; see Abdolmaleki et al.
(2018b)
for details.
D. C HOICE OF MPO
We believe that the choice of the underlying learning algorithm (MPO) is not crucial for our results.
In fact, we expect that using any state-of-the-art off-policy algorithm such as Soft Actor Critic (SAC) (Haarnoja et al., 2018) should lead to similar results across the tasks considered.
We choose MPO as our learning algorithm for two reasons.
First, performance-wise, it is comparable with state-of-the-art off-policy actor-critic algorithms and has been thoroughly evaluated across a variety of continuous control settings (Hoffman et al., 2020).
Second, the behavioral cloning objective (Eqn. 1) we add is compatible with the MPO objective (Eqn. 3) as they are both log-likelihood based objectives with a similar scale; setting a trade-off between the two is straightforward.
Adding the BC loss to MPO has also been explored in prior work (Abdolmaleki et al., 2021) with improvements in learning across different settings.
E E XPERIMENTAL SETUP
We work with a distributed set-up with 64 actors interacting with the environment, and a single learner that trains the policy, critic and model.
The actors collect data via the MPC loop described in
Alg. 1, where the probability of planning p plan controls the tradeoff between executing the action from either the planner or the learned proposal π θ .
This data is then stored in a replay buffer in the form of trajectories of state, action, reward triplets ( s 1: T , a 1: T − 1 , r 1: T ).
The learner samples batches from the replay buffer (we use a batch size of 256, sequence length T = 10 and a replay buffer size of 1 million trajectories in all our experiments) and runs a learning update where the parameters of the policy, critic and model are updated in tandem.
To study data efﬁciency in a data-limited regime, we additionally limit the number of actor steps per learner update step, i.e. the learner is blocked if the actor has not sampled enough data and the actor is blocked if the learner has not performed enough parameter updates.
This is implemented using the reverb framework (Cassirer et al., 2021).
We did a sweep over this “rate limit” parameter and found the best settings to be 8 actor steps per parameter update for the walking tasks and 16 actor steps per parameter update for the GTTP tasks (increasing the values further led to signiﬁcantly reduced performance).
This is held ﬁxed for all from scratch and transfer experiments with all variants.
Alg. 1 describes the overall agent including the acting and learning loops which run in a distributed fashion with 64 actors and a single learner as mentioned earlier.
The data from the actors which run
MPC, warm-started with the learned policy (see MPC loop of Alg. 1) is added to the replay buffer; the learner asynchronously samples batches of data from this buffer and updates the parameters of the policy, critic, model and optionally, the task-agnostic proposal in an iterative fashion.
E. P OLICY AND C RITIC TRAINING
Our policy architecture is presented in Fig.
6 (top left, (a))
; the policy takes as input both the proprioceptive and task-speciﬁc observations and outputs the parameters of a Gaussian distribution from which actions are sampled.
We use a distributional critic (Fig.
6 - bottom left, (b)) that takes as input both the observations and actions and returns a discrete distribution (101 bins) that parameterizes a histogram of Q-values (from 0-300), similar to the one proposed in Hoffman et al.
Given a batch sampled from the replay buffer we compute Q targets using Retrace (Munos et al., 2016) with a terminal value bootstrapped from a target critic network (both the target policy and critic network parameters are updated once every 200 steps).
The use of a distributional critic makes value backups hard; we circumvent this by bootstrapping using the mean Q-value.
We then convert the value targets into a two-hot representation (only two bins in the histogram are non-zero, rest are zero), which we ﬁt with our critic by minimizing the softmax cross-entropy loss between the critic
[Figure]
predictions and the two-hot targets.
As discussed in the main text (Sec. 4.3 and Eqn. 4), we use a weighted combination of the model-free MPO objective and a Behavioral Cloning objective (BC) for training the policy and critic.
We use separate Adam optimizers (Kingma & Ba, 2014) with learning rates of 3e-4 for training the policy and critic, and use an Adam optimizer with a learning rate of 1eto optimize the temperature and dual variables in the MPO algorithm (Abdolmaleki et al., 2018b).
E. T RAINING A TASK - AGNOSTIC PROPOSAL
In addition to the policy and critic, we also train a task-agnostic proposal (Fig.
6 - center right, (c))
that takes only proprioceptive observations as input and returns a Gaussian distribution for sampling actions.
Since this proposal lacks any task-speciﬁc information it can only capture average behavior (e.g. walking randomly in the GTTP task, as opposed to directed walking to a target).
Since proprioceptive observations are the same across tasks for a given embodiment this proposal can be transferred freely across tasks; we used this proposal for the initial planner experiments (Sec. 5.1) and the proposal transfer experiments (Sec. 5.4).
We train this task-agnostic proposal alongside our policy and critic using a Behavioral Cloning (BC) objective only.
We sample batches from the replay buffer (which have both good and bad data) and ﬁt the policy by maximizing the log-prob of executed actions under the proposal distribution (similar to Eqn. 1).
We use an Adam optimizer with a learning rate of 1e-4 for training this task-agnostic proposal.
We present a few videos showing the behavior learned by the task-agnostic proposals in the supplementary material and the website.
Interestingly, they likely because tends to would wander away gives not
Since close
F F URTHER EXPERIMENTAL RESULTS F. P LANNER RESULTS
For our result with pre-trained models and proposal (Table 1) we ran a hyper parameter sweep over the common planner parameters:
• Planning horizon ( H ): 5, 10, 15, 20, 25,• Number of samples ( S ): 50, 100, 200, 500, 1000,
For the Gaussian proposal we also swept over the standard deviation of the Gaussian (0.5, 1.0 and 2.0).
In addition we also swept over planner-speciﬁc parameters.
For SMC:
• Planner temperature ( τ ): 0.001, 0.01, 0.
For CEM:
• Elite fraction ( E ): 0.15, 0.
• Number of iterations ( I ): 1, 2,• Noise standard deviation ( σ init ): 0.1, 0.25, 0.5, 1.
Table 1 presents the best results for each planner, which use a planning horizon of H = 30 and number of samples S = 2000 .
We note here that best performance with CEM is obtained with I > 1 , which effectively means that CEM uses a signiﬁcantly larger computational budget than our SMC planner which is non-iterative; in spite of this SMC is still quite competitive with CEM across all tasks as can be seen from the results in Table 1.
In a further set of experiments we evaluated the performance of SMC with a wider range of samples and ground truth dynamics on the Ant walking, Ant GTTP and OP3 GTTP tasks.
The results are shown in Figures 7, 8, 9.
In the simpler, lower dimensional Ant domains, increasing the number of samples used by the planner improves performance, irrespective of the proposal distribution used for planning.
In the higher-dimensional OP3 GTTP task the results are qualitatively different.
When using longer horizons and more samples from a pre-trained, but goal-agnostic, proposal the performance increases.
However even with 25600 samples from a standard Gaussian proposal the planner fails to make any progress.
This is due to the difﬁculty of sample-based optimization in high-dimensional spaces, which further validates our choice of the use of a learned proposal distribution for the planner.
Overall the planner results show the challenges of planning in challenging high dimensional tasks even with access to ground truth models, rewards and large computational budgets.
Using a proposal is crucial for progress on these domains.
While the best planner only results use a high computational budget as presented above, we use SMC with S = 250 samples and planning horizon 10 for all our learning experiments for improved learning speed (wall-clock time) and to ensure that our computational budgets are realistic with respect to real-world considerations (assuming a 20Hz run-rate of the actor).
Interestingly, this setting does not result in performance reduction compared to those with more compute as we show in ablation experiments (Fig. 10 and 11).
Details on other planner hyper-parameters for individual experiments are detailed below.
F. F ROM SCRATCH RESULTS
In initial experiments we determined the number of actor steps per learner step (using the MPO baseline).
We tried 4, 8 and 16 actor steps per learner step and found that the performance degraded when using lower than 8 actor steps per learner step for the walking tasks and 16 actor steps per learner step for the GTTP experiments.
We used these settings for all subsequent experiments.
We tuned each algorithmic variant independently per task, running sweeps over MPO hyper parameters and, where applicable, BC objective weight β and planner temperature τ .
We show the best results of each variant in the ﬁgures.
For MPO hyper-parameters we ran sweeps over the constraint parameter (cid:15) = 0 .
5 in Equation 2.
Additionally we sweep over several different settings for the additional trust-region constraint in the policy mentioned to in the discussion of Equation 3.
In practice, MPO constraints the mean and variance of the policy separately, with constraint parameters
(cid:15) µ and (cid:15) Σ .
This helps avoid premature convergence in some cases.
See Abdolmaleki et al.
(2018b)
for details.
We used (cid:15) Σ = 10 − 5 and swept over 3 values from (cid:15) µ = 5 · 10 − 4 , 10 − 4 , 5 · 10 − 3 , 10 − depending on the task.
For algorithmic variants involving BC we ran hyper-parameter sweeps varying the BC objective weight β
= 0 .
For algorithmic variants involving MPC we ran sweeps over the planner temperature τ = 0 .
1 ( τ = 0 .
01 worked best for all tasks except the Ant walking tasks where τ = 0 .
1 worked best).
We use SMC with S = 250 samples and planning horizon of H = 10 for all our experiments with
MPC.
We determined these to be the best planner settings based on a hyperparameter sweep over planner parameters when learning from scratch with MPC on the OP3 GTTP task; these were also chosen with real-world considerations to ensure that planning could be run at 20Hz on practical systems.
Fig.
10 presents the results when training MPC+MPO+BC with different number of samples S for the SMC planner with a ﬁxed horizon of 10, and the best temperature setting ( τ = 0 . 01 ) from the from scratch experiment (see Fig. 2).
Increasing the number of samples or setting p plan = 1 . 0 does not make a big difference.
Fig.
11 shows a related experiment, but with a sweep over the planning horizon H instead with S = 250 samples.
Once again, there is little variation in performance with a longer or shorter horizon.
We also did a quick sweep over bootstrapping with the value function during planning, training
MPC+MPO+BC from scratch on several tasks with and without bootstrapping.
When bootstrapping from the value function, we use the learned critic and average the Q values from the given state and 10 actions sampled from the policy at the given state to compute the value V .
Fig.
12 presents this ablation.
In general, boostrapping with the value function helps on easier tasks such as walking (Fig.
12, bottom row)
but hurts performance on the harder GTTP task (Fig.
12, top row) where learning a good value function can be difﬁcult.
We use the best settings for each task in our experiments (no bootstrapping for the GTTP tasks and bootstrapping for the walking tasks).
We additionally present results for learning from scratch on the walking forward tasks in Fig.
As with the backward walking results, MPC+MPO does better than MPO early on during learning and the addition of the BC objective improves performance compared to the non-BC counterpart especially on the OP3 task.
All variants converge to similar asymptotic performance.
F. T RANSFER RESULTS
Model transfer: As discussed in the main text we test two variants of model transfer where the transferred model is kept frozen on the target task ( MPC+MPO+BC+Frozen ), or ﬁnetuned on the target task ( MPC+MPO+BC+Finetune ), along with a baseline where we learn the model from scratch.
For all the model transfer experiments we transfer the models from the best performing
MPC+MPO+BC agent on the source task (and its hyperparameter settings).
For all transfer settings except GTTP to GTTP we transfer only the forward dynamics model; for GTTP to GTTP we transfer both the forward dynamics and kinematics models.
For ﬁnetuning we use an Adam optimizer with a learning rate of 1e-4 for both models (similar to from scratch training).
Proposal transfer: In the proposal transfer experiments we transfer the task-agnostic proposal which depends only on proprioceptive observations to the target task.
We use a mixture of this reloaded proposal from the source task and the learned amortized policy (which has access to all observations) on the target task as the proposal distribution ( π θ ) for planning in algorithm 1.
The mixture weight of this reloaded proposal is annealed linearly from 1 to 0 in a ﬁxed number of learning steps M .
At the start of learning on the target task we sample exclusively from the reloaded task-agnostic proposal on the actor and as learning progresses we sample less and less from this proposal; when the mixture weight reaches 0 we stop sampling from the reloaded proposal and revert to using only the amortized policy on the target task as the proposal distribution on the actor.
We tuned the hyper-parameter M (in terms of learner steps) which controls the slope of this annealing of the mixture probability on all our transfer tasks and presented results from the best setting of M .
We choose all other parameters based on the best from scratch results on the target task, with the BC objective added ( MPC+MPO+BC ).
The per-task sweeps are:
• OP3 GTTP to GTTP:
M = 1.25e5, 2.5e5, 5e5, 1e6, 2e • Ant GTTP to GTTP:
M = 15625, 31250, 62500, 125000,• OP3 Forward Walk to GTTP: M = 1.25e5, 2.5e5, 5e5, 1e6, 2e
• Ant Forward Walk to GTTP:
M = 15625, 31250, 62500, 125000,• OP3 GTTP to Forward/Backward Walk: M = 15625, 31250, 62500,• Ant GTTP to Forward/Backward Walk: M = 5000, 10000, 20000,
Additional transfer results: Fig.
14 shows the performance of the learned amortized policy for the different transfer variants on several of our transfer tasks.
This is the counterpart to Fig.
3 which shows the actor performance on the same tasks; in general the amortized policy is able to match or exceed the performance of the actor for all these settings due to the addition of the BC objective.
Finally, Fig. 15 and Fig.
16 show the performance of the different variants when models and/or proposals are trained on the GTTP task and transferred to the Forward Walking task.
As in previous transfer results (Sec. 5.3), model transfer (solid lines) does not lead to signiﬁcant improvements w.r.t learning speed or asymptotic performance.
Similar to the results in Sec. 5.4, proposal transfer (dotted lines) can help speed up learning early on but can also potentially hurt asymptotic performance especially on the more complex OP3 tasks.
F. M ODEL A BLATION - U SING A S TOCHASTIC M ODEL E NSEMBLE
In our work we primarily used deterministic dynamics models (Sec. C), which is in contrast with some of the recent literature on model-based reinforcement learning that advocates for ensembles of stochastic models (e.g. PETS (Chua et al., 2018)).
In this section we ablate this design choice.
In this ablation, we use the same architecture as for our deterministic models but with Gaussian output distributions (parameterized by a mean and log-variance) instead of single predictions.
Similar to
PETS (Chua et al., 2018), we bound the predicted variance to ensure stability and use a single step negative log-likelihood loss for training the stochastic models instead of a multi-step loss as for the deterministic models.
For ensembles, we initialize each member randomly and train each ensemble member on a randomly selected subset (of half the batchsize) of each batch to encourage diversity of model predictions.
For planning with the stochastic ensemble, we split the different particles in SMC equally across the ensemble members; each particle uses the same ensemble member throughout entire rollout (similar to the TS ∞ approach from PETS (Chua et al., 2018)).
This approach ensures that we keep the overall computational budget for planning the same irrespective of the ensemble size.
The results of this ablation study are presented in Fig.
17 and Fig.
18 where we compare the performance of MPC+MPO and MPC+MPO+BC when learning from scratch on all our tasks.
Fig.shows the performance of the actor; using ensembles of stochastic models does not lead to a signiﬁcant difference in performance on the simpler Ant tasks and provides a small advantage on the OP3 tasks.
This advantage is reduced when looking at the performance of the amortized proposal in Fig.
18; only the performance of MPC+MPO on the OP3 GTTP task is slightly improved through the use of the stochastic ensemble.
All results used an ensemble size of 3 except the results on the Ant GTTP task which uses a single stochastic model.
We additionally ran an ablation on the size of the stochastic ensemble, considering ensembles of 1, and 5 models.
Fig.
19 shows both the actor and amortized proposal performance of MPC+MPO and
MPC+MPO+BC on the OP3 GTTP task.
We observed no meaningful difference between 1, 3 andensemble members with the same computational budget for planning; similar results were observed for all other tasks considered in this paper.
Overall, these results validate our choice of deterministic models showing that a single well structured deterministic model can perform as well as stochastic ensemble across challenging locomotion tasks involving high-dimensional, non-linear contact dynamics
F. R ESULTS ON D EEP M IND C ONTROL S UITE
In addition to the tasks presented in the paper we also evaluated our approach on externally published and open-sourced locomotion tasks which are a part of the DeepMind Control Suite (Tassa et al., 2018).
Speciﬁcally we tested our approach on three locomotion tasks from the DeepMind Control
Suite: stand , walk and run , with two embodiments: a 6-DOF walker (18-D state space) and a 21-DOF humanoid (54-D state space) forming a total of six tasks.
We ﬁrst tuned MPO (Abdolmaleki et al., 2018b) to match the performance of externally published benchmarks in the literature (Hoffman et al., 2020) across all these tasks.
Next, we evaluated our approach of combining MPC with MPO (using a deterministic model) on these tasks; Fig 20 shows the performance of all the algorithmic variants in Sec.
5.2 on these six tasks.
All variants perform similarly on the simpler walker embodiment; there is no advantage to using MPC across all three tasks for this embodiment.
On the humanoid , the MPC variants slightly outperform the model-free
MPO variants, especially on the challenging humanoid-run task.
Overall though, the results are broadly similar to those presented in the main text and further lend strength to the central message of the paper; while model-based methods can perform really well on high-dimensional continuous control tasks, well-tuned model-free methods are strong baselines.
Only on the most challenging tasks and, especially, in multi-goal/multi-task settings, do model-based methods substantially outperform their model-free counterparts.