You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Commit bf71b83
Browse filesBrowse the repository at this point in the historyBrowse files
We propose a semantic compositional diffusion transformer that factorizes transitions into robot-, object-, obstacle-, and objective-specific components and learns their interactions through attention.
133
133
Trained on a limited subset of tasks, our model can zero-shot generate high-quality transitions from which we learn control policies for unseen task combinations.
134
134
We then introduce an iterative self-improvement procedure in which synthetic data is validated via offline reinforcement learning and incorporated into subsequent training rounds.
135
-
Our approach substantially improves zero-shot performance, ultimately solving nearly all held-out tasks and demonstrating the emergence of meaningful compositional structure in the learned representations.
135
+
Our approach substantially improves zero-shot performance, ultimately solving nearly all held-out tasks and demonstrating the emergence of compositional structure in the learned representations.
<pclass="has-text-centered is-size-6 mt-3">Robotic tasks grow <strong>combinatorially</strong>, but collecting demonstrations still scales <strong>linearly</strong>. We learn a <strong>semantic compositional diffusion transformer</strong> that generates expert-style transitions for <em>unseen</em>combinations of robot, object, obstacle, and objective - then <strong>iteratively</strong> validates synthetic data with offline RL and snowballs coverage across the task space - without requiring new data collection efforts.</p>
144
+
<pclass="has-text-centered is-size-6 mt-3">Robotic tasks grow <strong>combinatorially</strong>, but collecting demonstrations still scales <strong>linearly</strong>. Having access to a limited subset of tasks, we learn a <strong>semantic compositional diffusion transformer</strong> that generates expert-style transitions for <em>unseen</em>task combinations - then <strong>iteratively</strong> validates synthetic data with offline RL and snowballs coverage across the task space - without requiring new data collection efforts.</p>
<pclass="mb-0">To ground this idea, we use <strong><ahref="https://arxiv.org/abs/2207.04136" target="_blank">CompoSuite</a></strong> (Mendez et al., 2022), which provides <strong>4×4×4×4 = 256 tasks</strong> by composing one element from each axis: robot, object, obstacle, and objective. Our main experiments study a <strong>14-task training regime</strong> where the robot is fixed to IIWA (the remaining three axes vary), providing access to 4×4×4 = 64 possible task combinations. We evaluate our approach on 32 tasks that are not used during training.</p>
154
+
<pclass="mb-0">To ground this idea, we use <strong><ahref="https://arxiv.org/abs/2207.04136" target="_blank">CompoSuite</a></strong> (Mendez et al., 2022), which provides <strong>4×4×4×4 = 256 tasks</strong> by composing one element from each axis: robot, object, obstacle, and objective. Our main experiments study a <strong>14-task training regime</strong> where the robot is fixed to IIWA (the remaining three axes vary), providing access to 4×4×4 = 64 possible task combinations. We <strong>evaluate</strong> our approach on <strong>32 held-out tasks that are not used during training</strong>.</p>
155
155
</div>
156
156
<divclass="compo-grid">
157
157
<figureclass="has-text-centered">
@@ -200,7 +200,7 @@ <h3 class="title is-4 mt-5">Generating unseen tasks with a semantic, composition
200
200
<divclass="callout-new content is-size-5">
201
201
<ulstyle="margin-left: 1.5rem;">
202
202
<li><strong>Transformers as compositional graphs:</strong> Instead of hard-coding task structure, which requires substantial domain knowledge and may be suboptimal, we propose to <strong>learn the graph structure directly from data</strong>. A transformer's self-attention naturally implements message passing over a graph where each component (robot, object, obstacle, objective, action, reward) is a node. By learning attention weights, the transformer discovers which components influence each other and learns the compositional structure automatically.</li>
203
-
<listyle="margin-top: 0.75rem;"><strong>Semantic, compositional tokenization:</strong> We factorize transitions into <strong>component-specific tokens</strong>, where each element has its own encoder-decoder. Self-attention enables each component to attend to others at every step, implementing graph-compositional inference. Crucially, this ensures synthetic data generation only updates the components involved in each task, preventing cross-task corruption while enabling generalization to unseen combinations.</li>
203
+
<listyle="margin-top: 0.75rem;"><strong>Semantic, compositional tokenization:</strong> We factorize transitions into <strong>component-specific tokens</strong>, where each element has its own encoder-decoder. Self-attention enables each component to attend to others at every step, implementing graph-compositional inference. Crucially, this <strong>ensures synthetic data generation only updates the components involved in each task</strong>, preventing cross-task corruption while enabling generalization to unseen combinations.</li>
<pclass="mb-0">Policies trained on our <strong>synthetic data</strong>outperform <strong>monolithic</strong> and <strong>standard DiT</strong> baselines and, after several refinement rounds, surpass strong <strong>multi-task RL</strong> baselines - without collecting new real trajectories.</p>
234
+
<pclass="mb-0">Policies trained on our synthetic data<strong>outperform</strong><strong>monolithic</strong> and <strong>standard DiT</strong> baselines and, after several refinement rounds, <strong>surpass strong multi-task RL</strong> baselines - without collecting new real trajectories.</p>
<figcaptionclass="has-text-centered">Zero-shot success rate for non-iterable RL models and RL models trained on synthetic data of a generative model at iteration 0, and best zero-shot success rate after 4 iterations of self-improvement. Overall, the semantic compositional architecture leads to large improvements in the RL as well as generative model setting.</figcaption>
240
+
<figcaptionclass="has-text-centered">Zero-shot success rate for non-iterable RL models and RL models trained on synthetic data of a generative model at iteration 0, and best success rate after 4 iterations of self-improvement, evaluated on 32 held-out tasks. Overall, the semantic compositional architecture leads to large improvements in the RL as well as generative model setting.</figcaption>
241
241
</figure>
242
242
243
243
<h3class="title is-4 mt-5">Iterative self-improvement (14 training tasks / 64 possible combinations)</h3>
<figcaptionclass="has-text-centered">Performance of different diffusion architectures over iterations of our self-improvement procedure. Best success rate, raw per-iteration success, solved tasks, and dataset coverage growth over 4 iterations.</figcaption>
246
+
<figcaptionclass="has-text-centered">Performance of different diffusion architectures over iterations of our self-improvement procedure. From left to right, top to bottom: best success rate, per-iteration success rate, best success rate separated by initial task difficulty, percentage of solved tasks, and dataset coverage across 4 iterations.</figcaption>
247
247
</figure>
248
248
</div>
249
249
</section>
@@ -253,7 +253,7 @@ <h3 class="title is-4 mt-5">Iterative self-improvement (14 training tasks / 64 p
<pclass="mb-2">Iterative compositional generation still needs <strong>simulator rollouts</strong> to score candidate policies - but those steps only <strong>evaluate datasets</strong>, they do not imitate data collection or online learning.</p>
256
-
<pclass="mb-0">We compare to <strong><ahref="https://arxiv.org/abs/2302.02948" target="_blank">RLPD</a></strong> (Ball et al., 2023) (offline + online RL) starting from the same 14-task expert pool, allowing up to <strong>100k</strong> on-policy steps per held-out task. Our method reaches much higher success and return with only <strong>~20k</strong> evaluation interactions per task after four refinement rounds - demonstrating that a small amount of structured evaluation can validate massive parallel synthetic rollouts (~1M model-generated transitions per task).</p>
256
+
<pclass="mb-0">We compare to <strong><ahref="https://arxiv.org/abs/2302.02948" target="_blank">RLPD</a></strong> (Ball et al., 2023) (offline + online RL) starting from the same 14-task expert pool, allowing up to <strong>100k</strong> on-policy steps per held-out task. <strong>Our method reaches much higher success and return</strong> with only <strong>~20k</strong> evaluation interactions per task after four refinement rounds - demonstrating that a small amount of structured evaluation can validate massive parallel synthetic rollouts (~1M model-generated transitions per task).</p>
257
257
<pclass="mt-3 mb-0">Note that <strong>our approach is not tied to this particular choice of data evaluation</strong> and <strong>any suitable scoring function that assesses the utility of generated data could be used instead</strong>.</p>
<h2class="title is-3 has-text-centered section-heading">Utility of rare successful synthetic trajectories</h2>
270
270
<divclass="callout-new content is-size-5">
271
-
<pclass="mb-0">While our iterative generation process often produces datasets with <strong>high success rates</strong>, some tasks may yield synthetic datasets with relatively low success rates. We show that even datasets containing only occasional successful trajectories provide valuable signal for downstream reinforcement learning. Using <strong><ahref="https://arxiv.org/abs/2302.02948" target="_blank">RLPD</a></strong> (Ball et al., 2023) initialized with the lowest non-zero success synthetic datasets <strong>(only 8-12% offline RL success rates)</strong>, we reach <strong>80-100% success</strong> within <strong>500,000 environment interactions</strong>, whereas standard <strong>online SAC</strong> fails even after <strong>5 million interactions</strong> - demonstrating that rare successful trajectories significantly accelerate online learning.</p>
271
+
<pclass="mb-0">While our iterative generation process often produces datasets with high success rates, some tasks may yield synthetic datasets with relatively <strong>low success rates</strong>. We show that even datasets containing only occasional successful trajectories provide valuable signal for downstream reinforcement learning. Using <strong><ahref="https://arxiv.org/abs/2302.02948" target="_blank">RLPD</a></strong> (Ball et al., 2023) initialized with the lowest non-zero success synthetic datasets <strong>(only 8-12% offline RL success rates)</strong>, we reach <strong>80-100% success</strong> within <strong>500,000 environment interactions</strong>, whereas standard <strong>online SAC</strong> fails even after <strong>5 million interactions</strong> - demonstrating that rare successful trajectories significantly accelerate online learning.</p>
<h2class="title is-3 has-text-centered section-heading">What structure does the model learn?</h2>
283
283
<divclass="callout-new content is-size-5">
284
-
<pclass="mb-0"><strong>Intervention tests</strong> and <strong>attention maps</strong> show emergent dependencies between task components - different from prior hand-designed compositional stacks and stable across depth and iterative training.</p>
284
+
<pclass="mb-0"><strong>Intervention tests</strong> and <strong>attention maps</strong> show <strong>emergent dependencies between task components</strong> - different from prior hand-designed compositional stacks and stable across depth and iterative training.</p>
285
285
</div>
286
286
287
287
<divclass="columns">
@@ -310,12 +310,12 @@ <h2 class="title is-3 has-text-centered section-heading">What structure does the
310
310
<divclass="container is-max-desktop">
311
311
<h2class="title is-3 has-text-centered section-heading">Scaling to larger compositional spaces</h2>
312
312
<divclass="callout-new content is-size-5">
313
-
<pclass="mb-0">In a <strong>56-task training regime (all four compositional axes varied, 256 possible combinations in total)</strong>, performance scales consistently as the number of axes increases.</p>
313
+
<pclass="mb-0">In a <strong>56-task training regime</strong> (all four compositional axes varied, 256 possible combinations in total), <strong>performance scales consistently</strong> as the number of axes increases.</p>
<figcaptionclass="has-text-centered">Performance in the 56/256-task regime. Semantic compositional model consistently achieves higher success and coverage while maintaining data quality.</figcaption>
318
+
<figcaptionclass="has-text-centered">Performance in the 56/256-task regime. From left to right, top to bottom: best success rate, per-iteration success rate, best success rate separated by initial task difficulty, percentage of solved tasks, and dataset coverage across 4 iterations. Semantic compositional model consistently achieves higher success and coverage while maintaining data quality.</figcaption>
319
319
</figure>
320
320
321
321
<h3class="title is-4 mt-4">When does monolithic generation fall short?</h3>
0 commit comments