Month: February 2021

Posted by Danijar Hafner, Student Researcher, Google Research

Deep reinforcement learning (RL) enables artificial agents to improve their decisions over time. Traditional model-free approaches learn which of the actions are successful in different situations by interacting with the environment through a large amount of trial and error. In contrast, recent advances in deep RL have enabled model-based approaches to learn accurate world models from image inputs and use them for planning. World models can learn from fewer interactions, facilitate generalization from offline data, enable forward-looking exploration, and allow reusing knowledge across multiple tasks.

Despite their intriguing benefits, existing world models (such as SimPLe) have not been accurate enough to compete with the top model-free approaches on the most competitive reinforcement learning benchmarks — to date, the well-established Atari benchmark requires model-free algorithms, such as DQN, IQN, and Rainbow, to reach human-level performance. As a result, many researchers have focused instead on developing task-specific planning methods, such as VPN and MuZero, which learn by predicting sums of expected task rewards. However, these methods are specific to individual tasks and it is unclear how well they would generalize to new tasks or learn from unsupervised datasets. Similar to the recent breakthrough of unsupervised representation learning in computer vision [1, 2], world models aim to learn patterns in the environment that are more general than any particular task to later solve tasks more efficiently.

Today, in collaboration with DeepMind and the University of Toronto, we introduce DreamerV2, the first RL agent based on a world model to achieve human-level performance on the Atari benchmark. It constitutes the second generation of the Dreamer agent that learns behaviors purely within the latent space of a world model trained from pixels. DreamerV2 relies exclusively on general information from the images and accurately predicts future task rewards even when its representations were not influenced by those rewards. Using a single GPU, DreamerV2 outperforms top model-free algorithms with the same compute and sample budget.

Gamer normalized median score across the 55 Atari games after 200 million steps. DreamerV2 substantially outperforms previous world models. Moreover, it exceeds top model-free agents within the same compute and sample budget.

Behaviors learned by DreamerV2 for some of the 55 Atari games. These videos show images from the environment. Video predictions are shown below in the blog post.

An Abstract Model of the World
Just like its predecessor, DreamerV2 learns a world model and uses it to train actor-critic behaviors purely from predicted trajectories. The world model automatically learns to compute compact representations of its images that discover useful concepts, such as object positions, and learns how these concepts change in response to different actions. This lets the agent generate abstractions of its images that ignore irrelevant details and enables massively parallel predictions on a single GPU. During 200 million environment steps, DreamerV2 predicts 468 billion compact states for learning its behavior.

DreamerV2 builds upon the Recurrent State-Space Model (RSSM) that we introduced for PlaNet and was also used for DreamerV1. During training, an encoder turns each image into a stochastic representation that is incorporated into the recurrent state of the world model. Because the representations are stochastic, they do not have access to perfect information about the images and instead extract only what is necessary to make predictions, making the agent robust to unseen images. From each state, a decoder reconstructs the corresponding image to learn general representations. Moreover, a small reward network is trained to rank outcomes during planning. To enable planning without generating images, a predictor learns to guess the stochastic representations without access to the images from which they were computed.

Learning process of the world model used by DreamerV2. The world model maintains recurrent states (h1–h3) that receive actions (a1–a2) and incorporate information about the images (x1–x3) via stochastic representations (z1–z3). A predictor guesses the representations as (ẑ1–ẑ3) without access to the images from which they were generated.

Importantly, DreamerV2 introduces two new techniques to RSSM that lead to a substantially more accurate world model for learning successful policies. The first technique is to represent each image with multiple categorical variables instead of the Gaussian variables used by PlaNet, DreamerV1, and many more world models in the literature [1, 2, 3, 4, 5]. This leads the world model to reason about the world in terms of discrete concepts and enables more accurate predictions of future representations.

The encoder turns each image into 32 distributions over 32 classes each, the meanings of which are determined automatically as the world model learns. The one-hot vectors sampled from these distributions are concatenated to a sparse representation that is passed on to the recurrent state. To backpropagate through the samples, we use straight-through gradients that are easy to implement using automatic differentiation. Representing images with categorical variables allows the predictor to accurately learn the distribution over the one-hot vectors of the possible next images. In contrast, earlier world models that use Gaussian predictors cannot accurately match the distribution over multiple Gaussian representations for the possible next images.

Multiple categoricals that represent possible next images can be accurately predicted by a categorical predictor, whereas a Gaussian predictor is not flexible enough to accurately predict multiple possible Gaussian representations.

The second new technique of DreamerV2 is KL balancing. Many previous world models use the ELBO objective that encourages accurate reconstructions while keeping the stochastic representations (posteriors) close to their predictions (priors) to regularize the amount of information extracted from each image and facilitate generalization. Because the objective is optimized end-to-end, the stochastic representations and their predictions can be made more similar by bringing either of the two towards the other. However, bringing the representations towards their predictions can be problematic when the predictor is not yet accurate. KL balancing lets the predictions move faster toward the representations than vice versa. This results in more accurate predictions, a key to successful planning.

Long-term video predictions of the world model for holdout sequences. Each model receives 5 frames as input (not shown) and then predicts 45 steps forward given only actions. The video predictions are only used to gain insights into the quality of the world model. During planning, only compact representations are predicted, not images.

Measuring Atari Performance
DreamerV2 is the first world model that enables learning successful behaviors with human-level performance on the well-established and competitive Atari benchmark. We select the 55 games that many previous studies have in common and recommend this set of games for future work. Following the standard evaluation protocol, the agents are allowed 200M environment interactions using an action repeat of 4 and sticky actions (25% chance that an action is ignored and the previous action is repeated instead). We compare to the top model-free agents IQN and Rainbow, as well as to the well-known C51 and DQN agents implemented in the Dopamine framework.

Different standards exist for aggregating the scores across the 55 games. Ideally, a new algorithm would perform better under all conditions. For all four aggregation methods, DreamerV2 indeed outperforms all compared model-free algorithms while using the same computational budget.

DreamerV2 outperforms the top model-free agents according to four methods for aggregating scores across the 55 Atari games. We introduce and recommend the Clipped Record Mean (right-most plot) as an informative and robust performance metric.

The first three aggregation methods were previously proposed in the literature. We identify important drawbacks in each and recommend a new aggregation method, the clipped record mean to overcome their drawbacks.

• Gamer Median. Most commonly, scores for each game are normalized by the performance of a human gamer that was assessed for the DQN paper and the median of the normalized scores of all games is reported. Unfortunately, the median ignores the scores of many simpler and harder games.
• Gamer Mean. The mean takes the scores for all games into account but is mainly influenced by a small number of games where the human gamer performed poorly. This makes it easy for an algorithm to achieve large normalized scores on some games (e.g., James Bond, Video Pinball) that then dominate the mean.
• Record Mean. Prior work recommends normalization based on the human world record instead, but such a metric is still overly influenced by a small number of games where it is easy for the artificial agents to outscore the human record.
• Clipped Record Mean. We introduce a new metric that normalizes scores by the world record and clips them to not exceed the record. This yields an informative and robust metric that takes the performance on all games into account to an approximately equal amount.

While many current algorithms exceed the human gamer baseline, they are still quite far behind the human world record. As shown in the right-most plot above, DreamerV2 leads by achieving 25% of the human record on average across games. Clipping the scores at the record line lets us focus our efforts on developing methods that come closer to the human world record on all of the games rather than exceeding it on just a few games.

What matters and what doesn’t
To gain insights into the important components of DreamerV2, we conduct an extensive ablation study. Importantly, we find that categorical representations offer a clear advantage over Gaussian representations despite the fact that Gaussians have been used extensively in prior works. KL balancing provides an even more substantial advantage over the KL regularizer used by most generative models.

By preventing the image reconstruction or reward prediction gradients from shaping the model states, we study their importance for learning successful representations. We find that DreamerV2 relies completely on universal information from the high-dimensional input images and its representations enable accurate reward predictions even when they were not trained using information about the reward. This mirrors the success of unsupervised representation learning in the computer vision community.

Atari performance for various ablations of DreamerV2 (Clipped Record Mean). Categorical representations, KL balancing, and learning about the images are crucial for the success of DreamerV2. Using reward information, that is specific to narrow tasks, offers no additional benefits for learning the world model.

Conclusion
We show how to learn a powerful world model to achieve human-level performance on the competitive Atari benchmark and outperform the top model-free agents. This result demonstrates that world models are a powerful approach for achieving high performance on reinforcement learning problems and are ready to use for practitioners and researchers. We see this as an indication that the success of unsupervised representation learning in computer vision [1, 2] is now starting to be realized in reinforcement learning in the form of world models. An unofficial implementation of DreamerV2 is available on Github and provides a productive starting point for future research projects. We see world models that leverage large offline datasets, long-term memory, hierarchical planning, and directed exploration as exciting avenues for future research.

Acknowledgements
This project is a collaboration with Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. We further thank everybody on the Brain Team and beyond who commented on our paper draft and provided feedback at any point throughout the project.

Thank you everyone who participated in our first mlverse survey!

Wait: What even is the mlverse?

The mlverse originated as an abbreviation of multiverse1, which, on its part, came into being as an intended allusion to the well-known tidyverse. As such, although mlverse software aims for seamless interoperability with the tidyverse, or even integration when feasible (see our recent post featuring a wholly tidymodels-integrated torch network architecture), the priorities are probably a bit different: Often, mlverse software’s raison d’être is to allow R users to do things that are commonly known to be done with other languages, such as Python.

As of today, mlverse development takes place mainly in two broad areas: deep learning, and distributed computing / ML automation. By its very nature, though, it is open to changing user interests and demands. Which leads us to the topic of this post.

The survey

GitHub issues and community questions are valuable feedback, but we wanted something more direct. We wanted a way to find out how you, our users, employ the software, and what for; what you think could be improved; what you wish existed but is not there (yet). To that end, we created a survey. Complementing software- and application-related questions for the above-mentioned broad areas, the survey had a third section, asking about how you perceive ethical and social implications of AI as applied in the “real world”.

A few things upfront:

Firstly, the survey was completely anonymous, in that we asked for neither identifiers (such as e-mail addresses) nor things that render one identifiable, such as gender or geographic location. In the same vein, we had collection of IP addresses disabled on purpose.

Secondly, just like GitHub issues are a biased sample, this survey’s participants must be. Main venues of promotion were rstudio::global, Twitter, LinkedIn, and RStudio Community. As this was the first time we did such a thing (and under significant time constraints), not everything was planned to perfection – not wording-wise and not distribution-wise. Nevertheless, we got a lot of interesting, helpful, and often very detailed answers, – and for the next time we do this, we’ll have our lessons learned!

Thirdly, all questions were optional, naturally resulting in different numbers of valid answers per question. On the other hand, not having to select a bunch of “not applicable” boxes freed respondents to spend time on topics that mattered to them.

As a final pre-remark, most questions allowed for multiple answers.

In sum, we ended up with 138 completed2 surveys. Thanks again everyone who participated, and especially, thank you for taking the time to answer the – many – free-form questions!

Deep learning Areas and applications

Our first goal was to find out in which settings, and for what kinds of applications, deep-learning software is being used.

Overall, 72 respondents reported using DL in their jobs in industry, followed by academia (23), studies (21), spare time (43), and not-actually-using-but-wanting-to (24).

Of those working with DL in industry, more than twenty said they worked in consulting, finance, and healthcare (each). IT, education, retail, pharma, and transportation were each mentioned more than ten times:

(#fig:unnamed-chunk-1)Number of users reporting to use DL in industry. Smaller groups not displayed.

In academia, dominant fields (as per survey participants) were bioinformatics, genomics, and IT, followed by biology, medicine, pharmacology, and social sciences:

(#fig:unnamed-chunk-2)Number of users reporting to use DL in academia. Smaller groups not displayed.

What application areas matter to larger subgroups of “our” users? Nearly a hundred (of 138!) respondents said they used DL for some kind of image-processing application (including classification, segmentation, and object detection). Next up was time-series forecasting, followed by unsupervised learning.

The popularity of unsupervised DL was a bit unexpected; had we anticipated this, we would have asked for more detail here. So if you’re one of the people who selected this – or if you didn’t participate, but do use DL for unsupervised learning – please let us know a bit more in the comments!

Next, NLP was about on par with the former; followed by DL on tabular data, and anomaly detection. Bayesian deep learning, reinforcement learning, recommendation systems, and audio processing were still mentioned frequently.

(#fig:unnamed-chunk-3)Applications deep learning is used for. Smaller groups not displayed.

Frameworks and skills

We also asked what frameworks and languages participants were using for deep learning, and what they were planning on using in the future. Single-time mentions (e.g., deeplearning4J) are not displayed.

(#fig:unnamed-chunk-4)Framework / language used for deep learning. Single mentions not displayed.

An important thing for any software developer or content creator to investigate is proficiency/levels of expertise present in their audiences. It (nearly) goes without saying that expertise is very different from self-reported expertise. I’d like to be very cautious, then, to interpret the below results.

While with regard to R skills3, the aggregate self-ratings look plausible (to me), I would have guessed a slightly different outcome re DL. Judging from other sources (like, e.g., GitHub issues), I tend to suspect more of a bimodal distribution (a far stronger version of the bimodality we’re already seeing, that is). To me, it seems like we have rather many users who know a lot about DL. In agreement with my gut feeling, though, is the bimodality itself – as opposed to, say, a Gaussian shape.

But of course, sample size is moderate, and sample bias is present.

(#fig:unnamed-chunk-5)Self-rated skills re R and deep learning.

Wishes and suggestions

Now, to the free-form questions. We wanted to know what we could do better.

I’ll address the most salient topics in order of frequency of mention.4 For DL, this is surprisingly easy (as opposed to Spark, as you’ll see).

“No Python”

The number one concern with deep learning from R, for survey respondents, clearly has to do not with R but with Python. This topic appeared in various forms, the most frequent being frustration over how hard it can be, dependent on the environment, to get Python dependencies for TensorFlow/Keras correct. (It also appeared as enthusiasm for torch, which we are very happy about.)

Let me clarify and add some context.

TensorFlow is a Python framework (nowadays subsuming Keras, which is why I’ll be addressing both of those as “TensorFlow” for simplicity) that is made available from R through packages tensorflow and keras . As with other Python libraries, objects are imported and accessible via reticulate . While tensorflow provides the low-level access, keras brings idiomatic-feeling, nice-to-use wrappers that let you forget about the chain of dependencies involved.

On the other hand, torch, a recent addition to mlverse software, is an R port of PyTorch that does not delegate to Python. Instead, its R layer directly calls into libtorch, the C++ library behind PyTorch. In that way, it is like a lot of high-duty R packages, making use of C++ for performance reasons.

Now, this is not the place for recommendations. Here are a few thoughts though.

Clearly, as one respondent remarked, as of today the torch ecosystem does not offer functionality on par with TensorFlow, and for that to change time and – hopefully! more on that below – your, the community’s, help is needed. Why? Because torch is so young, for one; but also, there is a “systemic” reason! With TensorFlow, as we can access any symbol via the tf object, it is always possible, if inelegant, to do from R what you see done in Python. Respective R wrappers nonexistent5, quite a few blog posts (see, e.g., https://blogs.rstudio.com/ai/posts/2020-04-29-encrypted_keras_with_syft/, or A first look at federated learning with TensorFlow) relied on this!

Switching to the topic of tensorflow’s Python dependencies causing problems with installation, my experience (from GitHub issues, as well as my own) has been that difficulties are quite system-dependent. On some OSes, complications seem to appear more often than on others; and low-control (to the individual user) environments like HPC clusters can make things especially difficult. In any case though, I have to (unfortunately) admit that when installation problems appear, they can be very tricky to solve.

tidymodels integration

The second most frequent mention clearly was the wish for tighter tidymodels integration. Here, we wholeheartedly agree. As of today, there is no automated way to accomplish this for torch models generically, but it can be done for specific model implementations.

Last week, torch, tidymodels, and high-energy physics featured the first tidymodels-integrated torch package. And there’s more to come. In fact, if you are developing a package in the torch ecosystem, why not consider doing the same? Should you run into problems, the growing torch community will be happy to help.

Documentation, examples, teaching materials

Thirdly, several respondents expressed the wish for more documentation, examples, and teaching materials. Here, the situation is different for TensorFlow than for torch.

For tensorflow, the website has a multitude of guides, tutorials, and examples. For torch, reflecting the discrepancy in respective lifecycles, materials are not that abundant (yet). However, after a recent refactoring, the website has a new, four-part Get started section addressed to both beginners in DL and experienced TensorFlow users curious to learn about torch. After this hands-on introduction, a good place to get more technical background would be the section on tensors, autograd, and neural network modules.

Truth be told, though, nothing would be more helpful here than contributions from the community. Whenever you solve even the tiniest problem (which is often how things appear to oneself), consider creating a vignette explaining what you did. Future users will be thankful, and a growing user base means that over time, it’ll be your turn to find that some things have already been solved for you!

Community, community, community

The remaining items discussed didn’t come up quite as often (individually), but taken together, they all have something in common: They all are wishes we happen to have, as well!

This definitely holds in the abstract – let me cite:

“Develop more of a DL community”

“Larger developer community and ecosystem. Rstudio has made great tools, but for applied work is has been hard to work against the momentum of working in Python.”

We wholeheartedly agree, and building a larger community is exactly what we’re trying to do. I like the formulation “a DL community” insofar it is framework-independent. In the end, frameworks are just tools, and what counts is our ability to usefully apply those tools to problems we need to solve.

Concrete wishes include

• More paper/model implementations (such as TabNet).

• Facilities for easy data reshaping and pre-processing (e.g., in order to pass data to RNNs or 1dd convnets in the expected 3-d format).

• Probabilistic programming for torch (analogously to TensorFlow Probability).

• A high-level library (such as fast.ai) based on torch.

In other words, there is a whole cosmos of useful things to create; and no small group alone can do it. This is where we hope we can build a community of people, each contributing what they’re most interested in, and to whatever extent they wish.

Spark Areas and applications

For Spark, questions broadly paralleled those asked about deep learning.

Overall, judging from this survey (and unsurprisingly), Spark is predominantly used in industry (n = 39). For academic staff and students (taken together), n = 8. Seventeen people reported using Spark in their spare time, while 34 said they wanted to use it in the future.

Looking at industry sectors, we again find finance, consulting, and healthcare dominating.6

(#fig:unnamed-chunk-6)Number of users reporting to use Spark in industry. Smaller groups not displayed.

What do survey respondents do with Spark? Analyses of tabular data and time series dominate:

(#fig:unnamed-chunk-7)Number of users reporting to use Spark in industry. Smaller groups not displayed.

Frameworks and skills

As with deep learning, we wanted to know what language people use to do Spark. If you look at the below graphic, you see R appearing twice: once in connection with sparklyr, once with SparkR. What’s that about?

Both sparklyr and SparkR are R interfaces for Apache Spark, each designed and built with a different set of priorities and, consequently, trade-offs in mind.

sparklyr, one the one hand, will appeal to data scientists at home in the tidyverse, as they’ll be able to use all the data manipulation interfaces they’re familiar with from packages such as dplyr, DBI, tidyr, or broom.

SparkR, on the other hand, is a light-weight R binding for Apache Spark, and is bundled with the same. It’s an excellent choice for practitioners who are well-versed in Apache Spark and just need a thin wrapper to access various Spark functionalities from R.

(#fig:unnamed-chunk-8)Language / language bindings used to do Spark.

When asked to rate their expertise in R7 and Spark, respectively, respondents showed similar behavior as observed for deep learning above: Most people seem to think more of their R skills than their theoretical Spark-related knowledge. However, even more caution should be exercised here than above: The number of responses here was significantly lower.

(#fig:unnamed-chunk-9)Self-rated skills re R and Spark.

Wishes and suggestions

Just like with DL, Spark users were asked what could be improved, and what they were hoping for.

Interestingly, answers were less “clustered” than for DL. While with DL, a few things cropped up again and again, and there were very few mentions of concrete technical features, here we see about the opposite: The great majority of wishes were concrete, technical, and often only came up once.

Probably though, this is not a coincidence.

Looking back at how sparklyr has evolved from 2016 until now, there is a persistent theme of it being the bridge that joins the Apache Spark ecosystem to numerous useful R interfaces, frameworks, and utilities (most notably, the tidyverse).

Many of our users’ suggestions were essentially a continuation of this theme. This holds, for example, for two features already available as of sparklyr 1.4 and 1.2, respectively: support for the Arrow serialization format and for Databricks Connect. It also holds for tidymodels integration (a frequent wish), a simple R interface for defining Spark UDFs (frequently desired, this one too), out-of-core direct computations on Parquet files, and extended time-series functionalities.

We’re thankful for the feedback and will evaluate carefully what could be done in each case. In general, integrating sparklyr with some feature X is a process to be planned carefully, as modifications could, in theory, be made in various places (sparklyr; X; both sparklyr and X; or even a newly-to-be-created extension). In fact, this is a topic deserving of much more detailed coverage, and has to be left to a future post.

Ethics and AI in society

To start, this is probably the section that will profit most from more preparation, the next time we do this survey. Due to time pressure, some (not all!) of the questions ended up being too suggestive, possibly resulting in social-desirability bias.8

Next time, we’ll try to avoid this, and questions in this area will likely look pretty different (more like scenarios or what-if stories)9. However, I was told by several people they’d been positively surprised by simply encountering this topic at all in the survey. So perhaps this is the main point – although there are a few results that I’m sure will be interesting by themselves!

Anticlimactically, the most non-obvious results are presented first.10

“Are you worried about societal/political impacts of how AI is used in the real world?”

For this question, we had four answer options, formulated in a way that left no real “middle ground”. (The labels in the graphic below verbatim reflect those options.)

(#fig:unnamed-chunk-10)Number of users responding to the question ‘Are you worried about societal/political impacts of how AI is used in the real world?’ with the answer options given.

The next question is definitely one to keep for future editions, as from all questions in this section, it definitely has the highest information content.

“When you think of the near future, are you more afraid of AI misuse or more hopeful about positive outcomes?”

Here, the answer was to be given by moving a slider, with -100 signifying “I tend to be more pessimistic”; and 100, “I tend to be more optimistic”. Although it would have been possible to remain undecided, choosing a value close to 0, we instead see a bimodal distribution:

(#fig:unnamed-chunk-11)When you think of the near future, are you more afraid of AI misuse or more hopeful about positive outcomes?

Why worry, and what about

The following two questions are those already alluded to as possibly being overly prone to social-desirability bias. They asked what applications people were worried about, and for what reasons, respectively. Both questions allowed to select however many responses one wanted, intentionally not forcing people to rank things that are not comparable (the way I see it). In both cases though, it was possible to explicitly indicate None (corresponding to “I don’t really find any of these problematic” and “I am not extensively worried”, respectively.)

What applications of AI do you feel are most problematic?

(#fig:unnamed-chunk-12)Number of users selecting the respective application in response to the question: What applications of AI do you feel are most problematic?

If you are worried about misuse and negative impacts, what exactly is it that worries you?

(#fig:unnamed-chunk-13)Number of users selecting the respective impact in response to the question: If you are worried about misuse and negative impacts, what exactly is it that worries you?

Complementing these questions, it was possible to enter further thoughts and concerns in free-form. Although I can’t cite everything that was mentioned here, recurring themes were:

• Misuse of AI to the wrong purposes, by the wrong people, and at scale.

• Not feeling responsible for how one’s algorithms are used (the I’m just a software engineer topos).

• Reluctance, in AI but in society overall as well, to even discuss the topic (ethics).

Finally, although this was mentioned just once, I’d like to relay a comment that went in a direction absent from all provided answer options, but that probably should have been there already: AI being used to construct social credit systems.

“It’s also that you somehow might have to learn to game the algorithm, which will make AI application forcing us to behave in some way to be scored good. That moment scares me when the algorithm is not only learning from our behavior but we behave so that the algorithm predicts us optimally (turning every use case around).”11

Conclusion

This has become a long text. But I think that seeing how much time respondents took to answer the many questions, often including lots of detail in the free-form answers, it seemed like a matter of decency to, in the analysis and report, go into some detail as well.

Thanks again to everyone who took part! We hope to make this a recurring thing, and will strive to design the next edition in a way that makes answers even more information-rich.

Thanks for reading!

Posted by Andy Zeng and Pete Florence, Research Scientists, Robotics at Google

Rearranging objects (such as organizing books on a bookshelf, moving utensils on a dinner table, or pushing piles of coffee beans) is a fundamental skill that can enable robots to physically interact with our diverse and unstructured world. While easy for people, accomplishing such tasks remains an open research challenge for embodied machine learning (ML) systems, as it requires both high-level and low-level perceptual reasoning. For example, when stacking a pile of books, one might consider where the books should be stacked, and in which order, while ensuring that the edges of the books align with each other to form a neat pile.

Across many application areas in ML, simple differences in model architecture can exhibit vastly different generalization properties. Therefore, one might ask whether there are certain deep network architectures that favor simple underlying elements of the rearrangement problem. Convolutional architectures, for example, are common in computer vision as they encode translational invariance, yielding the same response even if an image is shifted, while Transformer architectures are common in language processing because they exploit self-attention to capture long-range contextual dependencies. In robotics applications, one common architectural element is to use object-centric representations such as poses, keypoints, or object descriptors inside learned models, but these representations require additional training data (often manually annotated) and struggle to describe difficult scenarios such as deformables (e.g., playdough), fluids (honey), or piles of stuff (chopped onions).

Today, we present the Transporter Network, a simple model architecture for learning vision-based rearrangement tasks, which appeared as a publication and plenary talk during CoRL 2020. Transporter Nets use a novel approach to 3D spatial understanding that avoids reliance on object-centric representations, making them general for vision-based manipulation but far more sample efficient than benchmarked end-to-end alternatives. As a consequence, they are fast and practical to train on real robots. We are also releasing an accompanying open-source implementation of Transporter Nets together with Ravens, our new simulated benchmark suite of ten vision-based manipulation tasks.

Transporter Networks: Rearranging the Visual World for Robotic Manipulation
The key idea behind the Transporter Network architecture is that one can formulate the rearrangement problem as learning how to move a chunk of 3D space. Rather than relying on an explicit definition of objects (which is bound to struggle at capturing all edge cases), 3D space is a much broader definition for what could serve as the atomic units being rearranged, and can broadly encompass an object, part of an object, or multiple objects, etc. Transporter Nets leverage this structure by capturing a deep representation of the 3D visual world, then overlaying parts of it on itself to imagine various possible rearrangements of 3D space. It then chooses the rearrangements that best match those it has seen during training (e.g., from expert demonstrations), and uses them to parameterize robot actions. This formulation allows Transporter Nets to generalize to unseen objects and enables them to better exploit geometric symmetries in the data, so that they can extrapolate to new scene configurations. Transporter Nets are applicable to a wide variety of rearrangement tasks for robotic manipulation, expanding beyond our earlier models, such as affordance-based manipulation and TossingBot, that focus only on grasping and tossing.

Ravens Benchmark
To evaluate the performance of Transporter Nets in a consistent environment for fair comparisons to baselines and ablations, we developed Ravens, a benchmark suite of ten simulated vision-based rearrangement tasks. Ravens features a Gym API with a built-in stochastic oracle to evaluate the sample efficiency of imitation learning methods. Ravens avoids assumptions that cannot transfer to a real setup: observation data contains only RGB-D images and camera parameters; actions are end effector poses (transposed into joint positions with inverse kinematics).

Experiments on these ten tasks show that Transporter Nets are orders of magnitude more sample efficient than other end-to-end methods, and are capable of achieving over 90% success on many tasks with just 100 demonstrations, while the baselines struggle to generalize with the same amount of data. In practice, this makes collecting enough demonstrations a more viable option for training these models on real robots (which we show examples of below).

Our new Ravens benchmark includes ten simulated vision-based manipulation tasks, including pushing and pick-and-place, for which experiments show that Transporter Nets are orders of magnitude more sample efficient than other end-to-end methods. Ravens features a Gym API with a built-in stochastic oracle to evaluate the sample efficiency of imitation learning methods.

Our new Ravens benchmark includes ten simulated vision-based manipulation tasks, including pushing and pick-and-place, for which experiments show that Transporter Nets are orders of magnitude more sample efficient than other end-to-end methods. Ravens features a Gym API with a built-in stochastic oracle to evaluate the sample efficiency of imitation learning methods.

Highlights
Given 10 example demonstrations, Transporter Nets can learn pick and place tasks such as stacking plates (surprisingly easy to misplace!), multimodal tasks like aligning any corner of a box to a marker on the tabletop, or building a pyramid of blocks.

By leveraging closed-loop visual feedback, Transporter Nets have the capacity to learn various multi-step sequential tasks with a modest number of demonstrations: such as moving disks for Tower of Hanoi, palletizing boxes, or assembling kits of new objects not seen during training. These tasks have considerably “long horizons”, meaning that to solve the task the model must correctly sequence many individual choices. Policies also tend to learn emergent recovery behaviors.

One surprising thing about these results was that beyond just perception, the models were starting to learn behaviors that resemble high-level planning. For example, to solve Towers of Hanoi, the models have to pick which disk to move next, which requires recognizing the state of the board based on the current visible disks and their positions. With a box-palletizing task, the models must locate the empty spaces of the pallet, and identify how new boxes can fit into those voids. Such behaviors are exciting because they suggest that with all the baked-in invariances, the model can focus its capacity on learning the more high-level patterns in manipulation.

Transporter Nets can also learn tasks that use any motion primitive defined by two end effector poses, such as pushing piles of small objects into a target set, or reconfiguring a deformable rope to connect the two end-points of a 3-sided square. This suggests that rigid spatial displacements can serve as useful priors for nonrigid ones.

Conclusion
Transporter Nets bring a promising approach to learning vision-based manipulation, but are not without limitations. For example, they can be susceptible to noisy 3D data, we have only demonstrated them for sparse waypoint-based control with motion primitives, and it remains unclear how to extend them beyond spatial action spaces to force or torque-based actions. But overall, we are excited about this direction of work, and we hope that it provides inspiration for extensions beyond the applications we’ve discussed. For more details, please check out our paper.

Acknowledgements
This research was done by Andy Zeng, Pete Florence, Jonathan Tompson, Stefan Welker, Jonathan Chien, Maria Attarian, Travis Armstrong, Ivan Krasin, Dan Duong, Vikas Sindhwani, and Johnny Lee, with special thanks to Ken Goldberg, Razvan Surdulescu, Daniel Seita, Ayzaan Wahid, Vincent Vanhoucke, Anelia Angelova, Kendra Byrne, for helpful feedback on writing; Sean Snyder, Jonathan Vela, Larry Bisares, Michael Villanueva, Brandon Hurd for operations and hardware support; Robert Baruch for software infrastructure, Jared Braun for UI contributions; Erwin Coumans for PyBullet advice; Laura Graesser for video narration.