Example-Driven Model-Based Reinforcement Learning for Solving Long-Horizon Visuomotor Tasks


Bohan Wu, Suraj Nair, Li Fei-Fei*, Chelsea Finn*

* Equal advising and contribution

In this paper, we study the problem of learning a repertoire of low-level skills from raw images that can be sequenced to complete long-horizon multi-stage 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, grounding, and sequencing a large repertoire of skills 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 EMBER, a model-based RL method for learning primitive skills that are suitable for completing long-horizon visuomotor tasks. EMBER learns and plans using a learned model, critic, and success classifier, where the success classifier 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 efficiently 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 find that EMBER enables the robot to complete three long-horizon visuomotor tasks at 85% success rate, such as organizing an office desk, a file cabinet, and drawers, which require sequencing up to 12 skills, involve 14 unique learned primitives, and demand generalization to novel objects.

Robot Execution Videos (24x Speed Up)

Organize Desk and Cabinet

Putting three objects cluttered on the desk into the cabinet and close all drawers and the cabinet.

Organize Markers

Picking up three markers and inserting them into a transparent marker holder, and close all drawers and the cabinet. Marker reorientation policy is scripted.

Rearrange Objects

Picking up three objects and placing them on the desk, in the cabinet or drawers, and close all drawers and the cabinet.