RAG-Diff: Adapting Diffusion Policies to Dynamic Constraints with Retrieval-Augmented Guidance

Robots operating in unstructured environments must satisfy dynamic constraints that can change across tasks and even within a single execution. While diffusion policies can learn multimodal behaviors from demonstrations, adapting a trained policy at runtime to newly encountered or evolving constraints remains an open challenge.

We propose RAG-Diff (Retrieval-Augmented Guided Diffusion), a runtime adaptation framework for a frozen transformer diffusion policy that leverages retrieval-augmented memory. RAG-Diff maintains PrefMem, a memory bank of vision-language embeddings paired with (i) state-action snippets and (ii) constraint annotations. At test time, RAG-Diff retrieves the nearest entry and uses it to steer sampling in two complementary ways: I-Atten (In-place Attention recomputation) injects the retrieved snippet as additional cross-attention memory tokens with a classifier-free-guidance-style update, while a predictive value-guidance mechanism uses the retrieved constraint parameters to discourage violations during sampling.

We demonstrate RAG-Diff in physical robot caregiving, a domain with personalized, time-varying constraints — benchmarking on an adapted PushT environment and a suite of caregiving tasks (bed bathing, medicine delivery, shelf cleaning, feeding) spanning interaction, affordance, spatial, and semantic preferences, in both RCareWorld simulation and on a real robot, plus a real-world user study. RAG-Diff improves both task success and constraint satisfaction over unguided diffusion and other guidance- or sampling-based baselines.

We first train a transformer-based diffusion policy on offline demonstrations, following Diffusion Policy. We then steer the frozen policy at test time using PrefMem, a preference memory pool initialized from the same demonstrations. At each timestep, a frozen VLM encodes the recent observation history and retrieves the nearest PrefMem entry — a state-action snippet plus its constraint annotation — which steers denoising in two complementary ways:

• I-Atten (In-place Attention recomputation): injects the retrieved snippet as extra cross-attention memory tokens and mixes it with the base prediction via a classifier-free-guidance-style update — no retraining required.
• Value guidance: recovers the predicted clean action chunk at each denoising step and perturbs the noise prediction with gradients of a differentiable constraint value function.

"Old ladies have sensitive skin. I appreciate the robot following my preferences and being gentle with me."
"The robot really feels like it's listening to me and responding to what I need."

RAG-Diff grounds natural-language preferences (e.g., "Bring me the orange juice box") via text-to-image retrieval in the VLM embedding space — a capability none of the baselines support. It succeeds in all trials (5/5) with distinctive descriptions, though it can degrade under visual or linguistic ambiguity.

Performance depends on retrieval quality — misleading retrieval (visually similar but semantically different contexts) can yield wrong preference labels or noisy action snippets — and is sensitive to guidance hyperparameters. Retrieval assumes visually similar contexts imply similar actions and preferences, which can fail for latent, non-visual factors. Constraints are enforced as soft guidance, not hard safety guarantees; real-world deployment should add hard safety guardrails.