MoDem: Accelerating Visual Model-Based
Reinforcement Learning with Demonstrations

Reinforcement Learning (RL) provides a principled and complete abstraction for training agents in unknown environments. However, poor sample efficiency of existing algorithms prevent their applicability for real-world tasks like object manipulation with robots. This is further exacerbated in visuo-motor control tasks which present both the challenges of visual representation learning as well as motor control. Model-based RL (MBRL) can in principle (Brafman & Tennenholtz, 2002) improve the sample efficiency of RL by concurrently learning a world model and policy (Ha & Schmidhuber, 2018; Schrittwieser et al., 2020; Hafner et al., 2020; Hansen et al., 2022). The use of imaginary rollouts from the learned model can reduce the need for real environment interactions, and thus improve sample efficiency. However, a series of practical challenges like the difficulty of exploration, the need for shaped rewards, and the need for a high-quality visual representation, prevent MBRL from realizing its full potential. In this work, we seek to overcome these challenges from a practical standpoint, and we propose to do so by using expert demonstrations to accelerate MBRL.

Expert demonstrations for visuo-motor control tasks can be collected using human teleoperation, kinesthetic teaching, or scripted policies. While these demonstrations provide direct supervision to learn complex behaviors, they are hard to collect in large quantities due to human costs and the degree of expertise needed (Baker et al., 2022). However, even a small number of expert demonstrations can significantly accelerate RL by circumventing challenges related to exploration. Prior works have studied this in the context of model-free RL (MFRL) algorithms (Rajeswaran et al., 2018; Shah & Kumar, 2021; Zhan et al., 2020). In this work, we propose a new framework to accelerate model-based RL algorithms with demonstrations. On a suite of challenging visuo-motor control tasks, we find that our method can train policies that are approx. $160\%-250\%$ more successful than prior state-of-the-art (SOTA) baselines.

Off-Policy RL algorithms (Sutton & Barto, 1998) – both model-based and model-free – can in principle admit any dataset in the replay buffer. Consequently, it is tempting to naïvely append demonstrations to the replay buffer of an agent. However, we show that this is a poor choice (see Section 4), since the agent still starts with a random policy and must slowly incorporate the behavioral priors inherent in the demonstrations while learning in the environment. Simply initializing the policy by behavior cloning (Pomerleau, 1988) the demonstrations is also insufficient. Any future learning of the policy is directly impacted by the quality of world model and/or critic – training of which requires sufficiently exploratory datasets. To circumvent these challenges and enable stable and monotonic, yet sample-efficient learning, we propose Model-based Reinforcement Learning with Demonstrations (MoDem), a three-phase framework for visual model-based RL using only a handful of demonstrations. Our framework is summarized in Figure 2 and consists of:

• •

  Phase 1: Policy pretraining, where the visual representation and policy are pretrained on the demonstration dataset via behavior cloning (BC). While this pretraining by itself does not produce successful policies, it provides a strong prior through initialization.

• •

  Phase 2: Seeding, where the pretrained policy, with added exploration, is used to collect a small dataset from the environment. This dataset is used to pretrain the world model and critic. Empirically, data collected by the pretrained policy is far more useful for model and critic learning than random policies used in prior work, and is key to the success of our work as it ensures that the world model and critic benefit from the inductive biases provided by demonstrations. Without this phase, interactive learning can quickly cause policy collapse after the first few iterations of training, consequently erasing the benefits of policy pretraining.

• •

  Phase 3: Finetuning with interactive learning, where we interleave policy learning using synthetic rollouts and world model learning using data from all three phases including fresh environment interactions. Crucially, we aggressively oversample demonstration data during world model learning, and regularize with data augmentation in all phases.

Our Contributions.  Our primary contribution in this work is the development of MoDem, which we evaluate on $\mathbf{18}$ challenging visual manipulation tasks from Adroit (Rajeswaran et al., 2018) and Meta-World (Yu et al., 2019) suites with only sparse rewards, as well as locomotion tasks from DMControl (Tassa et al., 2018) that use dense rewards. Measured in terms of policy success after $\mathbf{100K}$ interaction steps (and using just 5 demonstrations), MoDem achieves $160\%-250\%$ higher relative success and $38\%-50\%$ higher absolute success compared to strong baselines. Through extensive empirical evaluations, we also elucidate the importance of each phase of MoDem, as well as the role of data augmentations and pre-trained visual representations.

We start by formalizing our problem setting, and then introduce two algorithms, TD-MPC and Behavior Cloning (BC), that were previously proposed for online RL and imitation learning, respectively, and are both central to our work.

Problem Setting   We model interaction between an agent and its environment as an infinite-horizon Markov Decision Process (MDP) defined as a tuple $\mathcal{M}\coloneqq\langle\mathcal{S},\mathcal{A},\mathcal{T},\mathcal{R},\gamma\rangle$, where $\mathcal{S}\in\mathbb{R}^{n},~{}\mathcal{A}\in\mathbb{R}^{m}$ are continuous state and action spaces, $\mathcal{T}\colon\mathcal{S}\times\mathcal{A}\mapsto\mathcal{S}$ is the environment transition function, $\mathcal{R}\colon\mathcal{S}\times\mathcal{A}\mapsto\mathbb{R}$ is a scalar reward function, and $\gamma\in[0,1)$ is the discount factor. $\mathcal{T},\mathcal{R}$ are assumed to be unknown. The goal for the agent is to learn a policy $\pi_{\theta}\colon\mathcal{S}\mapsto\mathcal{A}$ parameterized by $\theta$ that achieves high long term performance, i.e. $\max_{\theta}\mathbb{E}_{\pi_{\theta}}\left[\sum_{t=0}^{\infty}\gamma^{t}r_{t}\right]$, while using as few interactions with the environment as possible – referred to as sample efficiency.

In this work, we focus more specifically on high-dimensional MDPs with sparse rewards. This is motivated by applications in robotics and Embodied Intelligence where the state is not directly observable, but can be well-approximated through the combination: $\mathbf{s}=(\mathbf{x},\mathbf{q})$, where $\mathbf{x}$ denotes stacked RGB images observed by the agent’s camera, and $\mathbf{q}$ denotes proprioceptive sensory information, e.g., the joint pose of a robot. Furthermore, shaped reward functions can be hard to script for real-world applications (Singh et al., 2019) or result in undesirable artifacts or behaviors (Amodei et al., 2016; Burda et al., 2019). Thus, we desire to learn with simple sparse rewards that accurately capture task completion. The principal challenges with sample efficient learning in such MDPs are exploration due sparse rewards, and learning good representations of the high-dimensional state space (e.g. images). To help overcome these challenges, we assume access to a small number of successful demonstration trajectories $\mathcal{D}^{\mathrm{demo}}\coloneqq\{D_{1},D_{2},\dots,D_{N}\}$, obtained through human teleoperation, kinesthetic teaching, or other means. Our goal is to learn a successful policy using minimal environment interactions, by accelerating the learning process using the provided demonstration trajectories.

Behavior Cloning (BC) is a simple and widely used framework for imitation learning (IL). This relies entirely on $\mathcal{D}^{\mathrm{demo}}$ and presents an ideal scenario if successful policies can be trained – since it incurs zero interaction sample complexity. BC trains a parameterized policy $\pi_{\theta}\colon\mathcal{S}\mapsto\mathcal{A}$ to predict the demonstrated action from the corresponding observation (Pomerleau, 1988; Atkeson & Schaal, 1997). Behavior cloning with a small dataset is known to be challenging due to covariate shift (Ross et al., 2011a; Rajeswaran et al., 2018) and the challenge of visual representation learning (Parisi et al., 2022; Nair et al., 2022). Further, collecting a large demonstration dataset is often infeasible due to the required human costs and expertise (Baker et al., 2022). More importantly, BC cannot improve beyond the capabilities of the demonstrator since it lacks a notion of task success. Together, these considerations motivate the need for combining demonstrations with sample-efficient RL.

TD-MPC (Hansen et al., 2022) is a model-based RL algorithm that combines Model Predictive Control (MPC), a learned latent-space world model, and a terminal value function learned via temporal difference (TD) learning. Specifically, TD-MPC learns a representation $\mathbf{z}=h_{\theta}(\mathbf{s})$ that maps the high-dimensional state $(\mathbf{s})$ into a compact representation, and a dynamics model in this latent space $\mathbf{z}^{\prime}=d_{\theta}(\mathbf{z},\mathbf{a})$. In addition, TD-MPC also learns prediction heads, $R_{\theta},Q_{\theta},\pi_{\theta}$, for: (i) instantaneous reward $r=R_{\theta}(\mathbf{z},\mathbf{a})$, (ii) state-action value $Q_{\theta}(\mathbf{z},\mathbf{a})$, and (iii) action $\mathbf{a}\sim\pi_{\theta}(\mathbf{z})$. The policy $\pi_{\theta}$ serves to "guide" planning towards high-return trajectories, and is optimized to maximize temporally weighted $Q$-values. The remaining components are jointly optimized to minimize TD-errors, reward prediction errors, and latent state prediction errors. The objective is given by

where $\bar{\theta}$ is an exponential moving average of $\theta$, and $\mathbf{s}_{t}^{\prime},\mathbf{z}_{t}^{\prime}$ are the (latent) states at time $t+1$. During environment interaction, TD-MPC uses a sample-based planner (Williams et al., 2015) in conjunction with the learned latent world-model and value function (critic) for action selection. See Appendix B for additional background on TD-MPC. We use TD-MPC as our choice of visual MBRL algorithm due to its simplicity and strong empirical performance (see Appendix D), but our framework, in principle, can be instantiated with any MBRL algorithm.

In this work, our goal is to accelerate the sample efficiency of (visual) model-based RL with a handful of demonstrations. To this end, we propose Model-based Reinforcement Learning with Demonstrations (MoDem), a simple and intuitive framework for visual RL under a strict environment-interaction budget. Figure 2 provides an overview of our method. MoDem consists of three phases: (1) a policy pretraining phase where the policy is trained on a handful of demonstrations via behavior cloning, (2) a seeding phase where the pretrained policy is used to collect a small dataset for targeted world-model learning, and (3) an interactive learning phase where the agent iteratively collects new data and improves using data from all the phases, with special emphasis on the demonstration data. We describe each phase in more detail below.

Phase 1: policy pretraining. We start by learning a policy from the demonstration dataset ${\mathcal{D}^{\mathrm{\color[rgb]{0.80078125,0.25390625,0.14453125}demo}}\coloneqq\{D_{1},D_{2},\dots,D_{N}\}}$ where each demonstration $D_{i}$ consists of $\{\mathbf{s}_{0},\mathbf{a}_{0},\mathbf{s}_{1},\mathbf{a}_{1},\dots,\mathbf{s}_{T},\mathbf{a}_{T}\}$. In general, the demonstrations may be noisy or sub-optimal – we do not explicitly make any optimality assumptions. Let $h_{\theta}\colon\mathcal{S}\mapsto\mathbb{R}^{l}$ denote the encoder and $\pi_{\theta}\colon\mathbb{R}^{l}\mapsto\mathcal{A}$ denote the policy that maps from the latent state representation to the action space. In Phase 1, we pretrain both the policy and encoder using a behavior-cloning objective, given by

When $\pi_{\theta}(\cdot|\mathbf{z})$ is parameterized by an isotropic Gaussian distribution, as commonly used in practice, Eq. 2 simplifies to the standard MSE loss. As outlined in Section 2, behavior cloning with a small demonstration dataset is known to be difficult, especially from high-dimensional visual observations (Duan et al., 2017; Jang et al., 2021; Parisi et al., 2022). In Section 4, we indeed show that behavior cloning alone cannot train successful policies for the environments and datasets we study, even when using pre-trained visual representations (Parisi et al., 2022; Nair et al., 2022). However, policy pretraining can provide strong inductive priors that facilitate sample-efficient adaptation in subsequent phases outlined below (Peters et al., 2010).

Phase 2: seeding. In the previous phase, we only pretrained the policy. In Phase 2, our goal is to also pretrain the critic and world-model, which requires a “seeding” dataset with sufficient exploration. A random policy is conventionally used to collect such a dataset in algorithms like TD-MPC. However, for visual RL tasks with sparse rewards, a random policy is unlikely to yield successful trajectories or visit task-relevant parts of the state space. Thus, we collect a small dataset with additive exploration using the policy from phase 1.

Concretely, given $\pi_{\theta}^{\mathrm{\color[rgb]{0.80078125,0.25390625,0.14453125}P1}}$ and $h_{\theta}^{\mathrm{\color[rgb]{0.80078125,0.25390625,0.14453125}P1}}$ from the first phase, we collect a dataset $\mathcal{D}^{\mathrm{\color[rgb]{0.41796875,0.66015625,0.3125}seed}}=\{\tau_{1},\tau_{2},\ldots\tau_{K}\}$ by rolling out $\pi^{\mathrm{\color[rgb]{0.80078125,0.25390625,0.14453125}P1}}_{\theta}(h^{\mathrm{\color[rgb]{0.80078125,0.25390625,0.14453125}P1}}_{\theta}(\mathbf{s}))$. To ensure sufficient variability in trajectories, we add Gaussian noise to actions (Hansen et al., 2022). Let $\xi_{t}=(\mathbf{s}_{i},\mathbf{a}_{i},r_{i},\mathbf{s}^{\prime}_{i})_{i=t}^{t+H}$ be a generic trajectory snippet of length $H$. In this phase, we learn $\pi_{\theta}$, $h_{\theta}$, $d_{\theta},R_{\theta},Q_{\theta}$ – the policy, representation, dynamics, reward, and value models – by minimizing the loss

where $\kappa$ is an “oversampling” rate that provides more weight to the demonstration dataset. In summary, the seeding phase plays the key role of initializing the world model, reward, and critic in the task-relevant parts of the environment, both through data collection and demonstration oversampling. We find in Section 4 that the seeding phase is crucial for sample-efficient learning, without which the learning agent is unable to make best use of the inductive priors in the demonstrations.

Phase 3: interactive learning. After initial pretraining of model and policy, we continue to improve the agent using fresh interactions with the environment. To do so, we initialize the replay buffer from the second phase, i.e. $\mathcal{B}\leftarrow\mathcal{D}^{\mathrm{\color[rgb]{0.41796875,0.66015625,0.3125}seed}}$. A naïve approach to utilizing the demonstrations is to simply append them to the replay buffer. However, we find this to be ineffective in practice, since online interaction data quickly outnumbers demonstrations. In line with the seeding phase, we propose to aggressively oversample demonstration data throughout training, but progressively anneal away the oversampling through the course of training. Concretely, we minimize the loss

Finally, we find it highly effective to regularize the visual representation using data augmentation, which we apply in all phases. We detail our proposed algorithm in Algorithm 1.

Imitation Learning, and behavior cloning (Pomerleau, 1988; Atkeson & Schaal, 1997; Ross et al., 2011b)) in particular, has been successfully applied to a variety of robotics applications such as manipulation (Zhang et al., 2018b; Song et al., 2020; Young et al., 2020), locomotion (Nakanishi et al., 2003), and autonomous driving (Bojarski et al., 2016), where demonstration collection is easy. Recent works have also explored extensions like third-person imitation learning (Torabi et al., 2018; Radosavovic et al., 2021; Zakka et al., 2021; Kumar et al., 2022)) and multi-task learning (Duan et al., 2017; Sharma et al., 2018; Jang et al., 2021; Reed et al., 2022). Regardless, most prior work on IL require large datasets with hundreds of expert demonstrations to provide sufficient coverage of the state space. Instead, we require only a few demonstrations and allow a limited amount of interaction.

Sample-efficient RL. Algorithm design for RL has a large body of work  (Mnih et al., 2013; Lillicrap et al., 2016; Schrittwieser et al., 2020; Hansen et al., 2022). However, direct real-world applications remain limited due to poor sample efficiency. Various methods have been proposed to improve sample-efficiency of RL (Levine et al., 2016; Zhang et al., 2018a; Ebert et al., 2018; Zhan et al., 2020). For example, Ebert et al. (2018) learn a goal-conditioned world model via video prediction and use the learned model for planning, and other works improve sample-efficiency by using pretrained representations (Shah & Kumar, 2021; Pari et al., 2022; Parisi et al., 2022; Xiao et al., 2022), data augmentation (Yarats et al., 2021b; a; Hansen et al., 2021), and auxiliary objectives (Srinivas et al., 2020). Recently, Wu et al. (2022) have also demonstrated that modern model-based RL algorithms can solve locomotion tasks with only one hour of real-world training. Interestingly, they find that their choice of model-based RL algorithm – Dreamer-V2 (Hafner et al., 2020) – is consistently more sample-efficient than model-free alternatives across all tasks considered. Our work builds upon a state-of-the-art model-based method, TD-MPC (Hansen et al., 2022), and shows that leveraging demonstrations can further yield large gains in sample-efficiency. Our contributions are compatible with and largely orthogonal to most prior work on model-based RL.

RL with demonstrations. While IL and RL have been studied extensively in isolation, recent works have started to explore their intersection (Ho & Ermon, 2016; Rajeswaran et al., 2018; Nair et al., 2018; 2020; Zhan et al., 2020; Shah & Kumar, 2021; Mandlekar et al., 2021; Wang et al., 2022; Baker et al., 2022). Notably, Rajeswaran et al. (2018) augments a policy gradient algorithm with demonstration data to solve dexterous manipulation tasks from states. By virtue of using policy gradient, their method achieves stable learning and high asymptotic performance, but poor sample efficiency. Shah & Kumar (2021) extended this to visual spaces using a pre-trained visual representation network, but inherit the same limitations. Zhan et al. (2020) accelerate a SAC agent (Haarnoja et al., 2018) with demonstrations and contrastive learning. Baker et al. (2022) first learn a policy using imitation learning on a large, crowd-sourced dataset and then finetune the policy using on-policy RL. In contrast to all prior work, we accelerate a model-based algorithm with demonstrations, which we find leads to significantly improved results compared to prior model-free alternatives.

In this work, we studied the acceleration of MBRL algorithms with expert demonstrations to improve the sample efficiency. We showed that naively appending demonstration data to an MBRL agent’s replay buffer does not meaningfully accelerate the learning. We developed MoDem, a three phase framework to fully utilize demonstrations and accelerate MBRL. Through extensive experimental evaluation, we find that MoDem trains policies that are $250\%-350\%$ more successful in sparse reward visual RL tasks in the low data regime ($100K$ online interactions, $5$ demonstrations). In this process, we also elucidated the importance of different phases in MoDem, the importance of data augmentation for visual MBRL, and the utility of pre-trained visual representations.