Pre-Training for Robots: Offline RL Enables Learning New Tasks from a Handful of Trials

TL;DR

PTR uses CQL pre-training on the Bridge Dataset to learn unseen tasks on a real WidowX robot from only 10–15 demonstrations.

cs.RO 🔴 Advanced 2022-10-11 20 views
Aviral Kumar Anikait Singh Frederik Ebert Mitsuhiko Nakamoto Yanlai Yang Chelsea Finn Sergey Levine
offline reinforcement learning robot pre-training Conservative Q-Learning few-shot adaptation real-world robotics

Key Findings

Methodology

Pre-Training for Robots (PTR) first performs task-conditioned offline RL on diverse Bridge Dataset demonstrations, then fine-tunes the same CQL agent on a small target-task dataset. Conservative Q-Learning penalizes high values for policy actions relative to dataset actions, limiting extrapolation error on unseen actions. PTR combines a GroupNorm ResNet, learned spatial embeddings, and action injection into every fully connected decoder layer, while batch mixing controls the balance between prior and target data.

Key Results

  • On a real WidowX robot, PTR reportedly achieves the first RL-based demonstrations of learning new tasks in a new domain from as few as 10 target demonstrations. The target skills and environments are absent from pre-training, yet prior multi-task experience remains useful after fine-tuning.
  • Offline fine-tuning uses τ=0.8 to mix Bridge and target transitions; online fine-tuning mixes offline data with newly collected rollouts and performs 10 gradient steps per collected transition. PTR outperforms behavioral cloning, representation-pretraining methods, and RL baselines without pre-training.
  • Ablations identify GroupNorm, learned spatial embeddings, and repeated action injection as critical. The paper reports more than a 3.5× improvement in final real-world performance from the combined design choices; Q-value monotonicity also helps select checkpoints and avoid overtraining.

Significance

The paper addresses the high cost of collecting task-specific robot data and the weak transfer of narrow policies. It shows that large offline robot datasets can train not only imitation policies or visual representations, but also transferable value functions that reason about action consequences. Academically, PTR unifies multi-task pre-training, offline RL, and few-shot adaptation. Industrially, it suggests that new skills may require only a small calibration dataset rather than a complete data-collection campaign, while still allowing autonomous improvement from limited interaction.

Technical Contribution

PTR does not introduce a new offline-RL objective; its contribution is a careful system design that makes CQL viable for large-scale visual robot pre-training. GroupNorm replaces unstable BatchNorm in TD-trained vision networks. Learned spatial embeddings preserve image geometry instead of globally averaging feature maps. Actions are injected into every decoder layer so the Q-function cannot ignore them under narrow demonstrations. Batch ratios τ and β prevent target or online data from being overwhelmed. A held-out trajectory Q-value trend provides practical checkpoint selection.

Novelty

The central novelty is demonstrating that end-to-end offline RL can serve as a practical pre-training substrate for real-robot adaptation. Unlike approaches centered on behavioral cloning, representation learning, or thousands of online trials, PTR uses the same CQL mechanism for both phases and adapts to unseen tasks with at most 15 demonstrations. The authors characterize it as the first RL method to achieve this setting on a real WidowX robot.

Limitations

  • The experiments largely assume that target tasks and environments come from a related distribution and that demonstrations are successful. Under radically different cameras, dynamics, embodiments, or failure-heavy data, conservative values and the Q-monotonicity heuristic may not transfer.
  • Ten to fifteen demonstrations or online interaction are still required, while one-hot task identifiers provide little structured compositionality for open-world task combinations.
  • Checkpoint selection remains partly visual, and the paper does not fully quantify compute cost, scaling limits, or transfer across robot platforms.

Future Work

Future work should investigate language- or structure-conditioned task representations, automatic reward inference, compositional skills, and cross-embodiment transfer. Reliable automatic Q-monotonicity metrics, risk-sensitive offline RL, and stronger out-of-distribution detection are needed. Larger real-world datasets should also clarify computational efficiency, safety, and whether autonomous online improvement remains stable over long horizons.

AI Executive Summary

Teaching a robot a new skill usually requires many demonstrations or expensive trial-and-error. Behavioral cloning can reproduce observed motions but may not understand their consequences, while offline RL often overestimates actions absent from a narrow dataset. Robotics therefore lacks a dependable equivalent of supervised learning’s “pre-train, then rapidly adapt” paradigm.

Kumar and colleagues propose PTR, which pre-trains Conservative Q-Learning on the diverse Bridge Dataset and then fine-tunes on a new task. The practical breakthrough lies in architectural choices: GroupNorm stabilizes visual temporal-difference learning; learned spatial embeddings preserve object location; and action vectors are injected into every decoder layer so the Q-function compares actions rather than exploiting state-action correlations. During adaptation, τ=0.8 batch mixing keeps scarce target data visible, while β-style mixing supports online rollouts.

On a real WidowX robot, PTR learns unseen tasks from only 10–15 demonstrations and outperforms behavioral cloning, representation-pretraining, and no-pretraining RL baselines. The combined design choices improve final real-world performance by more than 3.5×, and limited autonomous interaction can further improve behavior without demonstrations. The result is not universal robot intelligence: it relies on related distributions, heuristic reward labeling, and costly training. Its importance is more practical—a broad offline dataset can become a reusable skill prior, turning new-task data collection from full policy training into targeted calibration.

Deep Analysis

Background

Robot learning has advanced through behavioral cloning, visual representation learning, and online RL, but every new task still demands costly real-world data. The Bridge Dataset offers multi-task demonstrations across toy kitchens and is a natural source for robot pre-training. Earlier methods mainly learn representations or imitate actions; they do not necessarily model action consequences. Offline RL is more general in principle, but visual Q-learning is vulnerable to distributional extrapolation and unstable normalization.

Core Problem

The goal is to use a broad offline dataset to solve target tasks absent from that dataset, given only 10–15 demonstrations or limited autonomous interaction. The bottlenecks are severe data imbalance, Q-functions that ignore actions because demonstrations correlate states and actions, BatchNorm-related TD instability, and the adaptation trade-off between under-training and catastrophic loss of pre-trained generalization.

Innovation

PTR applies one task-conditioned CQL system to both pre-training and adaptation. First, GroupNorm replaces BatchNorm to stabilize TD learning. Second, learned spatial embeddings preserve feature-map geometry, unlike global average pooling. Third, actions enter every fully connected decoder layer, improving action discrimination. Fourth, τ and β explicitly control prior, target, and online data exposure. Finally, target-trajectory Q-value trends provide a practical cross-validation signal for selecting checkpoints.

Methodology

  • �� Pre-training: learn Qϕ(s,a;i) and πθ(a|s,i), conditioned on a one-hot task identifier. CQL combines Bellman error with a conservative penalty that suppresses policy-action values relative to dataset-action values.
  • �� Policy improvement: maximize E[Qϕ]+βH(πθ), using a delayed target network for stability.
  • �� Visual encoding: use ResNet-34, GroupNorm, and learned spatial embeddings.
  • �� Decoding: inject the action vector into every MLP layer and provide the task identifier to the decoder.
  • �� Offline adaptation: mix Bridge and target transitions with τ=0.8 in most real-world experiments.
  • �� Online adaptation: mix offline and rollout data, applying 10 gradient updates per collected transition.
  • �� Rewards: label the final three transitions of successful demonstrations +1 and all others 0.

Experiments

Experiments use the Bridge Dataset, additional tasks, toy-kitchen environments, and a real WidowX robot. Baselines include behavioral cloning, representation-pretraining followed by imitation learning, RL without pre-training, and related meta-RL approaches. Target tasks generally provide only 10–15 demonstrations; online settings permit autonomous rollout collection. Main settings use CQL α=10.0 and τ=0.8, while online fine-tuning uses α=0.5. Ablations test network capacity, normalization, spatial embeddings, and action injection.

Results

PTR achieves what the authors describe as the first real-robot RL result in this setting with as few as 10 demonstrations, outperforming BC, representation-pretraining, and no-pretraining approaches across comparisons. The combined design decisions yield more than a 3.5× improvement in final real-world performance. Larger networks continue to help, including ResNet-50-scale models. Q-value inspection separates undertrained from overtrained checkpoints: a flatter curve fails on door opening, whereas a visibly increasing trajectory succeeds.

Applications

PTR is suited to warehouse picking, kitchen manipulation, laboratory automation, and flexible assembly. A company can pre-train on historical multi-task trajectories, then collect a small number of demonstrations for a new object, container, or placement goal. Requirements include related visual and dynamical distributions, compatible action spaces, successful demonstrations or safe interaction, and sufficient compute for visual offline RL. The industrial benefit is lower teleoperation cost per new skill or product.

Limitations & Outlook

Transfer is not guaranteed across unfamiliar cameras, robot morphologies, dynamics, or object distributions. One-hot task IDs offer limited compositional generalization and may scale poorly as task inventories grow. The +1 reward on the final three demonstration steps is a useful heuristic, not a general reward solution. Visual CQL training remains computationally demanding. Future systems should add automatic model selection, language conditioning, safe exploration, uncertainty estimation, and systematic cross-platform evaluation.

Plain Language Accessible to non-experts

Imagine a cook who has worked in many kitchens. Before being asked to prepare a new dish, the cook watches thousands of successful recordings: placing potatoes, eggs, or sushi into different pots, on different counters, under different camera angles. The cook does not merely memorize hand motions; they learn which choices usually lead to a successful finish. CQL acts like a cautious coach: if a choice was never supported by the recordings, it does not automatically receive a high score.

After building this general experience, the cook needs only 10–15 demonstrations of the new dish. Training keeps revisiting old recordings while repeatedly showing the new ones, so the new lesson is not drowned out by the large archive. The system also preserves where objects are located and makes sure the decision-maker pays attention to the proposed action, not just the current picture.

If allowed to try independently, the robot can add new attempts to its notebook and improve. The lesson is encouraging but limited: experience transfers well between similar kitchens, not necessarily to every kitchen, camera, or robot body.

ELI14 Explained like you're 14

Think of a robot game where every new level normally requires starting from zero. PTR first lets the robot play lots of related levels: picking up cups, moving fruit, and putting objects into pots. It learns a score for each possible move—basically, “If I do this now, how likely am I to finish successfully later?”

The main algorithm is called Conservative Q-Learning, or CQL. It is like a careful gaming coach. If the robot has never seen a move work, the coach does not declare it amazing just because the robot feels confident. That prevents the robot from chasing imaginary high scores. The robot also remembers where objects are on the screen and keeps telling its judge which move it is considering.

Now a brand-new level appears. Researchers show only 10–15 examples, such as putting sushi into a metal pot it never saw before. Instead of forgetting everything, the robot practices with both its old experience and the new examples. It can even try a few moves itself and learn from the results. Pretty useful, right? But it may still struggle in a completely strange room, with a new camera, or on a robot with a totally different body.

Glossary

Conservative Q-Learning (CQL)

An offline-RL algorithm that lowers values for actions insufficiently supported by the dataset, reducing overestimation. Technically, it adds a conservative regularizer to Bellman-error training.

PTR uses multi-task CQL for both Bridge pre-training and target-task fine-tuning.

Offline reinforcement learning

Learning a policy from a fixed dataset without continual environment interaction. It reduces real-robot trial cost but is constrained by dataset coverage.

PTR performs offline pre-training and, in one setting, entirely offline adaptation.

Bridge Dataset

A diverse robot-manipulation dataset containing demonstrations across toy kitchens, tasks, and viewpoints. The paper overview cites roughly 10 domains, 100 tasks, and 12,000 demonstrations.

It provides PTR’s principal pre-training data.

Q-function

A function estimating discounted long-term return from a state-action pair. It guides policy improvement by ranking candidate actions.

PTR learns Qϕ(s,a;i), conditioned on the task identifier i.

Group Normalization

A feature normalization layer that computes statistics within channel groups rather than across a minibatch. It is often more stable than BatchNorm for TD-trained visual networks.

PTR uses GroupNorm to prevent ResNet Q-network divergence.

Learned spatial embedding

A learned spatial weighting mechanism that aggregates feature maps while retaining location information. It avoids the positional information loss of global average pooling.

It forms part of PTR’s visual Q-function architecture.

Open Questions Unanswered questions from this research

  • 1 The method is strong mainly within related distributions; systematic evidence is still missing for radically novel objects, cameras, dynamics, and failure-heavy datasets.
  • 2 One-hot task identifiers provide little compositional structure. Language, causal task graphs, or reusable skill representations may be needed for open-world combinations.
  • 3 Q-value monotonicity is selected largely by visual inspection. A reliable automated metric robust to noisy demonstrations remains unresolved.

Applications

Immediate Applications

Flexible grasping and kitchen manipulation

Robot developers can pre-train on historical grasp-and-place demonstrations, then collect 10–15 successful examples for a new object or container. The approach is most suitable when robot hardware, camera geometry, and task distributions remain related; expected benefits include less teleoperation and faster deployment.

Laboratory and warehouse automation

Labs and warehouses can aggregate offline trajectories from multiple stations, then adapt to a new kit, package, or placement location. If safe rollouts are available, online mixing can incorporate autonomous trials and improve behavior beyond the initial demonstrations.

Long-term Vision

Transferable robot skill foundation models

A longer-term vision is a cross-robot, cross-scene skill model conditioned by language or explicit goals. Realizing it will require larger datasets, automatic reward inference, robust out-of-distribution safety checks, and inexpensive adaptation across embodiments.

Abstract

Progress in deep learning highlights the tremendous potential of utilizing diverse robotic datasets for attaining effective generalization and makes it enticing to consider leveraging broad datasets for attaining robust generalization in robotic learning as well. However, in practice, we often want to learn a new skill in a new environment that is unlikely to be contained in the prior data. Therefore we ask: how can we leverage existing diverse offline datasets in combination with small amounts of task-specific data to solve new tasks, while still enjoying the generalization benefits of training on large amounts of data? In this paper, we demonstrate that end-to-end offline RL can be an effective approach for doing this, without the need for any representation learning or vision-based pre-training. We present pre-training for robots (PTR), a framework based on offline RL that attempts to effectively learn new tasks by combining pre-training on existing robotic datasets with rapid fine-tuning on a new task, with as few as 10 demonstrations. PTR utilizes an existing offline RL method, conservative Q-learning (CQL), but extends it to include several crucial design decisions that enable PTR to actually work and outperform a variety of prior methods. To our knowledge, PTR is the first RL method that succeeds at learning new tasks in a new domain on a real WidowX robot with as few as 10 task demonstrations, by effectively leveraging an existing dataset of diverse multi-task robot data collected in a variety of toy kitchens. We also demonstrate that PTR can enable effective autonomous fine-tuning and improvement in a handful of trials, without needing any demonstrations. An accompanying overview video can be found in the supplementary material and at thi URL: https://sites.google.com/view/ptr-final/

cs.RO cs.LG