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

01 Motivation

扩散策略（diffusion policy）能从示教中学到多模态的操作行为，但一旦训练完成，它学到的是一个固定的行为分布。而现实中的机器人——尤其是物理护理场景——必须应对随时间演化的约束：从擦拭划痕转向擦拭易碎玻璃时要调整接触力，或者要在人的手臂突然伸进工作空间时立刻改道避让。作者指出，一个称职的机器人必须(i) 在线适应这些约束，且 (ii) 用最少的额外训练、试错或人工干预来做到这一点。Diffusion policies can learn multimodal manipulation behaviors from demonstrations, but once training is finished what they have learned is a fixed behavior distribution. Real-world robots — especially in physical caregiving — must cope with constraints that evolve over time: adjusting contact force when moving from wiping a scratch to wiping fragile glass, or immediately rerouting when a person's arm suddenly reaches into the workspace. The authors argue that a competent robot must (i) adapt to such constraints online, and (ii) do so with minimal additional training, trial and error, or human intervention.

论文比较了四类既有的策略适配思路——RL 微调、classifier-free (CF) 引导、sample-and-rank (SR)、classifier-guided (CG)——并指出：微调太慢、代价高，不适合安全关键场景；CF 引导无法跟随训练时未见过的约束；SR 无法超越基础策略的表征上限，且重复采样-评分效率低；CG 需要训练额外的分类器/奖励预测器，在约束隐式、难以密集标注，以及约束反复出现导致重复计算的场景下都会遇到瓶颈。The paper compares four existing families of policy adaptation — RL fine-tuning, classifier-free (CF) guidance, sample-and-rank (SR), and classifier-guided (CG) — and argues: fine-tuning is too slow and too costly for safety-critical settings; CF guidance cannot follow constraints unseen during training; SR cannot exceed the representational ceiling of the base policy, and its repeated sample-and-score loop is inefficient; CG requires training an extra classifier / reward predictor, and runs into bottlenecks when constraints are implicit and hard to annotate densely, and when recurring constraints force redundant computation.

“We propose RAG-Diff … a runtime adaptation framework for a frozen transformer diffusion policy that leverages retrieval-augmented memory.” —— RAG-Diff 冻结基础策略参数 θ，仅在推理时通过检索增强的记忆库转向满足动态、个体化的用户偏好，目标是“在最小化偏好违反的同时完成任务，且不更新基础策略参数”。“We propose RAG-Diff … a runtime adaptation framework for a frozen transformer diffusion policy that leverages retrieval-augmented memory.” — RAG-Diff freezes the base policy parameters θ and steers it, only at inference time, through a retrieval-augmented memory so as to satisfy dynamic, individualized user preferences; the stated goal is to “complete the task while minimizing preference violations, without updating the base policy parameters”.

02 Method

RAG-Diff 分两阶段：先用离线示教训练一个 transformer-based 的 base action diffusion policy（沿用 Diffusion Policy 的范式，以最近观测历史为条件、对动作块 action chunk 去噪）；再在测试时用两种互补的方式引导这个冻结策略——(1) I-Atten 针对隐式偏好（只能通过示例传达、没有显式参数），(2) value guidance 针对显式约束（如力阈值、禁区、目标位置）。两者共用同一个检索记忆库 PrefMem。RAG-Diff has two stages: first, offline demonstrations train a transformer-based base action diffusion policy (following the Diffusion Policy paradigm, denoising an action chunk conditioned on the recent observation history); then, at test time, two complementary mechanisms steer this frozen policy — (1) I-Atten addresses implicit preferences (conveyable only through examples, with no explicit parameters), and (2) value guidance addresses explicit constraints (such as force thresholds, regions to avoid, goal positions). Both share the same retrieval memory bank, PrefMem.

I-Atten：In-place Attention recomputationI-Atten: In-place Attention recomputation

从偏好记忆池 B 中检索到一个相关的 state-action 片段后，用与基础策略相同的编码器（ResNet 视觉特征 + 低维状态嵌入）把它编码成提示 token P，与原始上下文 C 拼接得到 C⁺ = [C; P]。denoiser 在同一个噪声动作块上跑两次，分别以 C 和 C⁺ 为交叉注意力上下文，得到 ε_base 与 ε_prompt，再以 classifier-free-guidance 风格混合：ε_attn = ε_base + w_attn(k)(ε_prompt − ε_base)。这一机制无需重新训练，只在推理时“原地”把检索到的行为模式作为额外记忆 token 重新计算注意力。After a relevant state-action snippet is retrieved from the preference memory pool B, the same encoder as the base policy (ResNet visual features + low-dimensional state embedding) encodes it into prompt tokens P, which are concatenated with the original context C to give C⁺ = [C; P]. The denoiser runs twice on the same noisy action chunk, using C and C⁺ respectively as cross-attention context, yielding ε_base and ε_prompt, and then blends them in classifier-free-guidance style: ε_attn = ε_base + w_attn(k)(ε_prompt − ε_base). The mechanism requires no retraining; it merely recomputes attention “in place” at inference time, treating the retrieved behavior pattern as additional memory tokens.

Value guidance：把约束变成可微值函数Value guidance: turning constraints into differentiable value functions

I-Atten 提供隐式引导，但不显式强制约束。为处理显式偏好，RAG-Diff 在采样当前去噪步 k 时，用 Tweedie 公式从 Aᵏ_t 和 ε_base 恢复出预测的“干净”动作块 Â⁰_t，定义一个可微值函数 J(Â⁰t) 度量沿动作轨迹的约束违反程度，再通过其梯度扰动噪声预测（ε_value = ε_base + w_value(k)Δε）来引导采样。论文考虑了四类约束，并用加权和 J = λ_F J_force + λ_A J_afford + λ_S J_goal 组合：I-Atten supplies implicit guidance but does not explicitly enforce constraints. To handle explicit preferences, at the current denoising step k RAG-Diff uses Tweedie's formula to recover the predicted “clean” action chunk Â⁰_t from Aᵏ_t and ε_base, defines a differentiable value function J(Â⁰t) measuring the degree of constraint violation along the action trajectory, and perturbs the noise prediction with its gradient (ε_value = ε_base + w_value(k)Δε) to guide sampling. The paper considers four constraint types, combined as a weighted sum J = λ_F J_force + λ_A J_afford + λ_S J_goal:

• Interaction (force) constraint：用一个可微 soft-hinge 值函数惩罚预测接触力超过用户舒适阈值 F_th 的部分（力由一个 MLP 力预测器 g_φ 估计）。Interaction (force) constraint: a differentiable soft-hinge value function penalizes the portion of the predicted contact force that exceeds the user's comfort threshold F_th (force is estimated by an MLP force predictor g_φ).
• Affordance (region-to-avoid) constraint：用可微 signed-distance function d_afford 编码禁区（如伤口、皮疹区域），惩罚轨迹进入或过于靠近该区域。Affordance (region-to-avoid) constraint: a differentiable signed-distance function d_afford encodes the region to avoid (such as a wound or a rash), penalizing trajectories that enter or come too close to that region.
• Spatial (goal position) constraint：把空间偏好建模为目标末端执行器位置 g，用到目标的距离作为平滑的可微引导信号。Spatial (goal position) constraint: spatial preference is modeled as a target end-effector position g, with the distance to that goal used as a smooth, differentiable guidance signal.
• Semantic (language-grounded) constraint：语言请求（如 "bring the orange juice"）不直接指定几何信息，用 VLM 从记忆库中检索相关片段，再以 value guidance 落地执行。Semantic (language-grounded) constraint: a language request (e.g. "bring the orange juice") does not specify geometry directly, so a VLM retrieves relevant snippets from the memory bank and value guidance grounds them into execution.

PrefMem：偏好记忆与检索PrefMem: preference memory and retrieval

I-Atten 与 value guidance 都依赖同一个共享记忆池 PrefMem，它初始化自训练基础策略用的同一批离线示教。每条记忆条目 b_i = (q_i, I_i, s_i, a_i, p_i) 存储一段观测窗口的 RGB 帧、低维状态、动作片段，以及关联的偏好标注（力阈值 / 禁区 / 语言偏好等）。查询时，用同一个 frozen VLM 对当前观测历史做逐帧编码并做平均池化，归一化后作为 query，按余弦相似度做 top-N 检索。当用户给出新偏好时，RAG-Diff 用余弦相似度阈值 δ 判断这是对已有记忆条目的更新，还是应作为分布外（out-of-distribution）新条目扩展进记忆库。Both I-Atten and value guidance rely on the same shared memory pool, PrefMem, initialized from the very same offline demonstrations used to train the base policy. Each memory entry b_i = (q_i, I_i, s_i, a_i, p_i) stores the RGB frames of an observation window, low-dimensional states, an action snippet, and the associated preference annotation (force threshold / region to avoid / language preference, etc.). At query time, the same frozen VLM encodes the current observation history frame by frame and mean-pools it; the normalized result serves as the query for top-N retrieval by cosine similarity. When the user states a new preference, RAG-Diff uses a cosine-similarity threshold δ to decide whether it is an update to an existing memory entry or should extend the memory bank as a new out-of-distribution entry.

03 Experiments

论文在仿真（RCareWorld 上的四个物理护理任务 + 一个改造版 Push-T 玩具环境）与真实机器人（Kinova Gen3 + Realsense D435，人体模型 manikin）上做了系统评测，覆盖 interaction / affordance / spatial 三类约束的五个任务：Push-T（仅仿真）、Bed Bathing、Medicine Handover、Serving、Shelf Cleaning。对比基线为 DP（无引导的 Diffusion Policy）、CG-DP（classifier-guided DP，强梯度引导基线）、SR（sample-and-rank，事后选择而非引导）；消融为 I-Atten DP（仅 I-Atten）与 V-DP（仅 value guidance）。The paper reports a systematic evaluation in simulation (four physical caregiving tasks in RCareWorld plus a modified Push-T toy environment) and on a real robot (Kinova Gen3 + Realsense D435, with a manikin), covering five tasks across the three constraint types interaction / affordance / spatial: Push-T (simulation only), Bed Bathing, Medicine Handover, Serving and Shelf Cleaning. The baselines are DP (unguided Diffusion Policy), CG-DP (classifier-guided DP, a strong gradient-guidance baseline) and SR (sample-and-rank, post-hoc selection rather than guidance); the ablations are I-Atten DP (I-Atten only) and V-DP (value guidance only).

完整数据见论文 TABLE II（含 Push-T Afford、Serving 等全部列）。注意 CG-DP 在 Bed Bathing 的 Afford↓ 指标上以 0.14±0.06 略优于 RAG-Diff 的 0.17±0.04——论文原话："CG-DP consistently improves constraint satisfaction but reduces task success."，这是一处如实呈现的“基线在单项指标上更优”的例子。Full data are in TABLE II of the paper (including all columns such as Push-T Afford and Serving). Note that on the Bed Bathing Afford↓ metric, CG-DP's 0.14±0.06 is slightly better than RAG-Diff's 0.17±0.04 — in the paper's words, "CG-DP consistently improves constraint satisfaction but reduces task success." This is an honestly reported instance of “a baseline winning on a single metric”.

量化结果：contact comfort 上 RAG-Diff 为 4.42±0.66，CG-DP 为 3.09±1.33；perceived improvement in comfort over time 上 RAG-Diff 为 4.38±0.85，CG-DP 为 2.68±1.60；preference adaptation 上 RAG-Diff 为 4.29±1.19，CG-DP 为 1.82±1.42。全部试验中参与者在 81.81% 的情形下更偏好 RAG-Diff；在偏好于交互中途发生变化的试验里，100% 的参与者更偏好 RAG-Diff。一位老年参与者反馈："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 and responding to what I need."Quantitative results: on contact comfort RAG-Diff scores 4.42±0.66 against 3.09±1.33 for CG-DP; on perceived improvement in comfort over time, 4.38±0.85 against 2.68±1.60; on preference adaptation, 4.29±1.19 against 1.82±1.42. Across all trials, participants preferred RAG-Diff in 81.81% of cases; in the trials where the preference changed mid-interaction, 100% of participants preferred RAG-Diff. One older participant commented, "Old ladies have sensitive skin. I appreciate the robot following my preferences and being gentle with me." Another said, "The robot really feels like it's listening and responding to what I need."

Ablations

论文指出："combining retrieved-context injection with value/constraint guidance yields complementary benefits."（结合两种引导方式收益互补）。仅用 I-Atten（I-Atten DP）通常能相对无引导策略提升成功率，但"does not reliably reduce violations, suggesting that attention-based prompting primarily selects favorable trajectory modes without explicitly optimizing constraint satisfaction"（更多是选中了更优的轨迹模式，而非显式优化约束满足）；仅用 value guidance（V-DP）则"offers a more balanced trade-off, improving constraint satisfaction with less impact on success than CG-DP, but remains worse than the full method."The paper notes: "combining retrieved-context injection with value/constraint guidance yields complementary benefits." (the two guidance mechanisms are complementary). I-Atten alone (I-Atten DP) generally raises the success rate relative to the unguided policy, but "does not reliably reduce violations, suggesting that attention-based prompting primarily selects favorable trajectory modes without explicitly optimizing constraint satisfaction" (it mostly picks better trajectory modes rather than explicitly optimizing constraint satisfaction); value guidance alone (V-DP) "offers a more balanced trade-off, improving constraint satisfaction with less impact on success than CG-DP, but remains worse than the full method."

语言引导（Sec. V-E）Language guidance (Sec. V-E)

在真实世界的取饮料任务中（桌上放可乐罐与橙汁盒，给出如 "Bring me the orange juice box." 的语言指令），RAG-Diff 借助 VLM 嵌入空间的文本-图像检索，在全部 5 次试验中都成功抓取正确物体（5/5），而所有基线都不支持语言引导。论文也坦承其局限："it degrades under visual similarity or ambiguous language (e.g., 'robot arm with a can') yields the same similarity to both orange juice box and coke can images due to the VLM limitation."In a real-world drink-fetching task (a coke can and an orange juice box on the table, with a language instruction such as "Bring me the orange juice box."), RAG-Diff uses text-image retrieval in the VLM embedding space and grasps the correct object in all 5 trials (5/5), while none of the baselines support language guidance. The paper is candid about the limitation: "it degrades under visual similarity or ambiguous language (e.g., 'robot arm with a can') yields the same similarity to both orange juice box and coke can images due to the VLM limitation."

04 Limitations