Source-linked AI summary
CACTI: A Framework for Scalable Multi-Task Multi-Scene Visual Imitation Learning
Zhao Mandi, Homanga Bharadhwaj, Vincent Moens, Shuran Song, Aravind Rajeswaran, Vikash Kumar
TL;DR
Robot learning lacks the large, diverse datasets and efficient training methods that have advanced computer vision and language processing. CACTI combines expert data collection, generative augmentation, pretrained visual representations, and imitation policy training. Its single policies perform across multiple real-robot and simulated kitchen tasks and scene variations.
Problem
Robot demonstrations are labor-intensive to collect and difficult to diversify across the objects and scenes needed for manipulation.
Source: §1 · p. 1Method
CACTI collects task-specific expert data, augments it for visual diversity, compresses images with pretrained representations, and trains one multi-task imitation policy.
Source: §1 · pp. 1–2Results
A single policy achieved ≈30% success across 10 real-robot tasks and ≈62% across 18 simulated tasks and 100 layout variations; end-to-end RL achieved 0% success.
Source: §3.4 · p. 6Takeaways & Limitations
CACTI provides a multi-task, multi-scene visual imitation framework and open-source multi-layout simulation benchmarks.
Source: §1 · p. 2Frozen visual backbones may compress away information in raw images that the policy needs.
Source: §3.3.1 · p. 6
Abstract
from arXiv · showhide
Large-scale training have propelled significant progress in various sub-fields of AI such as computer vision and natural language processing. However, building robot learning systems at a comparable scale remains challenging. To develop robots that can perform a wide range of skills and adapt to new scenarios, efficient methods for collecting vast and diverse amounts of data on physical robot systems are required, as well as the capability to train high-capacity policies using such datasets. In this work, we propose a framework for scaling robot learning, with specific focus on multi-task and multi-scene manipulation in kitchen environments, both in simulation and in the real world. Our proposed framework, CACTI, comprises four stages that separately handle: data collection, data augmentation, visual representation learning, and imitation policy training, to enable scalability in robot learning . We make use of state-of-the-art generative models as part of the data augmentation stage, and use pre-trained out-of-domain visual representations to improve training efficiency. Experimental results demonstrate the effectiveness of our approach. On a real robot setup, CACTI enables efficient training of a single policy that can perform 10 manipulation tasks involving kitchen objects, and is robust to varying layouts of distractors. In a simulated kitchen environment, CACTI trains a single policy to perform 18 semantic tasks across 100 layout variations for each individual task. We will release the simulation task benchmark and augmented datasets in both real and simulated environments to facilitate future research.
1. Introduction
CACTI addresses the cost and difficulty of collecting diverse robot demonstrations by splitting multi-task visual imitation learning into four stages. It combines limited expert data, generative augmentation, pretrained visual representations, and training a single policy across tasks and scenes.
- Robot demonstrations are labor-intensive, time-consuming, and difficult to diversify across the objects and scenes needed for manipulation.Source: §1 · p. 1
- CACTI separates robot learning into Collect, Augment, Compress, and TraIn stages, from task-specific expert data to one general multi-task agent.The stages respectively gather demonstrations, increase visual diversity, project images into a low-dimensional latent space, and train the policy.Source: §1 · pp. 1–2
- The framework uses pretrained out-of-domain visual representations and trains a single policy on frozen embeddings across multiple tasks.The paper also presents multi-layout, multi-task simulation benchmarks for release to the research community.Source: §1 · p. 2
- Generative augmentation expands limited in-domain data with semantic and visual variations, reducing the need for extensive physical data collection.In the real-robot setup, Stable Diffusion in-paints realistic object variations; experts curate the original data to maintain its quality.Source: §1 · p. 2
2. A Framework for Multi-Task Multi-Scene Visual Imitation Learning
CACTI is instantiated in kitchen environments using limited expert demonstrations, scene augmentation, compressed visual representations, and a single policy trained across tasks and scene variations. Its collection and augmentation procedures differ between real-robot and simulated settings.
2. A Framework for Multi-Task Multi-Scene Visual Imitation Learning
- CACTI separates data collection, visual augmentation, image compression, and multi-task policy training to address data diversity, training cost, and safety concerns.The framework is instantiated in real and simulated kitchen environments using standard techniques.Source: §2 · pp. 2–3
2.1. Collect: Small in-domain expert data collection
- Real-robot demonstrations use kinesthetic teaching, with five trajectories per task replayed 20 times while non-target objects are rearranged.Simulation experts are trained per task and layout; the study reports 900 policies filtered at a 90% success threshold.Source: §2.1.1 · p. 3, §2.1.2 · pp. 3–4, §2.1 · p. 3
2.2. Augment: Semantic scene variations for augmentation
- Augmentation varies visual attributes and object layouts, including changes to positions, orientations, and added objects.Source: §2.2 · p. 4
2.2.1. REAL-ROBOT ENVIRONMENT
- Real-robot augmentation uses masked, text-prompted Stable Diffusion in-painting to add plausible objects while leaving the rest of the scene unchanged.The process generates varied demonstrations without additional human data collection.Source: Figure 2 · p. 3, §2.2.1 · p. 4
2.2.2. SIMULATED KITCHEN
- In simulation, expert action trajectories are replayed with visual attribute changes, randomized distractor poses, and action noise.The procedure renders 50 episodes of 50 time-steps for each task-layout combination, producing 45,000 training episodes.Source: §2.2.2 · p. 4
2.3. Compress: Representations pre-trained on internet or in-lab data
- Compression maps image observations to low-dimensional embeddings, using R3M for real-robot data and MoCo-trained ResNet-50 for simulated in-domain data.The real setup’s limited data motivates fine-tuning R3M, while simulated data is used to train an in-domain baseline.Source: §2.3 · pp. 4–5
2.4. TraIn: Multi-Task Multi-Scene Visual Policy Learning
- The final stage trains one policy on the multi-task, multi-scene dataset using the stage-three visual backbone and behavior-cloning loss.For real-robot task conditioning, task-name token embeddings are concatenated with observation embeddings; augmented trajectory endpoints can serve as goal images.Source: §2.4 · p. 5
3. Experiments
The experiments assess task and scene performance, generalization to novel layouts, cost efficiency, and the effects of choices within CACTI’s stages across real-robot and simulated kitchens.
3.1. Real Robot Experiment Setup
- The real-robot setup uses 10 tabletop manipulation tasks, and evaluation varies distractor locations and target-object starting positions.The tasks use a Franka Emika Panda arm and household toy objects.Source: §3.1 · p. 5
3.2. Simulation Experiment Setup
- The simulated kitchen has 18 semantic tasks and 100 training layouts; evaluation adds 10 layouts per task and resamples visual properties.The released CACTI-Sim-100 benchmark includes 100 semantic variations per base task and 50 trajectories per variation.Source: §3.2 · pp. 5–6
3.3.1. VISUAL REPRESENTATIONS (COMPRESS)
- The representation ablation compares fine-tuned visual backbones with frozen backbones pretrained on in-domain or out-of-domain images.The motivation is that frozen out-of-domain features improve storage and computational efficiency, but may discard information present in raw images.Source: §3.3.1 · p. 6
3.3.2. IN-PAINTING AUGMENTATION (AUGMENT)
- The in-painting ablation compares semantic-augmented training data with standard crops and color jitter on the original data, using the same amount of expert-action information.The ablation varies image-level data while holding the amount of expert action trajectories constant.Source: §3.3.2 · p. 6
3.4. Results
- Semantic in-painting augmentation yields about 15–20% higher absolute success than standard color-jitter and random-crop augmentation variants.The comparison uses real-robot results; the no-augmentation variants use standard image augmentation.Source: Figure 5 · p. 6, §3.4 · p. 7
- CACTI achieves approximately 30% success across 10 real-robot tasks and approximately 62% across 18 simulated tasks and 100 variations; pixel-input end-to-end RL reaches 0% success in simulation.The authors describe the policy results as promising while noting that task performance remains far from perfect.Source: §3.4 · p. 6, §3.3.3 · p. 6, Figure 6 · p. 7
- Held-out-layout generalization increases from Sim-10 to Sim-100, while real-robot evaluation shuffles distractor layouts and includes novel distractor objects.The real-robot setup also varies target-object starting positions.Source: §3.1 · p. 5, §3.4 · p. 6, Table 1 · p. 7
- Frozen out-of-domain representations perform almost as well as in-domain frozen and fine-tuned representations on the real robot; the simulation gap is larger.The authors suggest the larger simulation difference may reflect visual differences between simulation and the real-world pretraining data.Source: §3.4 · p. 7, Figure 6 · p. 7, Table 1 · p. 7
4. Related Work
CACTI relates to prior approaches to scaling robot learning, visual policy learning, domain randomization, generative models, and representation learning. It combines offline imitation with semantic visual augmentation and pretrained representations.
4.1. Frameworks for scaling robot learning
- CACTI combines offline data, pretrained visual representations, and augmentation rather than collecting all demonstrations from experts, which is costly in robotics.Prior work found shared task learning can amortize behavior acquisition costs; CACTI draws on these findings while multiplying a small expert dataset.Source: §4.1 · pp. 7–8
4.2. Visual policy learning.
- Unlike prior work on simplified visual tasks or kitchen policies using simulator state, CACTI-Sim uses visual inputs in a realistic kitchen with more layout and visual diversity.The authors note visual inputs make it more amenable to real-world settings where ground-truth object states are unavailable.Source: §4.2 · p. 8
4.3. Domain randomization.
- CACTI extends domain randomization beyond color, camera, and texture changes to semantic distractor and layout variations through hindsight relabeling of limited demonstrations.Source: §4.3 · p. 8
4.4. Pretrained models for robot learning
- Generative-model in-painting adds realistic scene objects, and the authors report zero-shot generation on robotic data from strictly unseen lab settings [Ramesh et al., 2022].They contrast this with DALL-E image-goal generation for object rearrangement [Ramesh et al., 2022].Source: §4.4 · p. 8
4.5. Representation learning for control.
- CACTI uses pretrained visual-representation frameworks with both out-of-domain internet human-video data and in-domain augmented data for compressed representations.Source: §4.5 · p. 8
5. Discussion and Conclusion
CACTI is instantiated for multi-task, multi-scene visual imitation learning in simulation and on a real robot. The authors report that large out-of-domain generative and representational models can be leveraged to boost robot-learning generalization, and identify multimodal and multi-stage policies as future directions.
5. Discussion and Conclusion
- The framework combines efficient data collection, novel augmentation, compressed visual representations, and one control policy trained on augmented datasets.Source: §5 · p. 8
- Experiments instantiate the framework in a varied simulated kitchen and a real-robot environment with everyday kitchen objects, including generative-model in-painting augmentations.Source: §5 · p. 8
- The authors conclude that large out-of-domain generative and representational models can boost robot-learning generalization; future work includes multimodal data and multi-stage policies.They also propose investigating deeper connections between large out-of-domain models and robot learning.Source: §5 · p. 8
A.1.1. STAGE 1
The collection phase uses layout-specific simulation experts and kinesthetic demonstrations on the real robot.
A.1.1. STAGE 1
- Simulation collection produced 100 expert policies per task across layouts, totaling 1,800 task- and layout-specific policies.Source: §A.1.1 · p. 12
- Real-robot collection used 8 kinesthetic trajectories per task, totaling 80 expert trajectories for replay.Source: §A.1.1 · p. 12
A.1.2. STAGE 2
Real-robot augmentation replays demonstrations with visual and semantic scene changes, including controlled generative in-painting.
A.1.2. STAGE 2
- Replayed kinesthetic demonstrations vary scene attributes and record images at each timestep; visual augmentation adds color jitter, while action noise broadens coverage.Source: §A.1.2 · p. 12
- Manual semantic variation shuffles distractor positions and swaps some objects into or out of the scene.Source: §A.1.2 · p. 12
- Stable Diffusion takes a scene image and pixel-coordinate region to modify, generating plausible objects while keeping the rest unchanged and requiring no additional robot-operation hours.The authors describe this as dataset-level controlled generation that can produce visually augmented demonstrations without extra human data-collection effort.Source: §A.1.2 · p. 12
A.1.3. STAGE 3
The visual representation uses R3M pretrained on egocentric human videos, with its encoder kept frozen or fine-tuned during imitation learning.
- R3M’s pretrained visual representations were learned from egocentric human videos using time-contrastive and video-language alignment losses.Source: §A.1.3 · p. 12
- The authors use the original R3M pretrained model without adding a loss for fine-tuning on their collected data.Source: §A.1.3 · p. 12
- Fine-tuning updates the pretrained encoder through backpropagation during stage 4 imitation learning.Source: §A.1.3 · p. 12
A.1.4. STAGE 3
The section specifies the policy inputs and describes task and evaluation setups for real-robot and simulated kitchens.
A.1.4. STAGE 3
- The policy receives concatenated 1024-dimensional visual-goal and observation embeddings plus 8-dimensional joint-velocity and joint-pose vectors, and outputs an 8-dimensional Gaussian action distribution.Source: §A.1.4 · p. 12
A.2. Experiment Details on Real robot environment
- Real-robot tasks move tabletop objects to goal locations; Figure 7 shows 10 tasks, with each row sampling three frames from a 100-frame episode.Source: §A.2 · p. 12, Figure 7 · p. 14
- Real-robot evaluation uses 100-step episodes and a task-name embedding; generalization tests vary distractor and target-object positions and add unseen distractors.Source: §A.2 · pp. 12–13
A.3. Experiment Details on Simulation environment
- Simulation combines 18 semantic tasks with 100 layouts into 1,800 combinations, using 43-dimensional embeddings for target-object pose and layout-specific object arrangement.Source: §A.3 · p. 13
- Simulation evaluation tests 10 additional layouts per task with randomized color, lighting, and texture; the 100-layout policy gets five episodes per combination, versus ten for 10- and 50-layout policies.Source: §A.3 · p. 13