Example-Driven Model-Based Reinforcement Learning for Solving Long-Horizon Visuomotor Tasks
Abstract:
In this paper, we study the problem of learning a repertoire of lowlevel skills from raw images that can be sequenced to complete long-horizon visuomotor tasks.
Reinforcement learning (RL) is a promising approach for acquiring short-horizon skills autonomously.
However, the focus of RL algorithms has largely been on the success of those individual skills, more so than learning and grounding a large repertoire of skills that can be sequenced to complete extended multi-stage tasks.
The latter demands robustness and persistence, as errors in skills can compound over time, and may require the robot to have a number of primitive skills in its repertoire, rather than just one.
To this end, we introduce EMBR, a model-based RL method for learning primitive skills that are suitable for completing long-horizon visuomotor tasks.
EMBR learns and plans using a learned model, critic, and success classiﬁer, where the success classiﬁer serves both as a reward function for RL and as a grounding mechanism to continuously detect if the robot should retry a skill when unsuccessful or under perturbations.
Further, the learned model is task-agnostic and trained using data from all skills, enabling the robot to efﬁciently learn a number of distinct primitives.
These visuomotor primitive skills and their associated pre- and post-conditions can then be directly combined with off-the-shelf symbolic planners to complete long-horizon tasks.
On a Franka Emika robot arm, we ﬁnd that EMBR enables the robot to complete three long-horizon visuomotor tasks at 85% success rate, such as organizing a desk, a cabinet, and drawers, which require sequencing up to 12 skills, involve 14 unique learned primitives, and demand generalization to novel objects.
Introduction
We want robots to robustly complete a variety of long-horizon visuomotor manipulation tasks.
Such tasks can be completed by composing a sequence of low-level primitive skills, otherwise known as “manipulation primitives” [1] or “options” [2].
For example, the long-horizon task of “organizing an ofﬁce desk” often means performing a long sequence of shorter tasks, such as picking objects up, opening or closing drawers, placing objects into drawers, and putting pens into pen-holders.
However, completing such a long-horizon task may require the robot to learn many low-level visuomotor skills, compose them sequentially, and generalize across novel objects.
During task execution, the robot also needs to detect its own failures (e.g. accidentally dropping an object) and correct them (e.g. picking the object up again) before continuing to the remaining parts of the task. We hope to a r X i v : . 10312v [ c s .
R O ] S e p endow robots with such generalization and robustness.
Approaches such as symbolic planning [3], hierarchical RL [4, 5, 6, 7, 8], and hierarchical planning [9] can accomplish compositional task generalization by reasoning in an abstract or symbolic space.
While these approaches have shown promising results on tasks such as simulated tool use [10] and non-vision-based manipulation [9], they rely on the ability to acquire a repertoire of robust, low-level motor skills.
These low-level skills can in principle be acquired through reinforcement learning, but a number of challenges remain.
First, the acquisition of robust and persistent low-level visuomotor manipulation skills is still challenging, and the errors in completing any sub-sequence of a long-horizon task can compound over time and lead to overall task failure.
Second, acquiring the requisite low-level visuomotor skills requires data-efﬁcient learning and reward or task speciﬁcation across a broad set of skills.
Finally, sequencing the learned primitive skills to perform a long-horizon task requires persistent grounding of the robot’s visual observation within the broader abstract plan.
To address these challenges, we introduce example-driven model-based RL (EMBR), an algorithm that learns a dynamics model and a set of Q-functions and success classiﬁers to acquire a repertoire of visuomotor skills that are both robust and can be directly grounded into symbols for symbolic planning of a long-horizon visuomotor task.
More speciﬁcally, the success classiﬁers in EMBR enable both robustness and skill grounding, as they allow the robot to determine the pre- and postconditions of a skill.
For example, a reward function and post-condition for the skill of opening a drawer can be obtained from learning a binary image classiﬁer of “whether the drawer is open” given the robot’s pixel observation.
These success classiﬁers require a modest number of example images, making it easy to specify a variety of skills.
Finally, by learning a task-agnostic dynamics model, EMBR allows data to be shared across all tasks, leading to greater data-efﬁciency.
The main contribution of this work is a framework for learning grounded visuomotor skills that can be directly sequenced by off-the-shelf symbolic planners to complete long-horizon tasks.
Experiments on a Franka Emika Panda robot arm ﬁnd that EMBR enables the robot to complete three long-horizon tasks from raw pixels with 85% success, such as organizing a desk, a cabinet, and two drawers, which require up to 14 learned unique primitives, sequencing up to 12 skills, and generalization to novel objects.
Moreover, our ablation studies conﬁrm that all components of our method including the dynamics model, success classiﬁers, Q-functions, and stages are essential to performance for tackling long-horizon tasks.
Related Work
Reinforcement learning (RL) has a rich history of being used for robotic control problems on tasks from locomotion [11, 12, 13] to object manipulation [14, 15, 16].
In this work, we focus on learning vision-based robotic manipulation skills, a problem studied by many prior works [17, 18, 19, 20].
Model-free deep RL has been effective on short-horizon skills like object grasping [21, 20, 22], pushing and throwing [19, 23, 24], and multi-task learning [23, 25] from images.
Alternatively, model-based approaches [26, 27], which explicitly learn the environment forward dynamics from images have also been employed for multi-task vision-based robotic manipulation tasks, either with planning algorithms [28, 29, 30, 31, 32, 33, 34, 35] or for optimizing a parametric policy [36].
Unlike these prior works, our algorithm uses a model in conjunction with task-speciﬁc success classiﬁers and Q-functions to learn a wide range of skills suitable for sequencing into long-horizon tasks.
In
Section 5, we ﬁnd that both components of our approach are essential for robust performance.
Even with powerful reinforcement learning methods, task speciﬁcation on real robots remains challenging, as engineering rewards on physical systems can be costly and time-consuming [37].
Motivated by this, many works have studied the reward speciﬁcation, including inverse reinforcement learning [38] with robot demonstrations [39, 40, 41, 42, 43], learning rewards from user preferences [44, 45, 46, 47], and learning rewards from videos of humans [48, 49].
One common approach in visual RL is to learn to reach a goal image using a coarse measure of reward like negative `pixel distance [30, 50, 51] or temporal distance [52, 53, 34].
While these techniques have had some success on real robots, such rewards often provide a sparse and difﬁcult to optimize reward signal.
Alternatively, a number of recent works learn a classiﬁer from a modest number of examples goal states [54, 55, 24, 56, 57, 25] and tries to learn agents which maximize the classiﬁer score.
In this work, we extend this approach to multiple stages of human-provided goals, which we ﬁnd enables more effective learning of hard exploration skills, and ﬁnd that such reward speciﬁcation can be used to learn 14 unique visuomotor skills on a real robot without any demonstrations.
Like many prior works, the goal of this work is to complete challenging long-horizon visuomotor tasks.
One class of prior work is hierarchical RL, which aims to learn temporally-extended primitive skills and a high-level policy over them [4, 5, 6, 7, 8, 58].
Such approaches include jointly learning the primitive skills with the high-level policy through goal generation or an “options” framework [6, 8, 59], learning primitives through intrinsic motivation or other auxiliary objectives [60, 61, 62], from demonstration behavior [63, 64, 65, 66, 67], or from low-level, potentially goal conditioned, reward functions [7, 58, 68, 69].
EMBR is not a hierarchical RL algorithm, as it only learns lowlevel skills provided to the high-level planner and does not learn a model over options.
Nevertheless, our work is similar to the last group, except that we learn visuomotor primitives on a real robot from rewards derived from human examples.
Other works have also explored learning hierarchical policies on top of hand-designed skills [70, 71, 72].
These methods are complementary to EMBR, as skills learned via EMBR can be incorporated into most hierarchical policy learning algorithm.
Alternatively, a number of recent works have studied long horizon visual planning [52, 53, 73, 74, 75, 50, 76, 77, 78], often by combining a structured search approach with learned components like generative models [73, 50, 76] or distance functions [53, 34].
While these approaches have demonstrated impressive results on problems such as visual navigation [52, 53], unlike prior work we learn manipulation tasks on a real robots that consists of up to 12 unique skills and 250+ timesteps.
Lastly, the task and motion planning (TAMP) literature [79, 80] has extensively studied long-horizon robotic tasks.
Unlike these works, we do not assume pre-deﬁned state representations or accurate state estimators, making it possible to handle novel objects in cluttered scenes.
Like our work, a number of works have also explored combining learned models and skills with hand-designed or symbolic planning [10, 9, 81] and learning to ground skills for planning [82].
Unlike these techniques, a key insight in this work is that we can jointly learn visuomotor skills and the associated grounding of their pre- and post-conditions, allowing us to leverage off-the-shelf symbolic planners to execute long horizon tasks in a closed-loop fashion directly from pixels on a real robot.
Preliminaries
Tasks and skills .
We consider the problem of learning a repertoire of skills that can be sequenced to complete a long-horizon task M .
We model the robot’s environment as a controlled Markov process
E = h S , ⇢ 0 , A , T ,  , H i , with an image observation space s 2 S , an initial state distribution ⇢ 0 , an action space a 2 A , a dynamics model T : S ⇥
A ⇥ S !
R , a discount factor  2 [0 , 1) , and a ﬁnite horizon H .
We notice that many long-horizon tasks are composed of a sequence of lower-level skills.
For example, to organize an ofﬁce desk, the robot needs to acquire a number of skills such as grasping an object, opening a drawer, and placing an object in a drawer.
Let K denote the total number of unique skills the robot has acquired, and k 2 [1 , K ] denote the k th skill in the robot’s skill repertoire.
Depending on the environment state, the order in which the skills should be executed to successfully complete a long-horizon task varies.
For example, to organize a desk, the robot may not need to use the skill of closing a drawer if the drawer is already closed.
Here, each skill is an
MDP M k = h E ,
R k
i , where the robot’s environment E is shared across all tasks and skills, and
R k : S ⇥ A !
R is the reward function for the k th skill (e.g. the skill of opening a drawer).
Problem Deﬁnition .
Formally, EMBR takes as input N human-provided example images k
= s 1 , . . .
, s N for each skill k
we intend to robot to learn.
Our goal (i.e. EMBR’s output) is to learn the symbolic grounding of the pre- and post-conditions g kpre , g kpost : S ! { 0 , 1 } , and a policy ⇡ k : S !
⇧ ( A ) for each skill k , where ⇧ ( · ) deﬁnes a probability distribution over a set.
Each policy has a Q-function, Q k : S ⇥ A !
R for taking action a from s and following ⇡ k onward.
Variational Autoencoders (VAEs) .
VAEs compress high-dimensional observations s such as images into an embedding z .
VAEs can be optimized by maximizing the evidence lower bound (ELBO): max p,q L vae ( s ) , where L vae ( s ) =
E q ( z | s )
[Equation]
where p denotes the generative model, and q denotes the variational distribution.
Symbolic planning .
A symbolic planner performs task planning in the symbolic domain.
Concretely, given a set of candidate actions, each action’s associated pre- and post-conditions, the current condition h and a goal condition g , a symbolic planner outputs the sequence of actions that allows the robot to reach from h to g .
The pre-condition of an action is a set of predicates that deﬁnes the condition that needs to be satisﬁed in order for the action to be executable.
The post-condition of a skill is a set of predicates that deﬁnes the effect of executing an action.
For example, the pre- condition and post-condition of an “Open left drawer” action is “Drawer is closed” and “Drawer is open” respectively.
In our context, the candidate actions are the primitive skills learned by EMBR.
Example-Driven Model-Based RL (EMBR)
To learn a repertoire of skills that can be directly sequenced by an off-the-shelf symbolic planner to complete long-horizon visuomotor tasks, we propose Example-driven Model-Based RL, or EMBR.
Below we ﬁrst describe how EMBR learns each visuomotor skill without any demonstration skills or trajectories during training, and then elaborate on how these learned skills can be directly sequenced to complete a long-horizon task at test time.
4.1 Learning a Repertoire of Visuomotor Skills
During training, EMBR learns a repertoire of low-level persistent visuomotor skills, such as opening a drawer or inserting a marker, using a set of example success images provided by a person for each task.
The model learned by EMBR, denoted as f , is composed of four components: a global variational autoencoder (VAE) f vae , a global latent dynamics model f T , stage-speciﬁc binary reward functions f k R , and stage-speciﬁc Q-functions f kQ .
The VAE and latent dynamics model jointly learn a low-dimensional representation of the image observations to address the challenge of efﬁciently learning many primitive skills; both are trained on data collected across all skills.
The skill-speciﬁc binary classiﬁers allow us to learn a reward function for each skill from human-provided images.
We use the latent dynamics model to perform model predictive control with the skill-speciﬁc Q-function as the terminal planning objective.
The skill-speciﬁc Q-functions and reward functions are trained on skill-speciﬁc data.
Next, we elaborate on each of these components.
Learning a low-dimensional latent space with a VAE and a latent dynamics model .
Performing a long-horizon task using a large number of visuomotor skills demands the algorithm to be sampleefﬁcient and fast at execution time.
To address both of these demands, we learn a single VAE f vae = { f enc , f dec } to construct a latent space Z .
This VAE consists of an encoder f enc : S !
⇧ ( Z ) and a decoder f dec : Z ! ⇧ ( S ) .
We train f vae using data collected across all skills via Eq.
1, where p ( s | z ) ⌘ f dec ( s t | z t ) , p ( z ) ⇠ N (0 , 1) , and q ( z | s ) ⌘ f enc ( z t | s t ) .
By learning from all skills, the VAE can improve data efﬁciency.
In practice, we use the greedy hierarchical VAE [51] as f vae with the current image observation as a context frame.
Based on this latent space, EMBR learns a latent dynamics model f T jointly with f vae with data collected across all skills via l 2 loss:
[Equation]
where z t ⇠ f enc ( · | s t ) , L is the rollout length of f T and D is the dataset collected by the robot.
Obtaining rewards by learning image classiﬁers .
Obtaining a reward function in RL typically requires some forms of human supervision such as reward engineering or demonstrations.
In EMBR, binary image classiﬁers serve not only as reward functions for learning primitive visuomotor skills, but also as success detectors as well as grounding of raw pixels for long-horizon symbolic planning.
Concretely, we learn a reward function f k R for each skill M k by training binary image classiﬁers to distinguish human-provided example images from robot-collected non-example images using the binary cross-entropy objective [55, 56].
The reward learning objective for skill k is:
[Equation]
Here, D + is the dataset of images labeled as positive by a person (typically 100-200 images), s + are the images sampled from D + , and s  are sampled from all images in D , the dataset collected by the robot, and are considered negative by default.
In our experiments, these positive examples can be collected in 15-30 minutes per visuomotor skill, and 345 minutes in total for all 14 skills.
No demonstration trajectories are collected.
At test time, the learned image classiﬁers are re-purposed as post-condition grounding of each skill for long-horizon task execution.
Some of the pre-conditions the task planner requires can be derived from existing post-conditions.
For example, for a drawer-opening skill, the pre-condition of “drawer is not open” can be derived from the post-condition of “drawer is open” by negating this post-condition.
For the remaining pre-conditions that cannot be derived from any existing postconditions, we train additional image classiﬁers using Eq. 3.
Learning Q-functions for model-based control .
While one can directly use the latent dynamics model for model-based planning, the difﬁculty of learning an accurate dynamics model for the full horizon of a primitive skill increases with the length of the skill itself.
To alleviate this challenge and further accelerate skill learning, we learn a Q-function f kQ in the latent space.
Concretely, we perform Q-learning via the following objective for skill k :
[Equation]
t +  max a t +1 f kQ ( z t +1 , a t +1 ) ◆# (4) where f k R ( z ) ⌘ 1 { f k R ( z ) > 0 . 5 } , z t ⇠ f vae ( · | s t ) , z t +1 ⇠ f vae ( · | s t +1 ) , and the Q-value target is computed by maximizing over m 1 randomly and uniformly sampled actions a 1: m 1 t
+1 .
At both training and test time, we perform model-based planning in the latent space by ﬁrst rolling out f T for L timesteps into the future and then evaluating the predicted future latent space observation at timestep t + L using the learned Q-function, a process that is illustrated in Fig. 2:
[Equation]
Here, the outer and inner maximization operations are taken over m 0 and m 1 samples respectively.
[Figure] Integrating these four components, EMBR is trained simultaneously across all components from scratch (i.e. without pre-training) for learning a repertoire of visuomotor primitive skills (Alg. 1).
Concretely, the robot begins with an empty dataset D and can learn each skill in almost any order (Appendix C.2).
During skill learning, the robot collects a new trajectory using actions selected from Eq.
5 and updates all model components using D .
This process repeats until the current skill reaches a desirable success rate as computed by the success classiﬁers, before the robot moves on to learn the next skill.
As the robot acquires more skills, the dataset D accumulates and improves the generalization of the visual dynamics model, which in-turn accelerates future skill learning.
Wallclock training time per skill and hyperparameters are detailed in Appendix G and H respectively.
4.2 Composing Learned Visuomotor Primitive Skills for Long-Horizon Task Execution
At test time, EMBR provides the learned primitive skills and the symbolic grounding of each skill’s post-conditions (in the form of success classiﬁers) to a high-level task planner (standard PDDL
STRIPS symbolic planner [83] in our case) to perform long-horizon task.
As elaborated in Alg. 2, the human ﬁrst speciﬁes a long-horizon task using a symbolic goal condition g .
Next, EMBR computes the current symbolic condition of the environment from raw pixel observations by passing the current image observations to each of the success classiﬁers.
Using both the symbolic goal condition and the current symbolic condition, the symbolic planner computes the sequence of skills to perform and executes the ﬁrst skill of this sequence using actions selected from Eq. 5.
This procedure repeats until the current condition matches the goal condition, after which task execution terminates.
Note that the arm position of the robot is reset after executing each skill.
The predicate deﬁnitions, the pre- and post-conditions and the full details of the symbolic planner used in Section 5 are in Appendix C.3.
Experiments
Our experiments aim to answer four key questions: 1) how does EMBR compare to prior RL algorithms on both performing primitive skills and when sequencing those skills to complete longhorizon tasks?
2) how necessary are stage-speciﬁc reward functions (as elaborated in Appendix ?? ) compared to a single skill-speciﬁc reward function?
3) how important is replanning using the classiﬁers for the robustness of the algorithm?
and 4) how important is each of the ﬁve camera observations to the performance of the primitive skills?
To answer these questions, we conduct experiments across 14 primitive skills and three long-horizon tasks, using the experimental setup in Fig.
3, in which a Franka Emika Panda robot has access to ﬁve cameras, each capturing 64 ⇥ 64 RGB images.
One of the ﬁve cameras is mounted on the wrist.
To test object generalization, all objects used in all evaluations are novel and unseen during training, except for the desk, drawers, cabinet, and the marker holder.
Fig.
3 visualizes the train, validation, and test objects used in all experiments.
See Appendix C for experimental details.
Project website at https://sites.google.com/view/embr-site .
Comparisons.
We compare EMBR to three prior methods: (1) EMBR w/o f T , which corresponds to
EMBR without the dynamics model, i.e. using only the Q-function and the VAE for action selection; (2) EMBR w/o f vae , f T (Qt-Opt) [20], which is our re-implemented version of Qt-Opt [20] with a
VICE-like success classiﬁer [55] that does not learn a latent space and uses only the Q-function for action selection; and (3) BEE [56], which uses the success classiﬁer instead of the Q-function [Table]
Here, EMBR outperforms the three prior by 15-25% and % respectively. [Table]
(65%) 0 (0%) 11 (55%) for model predictive control. All comparisons are representative of prior methods for vision-based robotic RL.
For all experiments, the data used to train each method is the same dataset collected by
EMBR using Alg. 1, which contains 300-400 trajectories collected in 16 hours per skill on average.
This provides a more direct comparison and in effect makes the prior methods stronger since the challenge of online exploration is addressed by EMBR during data collection.
5.1 Short-horizon Primitive Skill Performance
To answer the ﬁrst question, we compare all methods across 14 visuomotor skills.
In Table 1, we observe that “EMBR (Ours)” can reliably complete all 14 skills with at least 90% success.
Compared to “EMBR w/o f T ” and “EMBR w/o f vae , f T (Qt-Opt)”, “EMBR (Ours)” improves success rates by 10-20% and 15-25% across high-precision skills, 5-10% and 15-20% for medium-precision skills and 0% for low-precision skills.
We hypothesize that EMBR’s improvement over “EMBR w/o f T ” is due to its ability to alleviate overestimation of Q-values in the presence of limited training trajectories (300-400 per skill), and the addition success rate improvement of EMBR over “EMBR w/o f vae , f T ” is due to less overﬁtting of the success classiﬁers and Q-functions in the VAE latent space.
Finally, EMBR outperforms BEE [56] by 50-65% on high-precision tasks because BEE [56] cannot capture action consequences beyond its planning horizon and the model is not always accurate for the full horizon of a skill.
5.2 Long-horizon Task Performance
Next, we compare EMBR to the three prior methods when sequencing the learned visuomotor skills with symbolic planning to complete three challenging long-horizon tasks (Fig. 4):
Organize Desk and Cabinet (8 Skills) - putting three novel objects cluttered on the desk into the cabinet.
Organize Markers (9 Skills) - picking up three markers and inserting them into a holder.
Rearrange Objects (12 Skills) - picking up three objects and placing them to the desk, cabinet or drawers.
Detailed task descriptions are in Appendix C.7.
In Table 2, we ﬁnd that “EMBR (Ours)” completes these long-horizon visuomotor tasks with 85% success, including the most challenging tasks that require up to 12 unique skills.
Compared to “EMBR w/o f T ” and “EMBR w/o f vae , f T (Qt-Opt)”,
EMBR leads to 10-20% and 20-30% improvement in success rates respectively because EMBR achieves higher success rate on individual primitive skills that require higher precision, such as inserting marker and grasping object from desk or drawers.
Finally, BEE [56], is unable to perform any long-horizon task due to low performance on skills with medium or high precision.
This reﬂects the importance of learning robust, high-success-rate primitive skills using EMBR.
5.3 Ablations
Next, we conduct six ablations to examine the importance of various components of EMBR.
Ablation 3, which compares EMBR with vs. without any one camera observation, is in Appendix F.
Comparisons to the original BEE and training without stages are in Appendix D and E.
Ablation 1: Stage-speciﬁc vs. skill-speciﬁc reward function .
While EMBR accelerates skill acquisition by learning a reward function for each stage of a skill, one can in principle learn just one reward function for each skill.
In this ablation, we answer Question 2 by comparing EMBR to an ablated EMBR that learns each skill with a single skill-speciﬁc reward function.
This does not affect data collection and exploration, which stage-speciﬁc rewards can also help.
In Table 1, we ﬁnd that using stage-speciﬁc reward functions outperforms using a single skill-speciﬁc reward function by 10-30% for higher precision skills such as grasping an object in various locations, indicating the importance of using stage-speciﬁc reward functions in long-horizon or higher precision skills.
Ablation 2: Task performance with vs. without classiﬁer-based replanning .
While EMBR improves the robustness of long-horizon task execution by using success classiﬁers during planning, one can in principle perform model-based planning without these classiﬁers.
In this ablation, we answer Question 3 by comparing EMBR to an ablation that executes each skill of a task sequentially without using the success classiﬁers to determine when the skill is successful.
In Table 2, we ﬁnd that long-horizon task execution with classiﬁers improves performance by 30-35% success rate, due to the robot’s ability to detect success in a closed-loop manner until successful.
These ablations and
Table 2 suggest that all components of EMBR are critical to completing long-horizon tasks.
Conclusion and Future Work
In this work, we propose EMBR: an example-driven model-based reinforcement learning algorithm for learning visuomotor skills that can be directly sequenced to perform long-horizon tasks from raw pixels.
Across three long-horizon visuomotor tasks which require up to 14 unique primitives and sequencing up to 12 skills, we ﬁnd that EMBR consistently achieves 85% success rates, outperforms prior methods, and generalizes to novel objects.
Nonetheless, EMBR also has a number of important limitations, including the limited scope of tasks, the amount of supervision required to complete these complex long-horizon tasks, and the limited degree of generalization especially across different scenes.
More speciﬁcally, the visuomotor skill deﬁnitions in EMBR are currently object-speciﬁc, such as “open left drawer”, which will scale poorly in the presence of many objects.
We discuss limitations in more detail in Appendix A, each of which presents interesting directions for future work. Acknowledgments The authors would like to thank members of the IRIS and RAIL labs for providing valuable feedback. Suraj Nair is funded in part by an NSF GRFP.
This work was also supported in part by ONR grants N00014-20-1-2675 and N00014-21-1-2685.