🤖 machine learning

Iterative Compositional Data Generation for Robot Control

This paper proposes an iterative compositional data generation framework using a semantic compositional diffusion transformer to synthesize high-quality robotic transitions for unseen task combinations, which are further refined through offline reinforcement learning to achieve near-perfect zero-shot generalization.

Original authors: Anh-Quan Pham, Marcel Hussing, Shubhankar P. Patankar, Dani S. Bassett, Jorge Mendez-Mendez, Eric Eaton

Published 2026-04-15

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Anh-Quan Pham, Marcel Hussing, Shubhankar P. Patankar, Dani S. Bassett, Jorge Mendez-Mendez, Eric Eaton

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to do chores. The problem is that the world is full of endless combinations: a robot arm picking up a box, a different arm pushing a ball, a third arm avoiding a wall while stacking plates. If you tried to film a human expert demonstrating every single one of these millions of combinations, you'd never finish. It would take a lifetime and cost a fortune.

This paper proposes a clever solution: Don't film every chore. Teach the robot the "ingredients" of the chores, then let it imagine the rest.

Here is the breakdown of their method, using some everyday analogies.

1. The Problem: The "Infinite Menu"

Think of robotic tasks like a restaurant menu.

The Chef (Robot): Can be a human, a mechanical arm, or a drone.
The Ingredients (Objects): A box, a dumbbell, a plate.
The Obstacles: A wall, a door, or nothing.
The Goal: Pick it up, push it, or put it in a bin.

If you have 4 types of chefs, 4 types of ingredients, 4 types of obstacles, and 4 types of goals, you have $4 \times 4 \times 4 \times 4 = 256$ different "dishes" (tasks). In the real world, this number explodes into the millions. Collecting real-world data (filming experts) for every single dish is impossible.

2. The Solution: The "Master Chef" AI

The authors built a special AI called a Semantic Compositional Diffusion Transformer. Let's break down that scary name:

Diffusion: Think of this as a "denoising" process. Imagine you have a clear photo of a robot doing a task, and you slowly add static (noise) until it's just gray fuzz. The AI learns to reverse this: it starts with gray fuzz and slowly removes the noise to reconstruct the clear photo. In this case, it's reconstructing the movement of the robot.
Compositional: Instead of treating the whole task as one giant, unbreakable blob, the AI breaks it down into its "ingredients" (Robot, Object, Obstacle, Goal). It learns how a "Robot" interacts with an "Object" separately from how it interacts with an "Obstacle."
Transformer: This is the brain that connects the dots. It's like a super-smart translator that understands how the "Robot" token talks to the "Goal" token.

The Analogy:
Imagine you are learning to cook.

Old Way: You watch a video of someone making a "Spaghetti Carbonara." Then you watch another video for "Spaghetti Bolognese." You have to re-learn everything from scratch for every new dish.
This Paper's Way: You learn the rules of cooking. You learn how to boil pasta (Skill), how to chop onions (Object), and how to use a stove (Robot). Once you understand these ingredients, you can instantly imagine how to cook a dish you've never seen before, like "Spaghetti with a weird new sauce," without needing a video tutorial.

3. The Secret Sauce: "Iterative Self-Improvement"

This is the most exciting part. The AI doesn't just generate data once and stop. It plays a game of "Show and Tell" with itself.

Step 1: The Guess. The AI looks at the few tasks it knows (e.g., Robot A with Box B) and tries to imagine what Robot A would do with Box C (a task it has never seen). It generates a fake video of this new task.
Step 2: The Test. It takes this fake video and tries to train a robot controller on it. Then, it puts that controller in a simulation to see if it actually works.
Step 3: The Filter.
- If the controller fails miserably, the AI says, "That fake video was garbage," and throws it away.
- If the controller succeeds, the AI says, "Great! That fake video was actually good data." It adds this new video to its library.
Step 4: The Loop. Now the AI has more data (the real ones + the good fake ones). It retrains itself on this bigger library, gets smarter, and tries to imagine even more difficult tasks.

It's like a student who takes a practice test. If they get a question right, they study that type of question again to master it. If they get it wrong, they realize they need to study the basics. Over time, they get better at taking tests they've never seen before.

4. The Results: From "Zero" to "Hero"

The researchers tested this on a benchmark with 256 different robotic tasks. They only gave the AI real data for 14 of them.

The Result: The AI used its "imagination" (generative model) to create high-quality fake data for the other 242 tasks.
The Outcome: It learned to solve nearly all of the unseen tasks, outperforming other methods that tried to memorize tasks or use rigid, pre-programmed rules.

Why This Matters

Saves Time and Money: We don't need to hire humans to film robots doing every possible task in the world.
Generalization: It proves that if you teach a robot the structure of the world (how objects relate to each other), it can figure out new situations on its own.
Safety: Since the "imagination" happens in a computer simulation, the robot doesn't break anything while learning.

In a nutshell: This paper teaches robots to be creative. Instead of just memorizing a list of instructions, the robot learns the "grammar" of movement. Once it knows the grammar, it can write its own sentences (solve new tasks) without needing a teacher to show it every single time.

1. Problem Statement

Collecting robotic manipulation data is prohibitively expensive, particularly in multi-object, multi-robot, and multi-environment settings where the task space is combinatorially large. While recent generative models can synthesize data for individual tasks, they fail to exploit the inherent compositional structure of robotic domains (e.g., tasks formed by combining specific robots, objects, obstacles, and goals). Consequently, these models struggle to generalize to unseen task combinations, requiring new data collection for every novel scenario. The paper addresses the challenge of how a robot learning system can iteratively improve its ability to solve unseen tasks by generating high-quality artificial training data for those tasks without requiring additional real-world demonstrations.

2. Methodology

The authors propose a framework centered on a Semantic Compositional Diffusion Transformer (DiT) combined with an Iterative Self-Improvement loop.

A. Semantic Compositional Diffusion Transformer

Instead of using a monolithic architecture or hard-coded compositional graphs, the model learns the task structure directly from data.

Factorization: Transitions are factorized into specific components: robot, object, obstacle, objective (subtask), action, reward, and termination signal.
Architecture: The model uses a Diffusion Transformer where each factor (e.g., a specific robot arm or object type) has its own dedicated encoder-decoder pair.
Graph Interpretation: The transformer is interpreted as a Graph Neural Network (GNN). Self-attention mechanisms allow the model to learn the underlying graph structure connecting these factors (e.g., how a specific robot interacts with a specific object) without pre-defining the edges.
Tokenization: Unlike standard DiTs that use shared encoders, this approach uses factor-specific tokenizers. This allows the model to learn distinct representation spaces for different elements (e.g., the "Jaco" robot encoder is separate from the "Panda" robot encoder) while sharing the transformer blocks that model their interactions.

B. Iterative Self-Improvement Procedure (Algorithm 1)

The system operates in a closed loop to bootstrap data generation:

Training: A DiT is trained on a limited subset of real expert datasets (e.g., 14 tasks).
Generation: The DiT generates synthetic transitions for unseen task combinations (zero-shot).
Validation: Offline Reinforcement Learning (RL) policies (TD3-BC) are trained on the generated data and evaluated online in the environment.
Filtering & Update: If a generated dataset yields a policy with a success rate above a threshold $\tau$ , the synthetic data is added to the training pool.
Refinement: The DiT is retrained on the expanded dataset (real + validated synthetic). If no tasks meet the threshold, the patience counter increases; eventually, the threshold $\tau$ is lowered to admit more data, ensuring the system continues to learn even from sparse successes.

3. Key Contributions

Learned Compositional Structure: The paper demonstrates that a transformer can automatically discover the optimal compositional graph structure from data, outperforming hard-coded architectures (like Mendez et al., 2022) which rely on pre-defined decompositions.
Factor-Specific Tokenization: The authors show that using separate encoders for each task element (robot, object, etc.) is crucial for zero-shot generalization, outperforming models that share a single encoder across different elements of the same category.
Iterative Data Bootstrapping: The method introduces a self-improving loop where synthetic data is validated via offline RL and fed back into the generator. This allows the model to solve nearly all held-out tasks without collecting new real-world data.
Emergent Structure: Analysis of the learned attention maps reveals that the model discovers a dependency structure (e.g., the robot arm is the most critical factor influencing other components) that differs from previous hard-coded assumptions, proving that effective compositional priors can emerge automatically.

4. Experimental Results

Experiments were conducted on the CompoSuite benchmark, a simulated environment with 256 tasks formed by combining 4 robots, 4 objects, 4 obstacles, and 4 objectives. The study focused on a low-data regime (14 training tasks out of 64 possible with a fixed robot).

Zero-Shot Performance: The Semantic Compositional DiT achieved a success rate nearly double that of hard-coded compositional RL baselines and significantly outperformed monolithic diffusion models (SynthER) and standard DiTs.
Iterative Improvement: After 4 rounds of self-improvement, the proposed method solved nearly 100% of the held-out tasks. In contrast, monolithic baselines showed much slower improvement and lower final coverage.
Data Efficiency: The method achieved high success rates using only 20,000 environment interactions (for evaluation) to validate 1 million synthetic transitions. This is vastly more efficient than online RL baselines (like RLPD) which required 100,000 interactions and still failed to solve the tasks effectively.
Rare Success Utility: Even synthetic datasets with very low success rates (8–12%) provided sufficient signal to bootstrap online RL, allowing agents to reach 80–100% success rates rapidly, whereas training from scratch failed completely.
Scalability: The approach maintained its advantages when scaled to a larger regime (56 training tasks out of 256), confirming that the benefits of compositional learning persist as task complexity increases.

5. Significance

This work addresses a fundamental bottleneck in robotics: the inability to scale learning to the vast combinatorial space of possible tasks. By treating the generative model as a learner of task structure rather than just a pattern matcher, the authors enable zero-shot synthesis of high-quality training data for entirely new task combinations.

The significance lies in:

Reducing Data Collection Costs: It drastically reduces the need for expensive human demonstrations or teleoperation for every new task variation.
Generalization: It proves that learning the structure of the domain (the graph of interactions) is more effective than learning specific task policies, enabling robust generalization to unseen environments.
Self-Improving Systems: It establishes a viable path for autonomous systems to curate their own training data, moving from static datasets to dynamic, self-refining learning loops.

The paper concludes that while currently state-based, this framework offers a scalable path toward sim-to-real transfer and could be extended to visual observations by incorporating structured perception modules.