PaperlayerView exact revision on arXiv

Source-linked AI summary

Imitating Task and Motion Planning with Visuomotor Transformers

Murtaza Dalal, Ajay Mandlekar, Caelan Garrett, Ankur Handa, Ruslan Salakhutdinov, Dieter Fox

arXiv:2305.16309v3cs.ROcs.CVcs.LG

TL;DR

Robot manipulation learning is constrained by the difficulty of collecting useful, diverse data at scale. OPTIMUS imitates TAMP demonstrations with visuomotor Transformer policies, and the paper reports success across hundreds of long-horizon tasks involving up to 72 objects.

  • Problem

    Robot manipulation lacks scalable access to useful, task-relevant, diverse data, while human-supervised collection is labor-intensive.

    Source: §1 · p. 1

  • Method

    OPTIMUS distills TAMP demonstrations into closed-loop visuomotor Transformer policies trained on curated data.

    Source: §1 · p. 2, §3.2 · p. 4

  • Results

    OPTIMUS outperforms the best baselines across PickPlace, Shelf, and Microwave, reaching 75%, 48%, and 47% on their 72-object variants.

    Source: §4.1 · p. 8

  • Takeaways & Limitations

    The trained policies solve challenging long-horizon manipulation tasks without task, state, or environment knowledge.

    Source: §5 · p. 8

    OPTIMUS requires TAMP to solve the task during training, which may limit tasks involving dynamics or contact-rich manipulation that traditional TAMP finds challenging.

    Source: §5 · p. 8

Abstract

from arXiv · show

Imitation learning is a powerful tool for training robot manipulation policies, allowing them to learn from expert demonstrations without manual programming or trial-and-error. However, common methods of data collection, such as human supervision, scale poorly, as they are time-consuming and labor-intensive. In contrast, Task and Motion Planning (TAMP) can autonomously generate large-scale datasets of diverse demonstrations. In this work, we show that the combination of large-scale datasets generated by TAMP supervisors and flexible Transformer models to fit them is a powerful paradigm for robot manipulation. To that end, we present a novel imitation learning system called OPTIMUS that trains large-scale visuomotor Transformer policies by imitating a TAMP agent. OPTIMUS introduces a pipeline for generating TAMP data that is specifically curated for imitation learning and can be used to train performant transformer-based policies. In this paper, we present a thorough study of the design decisions required to imitate TAMP and demonstrate that OPTIMUS can solve a wide variety of challenging vision-based manipulation tasks with over 70 different objects, ranging from long-horizon pick-and-place tasks, to shelf and articulated object manipulation, achieving 70 to 80% success rates. Video results and code at https://mihdalal.github.io/optimus/

1 Introduction

Robot manipulation learning lacks scalable access to diverse, task-relevant data. OPTIMUS uses TAMP-generated demonstrations to train visuomotor Transformer policies for challenging, long-horizon manipulation.

1 Introduction

  • Robot manipulation has not yet matched NLP and computer vision in large-scale data-driven learning, in part because useful, diverse data is difficult to collect.Robot data collection requires real-time interaction, and policies need diverse task instances to generalize broadly.Source: §1 · p. 1
  • Human teleoperation can demand substantial labor—RT-1 required 1.5 years of data collection—while reinforcement learning is sample-inefficient and computationally costly.TAMP offers an alternative by autonomously producing task-relevant, diverse demonstrations across task instances.Source: §1 · pp. 1–2
  • OPTIMUS policies solve over 300 long-horizon manipulation tasks involving up to 8 stages and 72 objects, with success rates above 70%.The system imitates a TAMP agent using visuomotor Transformer policies.Source: §1 · p. 2

2 Preliminaries

The preliminaries frame the problem as sparse-reward partial observability and describe how an observable-state TAMP planner produces demonstrations for imitation learning.

2 Preliminaries

  • The paper formulates its setting as a sparse-reward POMDP, with reward zero at terminal goal states and negative one elsewhere.Policies use current observations and histories of observations and actions.Source: §2 · pp. 2–3
  • For behavior cloning, a policy learns actions from observation-action histories in expert trajectories.The paper situates this offline learning setting alongside prior human-demonstration imitation work and distinguishes its TAMP supervisor.Source: §2 · pp. 2–3
  • TAMP plans object interactions, their grasp and placement constraints, and safe joint motions using an approximate deterministic model of the observable system state.A feedback controller tracks the plan and periodically replans.Source: §2 · p. 3
  • A shelf task illustrates TAMP’s parameterized move, pick, and place actions, whose plans search both action sequences and feasible motion parameters.The searched parameters must satisfy grasp, stability, kinematic, and collision constraints.Source: §2 · p. 3

3 Designing a TAMP Imitation System

OPTIMUS builds an imitation-learning pipeline around TAMP demonstrations, shaping plans and data for visuomotor policy learning across diverse manipulation tasks.

3 Designing a TAMP Imitation System

  • The system applies a TAMP supervisor to procedurally generated tasks spanning block stacking, pick-and-place, shelf manipulation, and articulated microwave manipulation.Tasks vary in objects, poses, and configurations; objects are drawn from ShapeNet.Source: §3 · pp. 3–4, §3.1 · p. 4, Figure 4 · p. 6

3.1 Cost-Minimizing TAMP

  • The two-stage planner first searches for low joint-space travel costs, then smooths BiRRT motions with cubic-spline short cutting under velocity and acceleration limits.The anytime planning stage has a five-second computation budget; the trajectory-smoothing procedure converges to a locally time-optimal trajectory.Source: §3.1 · p. 4

3.2 Generating Imitation Data from TAMP

  • OPTIMUS addresses privileged TAMP inputs, joint-space actions, and potentially long task-space paths by using visual observations, task-space control, and curated demonstrations.The data pipeline uses multiple camera views, including a wrist camera, and prunes trajectories that exceed two standard deviations from mean task-space length or leave a visible-workspace bounding box.Source: §3.2 · pp. 4–5

3.3 Training Imitation Policies at Scale

  • OPTIMUS combines a Transformer policy over observation histories with encoded camera images and proprioception to predict low-level task-space actions.Images use a ResNet-18 backbone; proprioception is encoded with an MLP, and the resulting action distribution models multimodal demonstrations with a five-component GMM.Source: §3.3 · pp. 5–6, Figure 3 · p. 5

4 Experimental Evaluation

The evaluation tests OPTIMUS and baselines on single- and multi-task manipulation, using unseen initial states and task success rate. OPTIMUS performs strongly across long-horizon and multi-object tasks, while performance falls as task difficulty and object count increase.

4 Experimental Evaluation

  • Evaluation spans block stacking, pick-and-place, shelf, and microwave tasks; single-task tests use 50 problems and three seeds, while multi-task tests use 10 problems per task and one seed.The multi-task datasets cover 19 or 72 objects, with 1K demonstrations per object.Source: §4 · p. 6

4.1 Learning Results

  • OPTIMUS solves long-horizon tasks with up to eight pick-and-place operations at 70% success and 40 TAMP high-level actions at 60%, outperforming baselines on block stacking.BC-RNN outperforms OPTIMUS on PickPlaceTwo, while OPTIMUS improves substantially on PickPlaceThree and PickPlaceFour.Source: Figure 5 · p. 7, §4.1 · p. 7
  • On Stack, OPTIMUS and baselines all reach 100%; on PickPlace-1, OPTIMUS reaches 100% while baselines achieve up to 97%.The authors note that BeT struggles to fit TAMP data on most tasks and hypothesize that inaccurate cluster centers contribute.Source: §4.1 · p. 7
  • On Robosuite PickPlace, OPTIMUS reaches 90% success versus Guided TAMP’s 88%, without privileged state, fixed ground actions, or online supervision.The comparison uses 25K demonstrations and four fixed objects placed together in one bin initially.Source: Table 1 · p. 7, §4.1 · p. 7
  • OPTIMUS achieves 91% success on Shelf-1 and 86% on Microwave-1, outperforming baselines by nearly 10% on the microwave task.These tasks require obstacle awareness and, for Microwave-1, manipulation of an articulated object.Source: §4.1 · p. 7
  • Scene adaptation yields 96% success on StackAdapt and 75% on MicrowaveAdapt, while held-out shelf and microwave sizes reach 80% and 70%.Source: §4.1 · p. 8
  • OPTIMUS beats the best baseline across 19- and 72-object PickPlace, Shelf, and Microwave tasks, though multi-task success trails single-task performance.For 19-object tasks, OPTIMUS scores 85%, 66%, and 61% versus 61%, 48%, and 41%; for 72-object tasks, it scores 75%, 48%, and 47% versus 50%, 36%, and 29%.Source: Table 2 · p. 8, §4.1 · p. 8
  • Across single-task datasets, OPTIMUS nearly doubles the TAMP supervisor’s average success rate (87% vs. 52%) and runs 5–7.5x faster (21/31ms vs. 150ms per action).Runtime was averaged per step across 100 trials.Source: §4.1 · p. 8

5 Limitations and Future Work

OPTIMUS depends on TAMP being able to solve a task during training, and performance drops as task difficulty and object count rise. The paper identifies real-world transfer and improved multi-task learning as directions for future work.

5 Limitations and Future Work

  • OPTIMUS can only solve tasks that TAMP can solve during training, potentially limiting use on dynamics-dependent or contact-rich tasks that challenge traditional TAMP.The authors note that integrated task planning and skill learning may extend TAMP supervision to examples such as pouring, scooping, stirring, peg insertion, or coffee making.Source: §5 · p. 8
  • Success drops significantly as task difficulty and object count increase, especially for multi-task Shelf and Microwave, leaving a gap for future work.Source: §5 · p. 8
  • Real-world extensions proposed include collecting TAMP data with pose estimation in a controlled setting or transferring simulation-trained imitation policies to real data.The authors leave these extensions to future work.Source: §5 · p. 8

A Table of Contents

  • The appendices cover additional learning results, ablations, environments, agent structure, experiment details, and related work.Source: A Table of Contents · p. 17

B Additional Learning Results

OPTIMUS is evaluated on multi-task category learning, scene adaptation, reinforcement-learning and Transformer baselines, and comparisons with its TAMP supervisor. These results also report execution-time advantages over TAMP.

B Additional Learning Results

  • 73% success: one OPTIMUS Transformer learns PickPlace, Shelf, and Microwave across object categories, outperforming the multi-task baselines.The dataset contains 15K trajectories, with five objects per category and 1K demonstrations per task.Source: B Additional Learning Results · p. 18, Table B.1 · p. 18
  • OPTIMUS adapts task plans to scene configurations and generalizes to held-out receptacle sizes, reaching 96% and 75% on adaptation tasks and 80% and 70% on held-out shelves and microwaves.For MicrowaveAdapt, the best baseline reaches 40%, compared with OPTIMUS at 75%.Source: B Additional Learning Results · p. 18, Table B.2 · p. 18
  • Across four Robosuite tasks, SAC, DRQ-v2, and MoDem each fail to exceed 10% success, with exploration identified as a significant challenge.The comparison trains each RL method with up to five million samples and five seeds.Source: B Additional Learning Results · p. 18, Table B.3 · p. 19
  • OPTIMUS outperforms Transformer-in-Transformer by an average of 16.8% across five tasks and executes 4–5x faster.The authors hypothesize that ResNets provide a useful spatial-locality bias for visuomotor control.Source: Table B.4 · p. 19, B Additional Learning Results · p. 19
  • OPTIMUS outperforms TAMP by 20% on PickPlace, 64% on Microwave-1, and 44% on PickPlaceFour; averaged success is 87% versus 52% for joint-space TAMP.The authors attribute this partly to distilling successful trajectories, which avoids learning observed grasp-execution failures and object slippage.Source: B Additional Learning Results · p. 19, Table B.5 · p. 20
  • 150 ms per action for TAMP versus 21 ms for 30M-parameter OPTIMUS and 31 ms for 100M-parameter OPTIMUS; OPTIMUS is 5–7.5x faster on average.TAMP additionally incurs a 2–5 second up-front cost, whereas OPTIMUS takes a constant time per action.Source: B Additional Learning Results · pp. 19–20, Table B.6 · p. 20

C Ablations

The ablations examine action space, trajectory filtering, gripper and stall handling, camera views, and output losses. The accompanying environment and agent descriptions specify the tasks and observation and control choices used.

C Ablations

  • 100% versus 86%: visual policies using task-space versus joint-space control on Microwave-1; for state-based policies, the corresponding rates are 100% and 98%.The authors hypothesize that visual policies find it difficult to learn the inverse-kinematics mapping from image locations to joint angles.Source: C Ablations · p. 21, Figure C.1 · p. 21
  • 10% improvement: combining duration and visible-workspace filtering yields the greatest success-rate gain, with workspace filtering having a greater impact than duration filtering.Policies trained on unfiltered TAMP data perform worse than policies trained on filtered data.Source: C Ablations · p. 21, Figure C.1 · p. 21
  • Discrete gripper control improves Stack success by 4%; reducing grasp-and-release stall regions from 25 to 5 steps raises success from 78% to 100%.The authors suggest shorter roll-outs ease learning and reduce exposure to unchanged observations and actions that can lead to freezing.Source: C Ablations · p. 21, Figure C.2 · p. 22
  • A suitably angled and positioned camera performs best on Stack; a wrist camera has little impact for blocks but significantly improves Microwave performance.Close-up views of the microwave handle and target object help expose grasp geometry, while other close or distant views impose perception trade-offs.Source: C Ablations · pp. 21–22, Figure C.2 · p. 22
  • 86% for GMM loss versus 66% for MSE and 70% for Gaussian log-likelihood; the authors attribute the advantage to modeling multimodal TAMP actions.The stochastic Gaussian improves performance by 4% over deterministic MSE, while GMM adds a further 16%.Source: C Ablations · p. 22

D Environments

  • The environment suite includes ordered block stacking, object pick-and-place, shelf placement, and microwave opening and object placement, with variants for scene adaptation and randomized receptacle sizes.StackAdapt sometimes requires moving the base block to a target platform; MicrowaveAdapt sometimes requires clearing an object before opening the door.Source: D Environments · pp. 23–24

E Agent Structure

  • Policies use proprioception and selected camera views, with task-space arm control and binary gripper actions; stall-region lengths are tuned to reduce policy freezing while preserving data-collection success.Shelf and Microwave use left and right shoulder views, and every task uses a wrist camera.Source: E Agent Structure · p. 25

F Experiment Details

The experiment details specify training configurations, a tuned BeT baseline, and an evaluation protocol based on held-out demonstration initial states.

F Experiment Details

  • Training uses a ResNet-18 with Spatial Softmax, learned positional embeddings, and AdamW; single-task models use batch size 16 on one V100 GPU, while multi-task training uses batch sizes 512–1024 across eight V100 GPUs.The learning rate is annealed from 10^-4 to 10^-5 over 100K gradient steps; the authors report that larger multi-task batches greatly improved results.Source: F Experiment Details · p. 26, Table F.2 · p. 26, Table F.1 · p. 26
  • The reimplemented BeT baseline uses OPTIMUS's network structure with a discrete-cluster-center and offset head; tuning selected 24 cluster centers and an offset-loss scale of 100.The original BeT hyperparameters did not perform well in this setup.Source: F Experiment Details · p. 26
  • Evaluation resets the simulator to initial states from held-out validation demonstrations, enabling comparisons across seeds and algorithms on the same states.These are states the TAMP solver can solve; the paper reports similar success rates from randomly sampled environment poses.Source: F Experiment Details · p. 27

G.1 Offline Learning from Demonstrations

This section situates the work in offline imitation learning from expert demonstrations. It contrasts human-provided supervision with TAMP-generated datasets for training policies.

G.1 Offline Learning from Demonstrations

  • Offline learning trains a policy on a demonstration dataset without additional interaction, commonly using behavior cloning to imitate demonstrated actions through supervised learning.The paper focuses on behavior cloning and notes that diffusion policies could be integrated into OPTIMUS.Source: §G.1 · p. 28
  • Human teaching and teleoperation provide demonstrations, but teleoperation can require months of collection and numerous operators; this work instead uses TAMP to generate large, diverse datasets without human supervision.The paper also describes intervention-based systems as a way to reduce operator burden and improve sample efficiency.Source: §G.1 · p. 28

G.2 Transformers for Robot Control

Transformer policies have shown strong results in robot manipulation, but discretized observations and actions can limit precise control. OPTIMUS instead uses raw observations and outputs low-level actions.

G.2 Transformers for Robot Control

  • Transformer manipulation policies have demonstrated results across varied tasks, but discretized observations and actions limit their applicability to precise manipulation.Examples include Gato, PerAct, VIMA, RT-1, and Behavior Transformer.Source: §G.2 · p. 28
  • PerAct and VIMA use abstracted actions that ease learning but reduce expressivity and execution speed.Source: §G.2 · p. 28
  • OPTIMUS uses a trainable, scalable Transformer that consumes raw observations and produces fast-to-execute low-level control actions.Unlike HiveFormer, it does not assume temporally extended motion-planner actions or privileged environment geometry for safety.Source: §G.2 · p. 28

G.3 Task and Motion Planning

TAMP combines discrete manipulation choices with continuous motions, but existing approaches can inherit human design bias and computational costs. OPTIMUS focuses on visuomotor imitation without fixed task-level actions and identifies considerations for using TAMP to generate training data.

G.3 Task and Motion Planning

  • TAMP searches combinations of manipulation models to achieve goals, but engineered or learned models can impose human design bias and be too computationally expensive for real-time reaction.TAMP realizes discrete manipulation types through continuous motions.Source: §G.3 · p. 28
  • Unlike Guided TAMP, which focuses primarily on privileged state and predicts fixed ground actions hierarchically, this work imitates TAMP through visuomotor learning without explicit task-level action prediction.A fixed set of ground actions can prevent deployment when the number of objects varies; OPTIMUS does not require grounding and fixing scene objects.Source: §G.3 · pp. 28–29
  • The work identifies new considerations for using TAMP as a data-generation pipeline.Source: §G.3 · p. 29

Related papers

Loading 2305.16309v3…