Exploring Conditions for Diffusion models in Robotic Control

AI summary

Overview

• Research area: Computer Vision / robot learning — specifically visual representation learning for imitation-based robotic control using pre-trained text-to-image diffusion models (Stable Diffusion v1.5).
• Technical level: Intermediate (some familiarity with diffusion models, behavior cloning, and vision-language concepts helps, but the core argument is intuitive).
• Scope: The paper investigates how to condition a frozen text-to-image diffusion model so that the visual features it produces become task-adaptive for robotic control, and proposes a method called ORCA that learns task prompts and visual prompts rather than relying on text captions.

Authors: Heeseong Shin, Byeongho Heo, Dongyoon Han, Seungryong Kim, Taekyung Kim. Affiliations listed: NAVER AI Lab and KAIST AI. Project page: https://orca-rc.github.io/

What This Paper Is About

Pre-trained visual encoders used in imitation learning are typically kept frozen, which makes the visual representation they provide the same for every downstream control task — that is, task-agnostic. The authors ask whether a pre-trained text-to-image diffusion model can instead provide a task-adaptive representation by conditioning it, without fine-tuning the diffusion model itself. They find that the obvious answer — feeding in text descriptions — largely fails in control environments, and propose a replacement based on learnable task prompts plus frame-level visual prompts.

Key Contributions

1. A negative result with an explanation. The authors show that text conditions generated by a state-of-the-art vision-language model (Gemini 2.5) yield minimal or even negative gains in robotic control, unlike in other vision tasks. They trace this to a domain gap between the diffusion model's web-image training data and simulated control environments, evidenced by noisy cross-attention maps in failing tasks.
2. A formulation of what conditions in control require. They argue that effective conditions for control should be task-grounded (attached to the specific environment) and frame-sensitive (varying across a dynamic video stream), because control tasks involve fine-grained movement of specific agent parts rather than simply categorizing objects.
3. ORCA, a framework that conditions a frozen Stable Diffusion v1.5 using learnable task prompts (shared across observations, capturing task-relevant objects or areas) and visual prompts (dense features from a vision encoder, projected through a small convolutional layer, capturing frame-level detail). Both are trained end-to-end with a standard behavior cloning loss during downstream policy learning.
4. State-of-the-art results across DeepMind Control, MetaWorld, Adroit, and LIBERO-Long, with ablations covering prompt components, U-Net layer selection, vision encoder choice, and diffusion backbone choice.

Main Findings

• Text conditions are ineffective or harmful. Comparing the null condition (an empty string containing only <eos> and <bos> tokens) against text conditions produced mixed results: text helped on some tasks (e.g., Button-press, Reacher) but degraded others (e.g., Cheetah-run).
• Cross-attention maps explain the failure. For Button-press, words such as press and button were well-grounded to relevant image regions. For Cheetah-run, words such as cheetah and run produced noisy attention maps, and the <eos> token of the null condition was already roughly grounded to the salient agent — which is how a sub-optimal text condition can be worse than no condition at all.
• ORCA is the top performer across all 12 tasks in the three main suites. On DeepMind Control, ORCA reaches a mean normalized score of 74.3, versus 70.7 for TADP, 70.6 for CoOp, 69.1 for Text (Simple), 68.9 for Text (Caption), 68.3 for SCR, 59.5 for CLIP, and 53.1 for VC-1. On MetaWorld, ORCA reaches a mean success rate of 95.2%, versus 93.1% (TADP), 92.5% (CoOp), 92.5% (Text Caption), 91.5% (Text Simple), 90.4% (SCR), 84.2% (VC-1), and 81.3% (CLIP). On Adroit, ORCA reaches a mean of 65.3%, versus 58.0% (SCR and CoOp), 57.3% (Text Simple, Text Caption, TADP), 51.4% (CLIP), and 47.3% (VC-1).
• No task-agnostic representation wins everywhere. Among task-agnostic baselines, SCR performed best overall, but VC-1 and CLIP outperformed it on certain tasks — which the authors use to illustrate the core limitation of frozen, task-agnostic features.
• Visual information helps. The task-adaptive baselines CoOp (learnable prefix tokens on a fixed task name) and TADP (a special token S* optimized via Textual Inversion to encapsulate visual style) generally beat plain text conditions, supporting the hypothesis that visual information is beneficial — though their gains on DeepMind Control and Adroit were limited because they were not designed for control.
• Full fine-tuning is harmful. On Adroit, fine-tuning VC-1 with 302.3M learnable parameters dropped its mean from 47.3% to 31.3%; fine-tuning SCR with 346.7M parameters collapsed its mean from 58.0% to 9.3%. The paper states success rates collapse by over 80% compared to the frozen counterpart and that full fine-tuning uses roughly 30× more parameters. Parameter-efficient variants did better: VC-1 + RoboAdapter (18.0M params) reached 59.3%, and SCR + LoRA (4.6M params) reached 60.0% — while ORCA reached 65.3% with 10.6M learnable parameters.
• Both prompt types are needed. On DeepMind Control, using neither prompt gave a mean of 68.3, task prompt alone 69.8, visual prompt alone 70.5, and both 74.3. The paper notes that individually each prompt can behave divergently across tasks (e.g., Reacher-easy benefiting more from visual prompts), but together they give consistent gains.
• Earlier U-Net layers are better. Layer-wise means on DeepMind Control: down_1 68.4, down_2 65.1, down_3 67.5, mid 64.6, up_0 59.0, up_1 60.4, up_2 49.7. Concatenating the best-performing blocks (denoted d+m) gave 74.3.
• U-Net based diffusion beats the alternatives tested. On Adroit, Stable Video Diffusion (SVD) reached a mean of 26.0% and Stable Diffusion 3.0 (SD 3.0, a multi-modal diffusion Transformer) reached 30.0%, versus 65.3% for ORCA with SD 1.5. The authors attribute SVD's weakness partly to restricting input to 3 frames instead of its native 8-frame setting and to the released version not supporting text conditioning; they attribute SD 3.0's weakness to limited empirical study of DiT-based models relative to U-Net architectures.
• It transfers to long-horizon tasks. On three LIBERO-Long tasks, ORCA achieved a single-task mean of 65.5 versus 57.2 (SCR), 43.3 (VC-1), and 36.7 (SigLIP). In multi-task training with a shared policy, ORCA reached 46.6 versus 26.7 (VC-1), 23.3 (SCR), and 16.7 (SigLIP).
• The vision encoder choice matters but is not the whole story. For visual prompts on Adroit, no vision encoder gave a mean of 54.7, SigLIP 55.3, CLIP 63.3, and SD-VAE 61.3; ORCA's own configuration uses pre-trained DINOv2 per the implementation details (the corresponding table row is truncated in the provided text).

Methodology in Plain English

The setup starts from an existing recipe in computer vision: take a pre-trained text-to-image diffusion model, feed it an image plus a text prompt, and pull intermediate features out of its denoising network to use as a visual representation for a downstream task. The paper applies this recipe to robot control, where a policy network must map visual observations to actions, and asks what the prompt should be.

Instead of writing captions, ORCA replaces the text prompt with two kinds of learnable parameters:

• Task prompts — a small set of learned tokens (length 4) shared across every observation in a task. They are trained with the policy so they implicitly latch onto whatever regions of the scene matter for that task, such as the button and the robot arm in Button-press.
• Visual prompts — dense features extracted from a separate vision encoder (pre-trained DINOv2) and projected through a small convolutional layer (length 16 tokens). Because these are dense rather than global, they carry per-frame spatial detail, so they can distinguish, for example, the front and back legs of an agent.

An observation is encoded by the VQGAN encoder of Stable Diffusion v1.5, noised at timestep t = 0, and passed through the U-Net. The visual representation is taken from the downsampling blocks and the bottleneck block, concatenated and pushed through a compression layer, and passed to the policy network. Both the prompt parameters and the policy network are trained end-to-end with a behavior cloning loss that regresses the demonstrated actions — the diffusion model itself is never fine-tuned. Because only the prompts change per task, prompting modules can be swapped between tasks while the diffusion backbone stays fixed.

Training details: 100 epochs per task with online evaluation every 10 epochs; 5 demonstrations for DeepMind Control, 5 for MetaWorld, and 2 for Adroit; proprioceptive data used except for DeepMind Control; a compression layer applied to all baselines for fair comparison.

Why This Matters

This work challenges a common assumption that techniques which work for perception tasks (text-conditioned diffusion features) transfer straightforwardly to control tasks, and it offers a concrete, cheap alternative: adapt through learned prompts instead of fine-tuning a large generative model per task. For the research community, it reframes "what should the condition be?" as a design question specific to the dynamics and fine-grained spatial demands of control, and it provides evidence that full fine-tuning on small imitation datasets overfits badly.

Real-world applications suggested by the evaluation suites:

• Robotic manipulation with a Sawyer-style arm — assembly, bin-picking, button pressing, drawer opening, and hammering.
• Dexterous manipulation with a multi-fingered hand — the Adroit tasks use a 28-DoF anthropomorphic hand for relocating and reorienting objects.
• Long-horizon, language-conditioned household tasks — the LIBERO-Long kitchen, living room, and study setups.
• Locomotion and balance control — standing, walking, reaching, running, and finger spinning from DeepMind Control.

Industry relevance: since the diffusion backbone is untouched, a single frozen model can serve many tasks by swapping lightweight prompting modules (ORCA uses 10.6M learnable parameters versus 302.3M for full fine-tuning of VC-1 and 346.7M for SCR), which matters for storage, deployment, and multi-task robot fleets.

Future Directions

• Better diffusion backbones for control. SD 3.0, a diffusion Transformer, and Stable Video Diffusion both underperformed the U-Net-based SD 1.5 in the paper's tests; the authors explicitly leave deeper exploration of video diffusion and DiT-based models as future work.
• Exploiting video models properly. SVD had to be restricted to 3 input frames instead of its native 8 and lacks publicly released text conditioning, so the question of whether a video-native backbone helps control remains open.
• Reducing the need for hand-designed conditions further. The paper notes that frame-by-frame text descriptions would be impractical to produce; further automation of task-grounded conditioning is a natural extension.
• Scaling and generalization. Multi-task LIBERO-Long results suggest joint learning of prompts across tasks is possible, but the provided text does not report evaluation on real physical robots, nor does it report scaling to the full task sets of these benchmarks.

Target Audience

Researchers and engineers working on robot learning, imitation learning, or pre-trained visual representations, particularly those already familiar with diffusion models and looking for efficient alternatives to fine-tuning them. It is also useful for practitioners evaluating whether prompting-based adaptation from vision transfers to control domains, and for readers interested in a careful case study of when a widely reported technique does not carry over.

Note on Unreported Material

The provided content lists additional appendix sections — results with different diffusion backbones, LIBERO-Long results, vision encoder ablation, an ablation on timesteps, comparison with stronger pre-trained encoders, an efficiency comparison, an analysis on the null condition, further discussion, limitations, and qualitative results — but the body text for the timestep ablation, the stronger-encoder comparison, the efficiency comparison, the null-condition analysis, the limitations section, and the qualitative results section is not included above. Values from those sections are therefore not reported here.

Authors’ abstract

While pre-trained visual representations have significantly advanced imitation learning, they are often task-agnostic as they remain frozen during policy learning. In this work, we explore leveraging pre-trained text-to-image diffusion models to obtain task-adaptive visual representations for robotic control, without fine-tuning the model itself. However, we find that naively applying textual conditions - a successful strategy in other vision domains - yields minimal or even negative gains in control tasks. We attribute this to the domain gap between the diffusion model's training data and robotic control environments, leading us to argue for conditions that consider the specific, dynamic visual information required for control. To this end, we propose ORCA, which introduces learnable task prompts that adapt to the control environment and visual prompts that capture fine-grained, frame-specific details. Through facilitating task-adaptive representations with our newly devised conditions, our approach achieves state-of-the-art performance on various robotic control benchmarks, significantly surpassing prior methods.