PaperlayerView exact revision on arXiv

Source-linked AI summary

RARA: Robust Policy Fine-Tuning through Retrieval-Anchored Representation Alignment for Robotic Manipulation

Xuhui Kang, Sung-Wook Lee, Yan Zhuang, Tianle Zhong, Yen-Ling Kuo

arXiv:2610.13578v1cs.ROcs.CVcs.LG

TL;DR

Limited-data fine-tuning can distort pre-trained robot-policy representations and weaken robustness to out-of-distribution shifts. RARA aligns target representations with relevant retrieved demonstrations, then anchors those demonstrations during fine-tuning. Across 13 tasks, it improves average success over fine-tuning in both in-distribution and out-of-distribution settings.

  • Problem

    Fine-tuning pre-trained policies on limited target data can distort learned representations and weaken generalization under out-of-distribution shifts.

    Source: §1 · p. 1

  • Method

    RARA retrieves task- and motion-relevant pre-training demonstrations, aligns target representations to them, and anchors their representations during target-task fine-tuning.

    Source: §1 · p. 2, §4 · p. 4, §4.1 · p. 5

  • Results

    Across 13 tasks, RARA improves average fine-tuning success by 9.5 % points in-distribution and 13.6 % points out-of-distribution.

    Source: §1 · p. 2

  • Takeaways & Limitations

    RARA keeps target-task representations closer to those of the pre-trained policy while anchoring only retrieved demonstrations.

    Source: §6 · p. 10

    RARA relies on pre-training data containing the same type of task as the target, and the experiments focus on visual change.

    Source: §6 · p. 10

Abstract

from arXiv · show

Pre-trained robot policies, such as Vision-Language-Action (VLA) models, provide a strong foundation for generalist robotic manipulation, but adapting them to new tasks with only a few demonstrations remains challenging. Directly fine-tuning a pre-trained policy often achieves strong in-distribution performance, yet it can distort the representations learned during pre-training and weaken the model's generalization ability under out-of-distribution shifts. We propose Retrieval-Anchored Representation Alignment (RARA), a two-stage adaptation method, to address this issue. First, RARA retrieves representation features from pretraining with similar task intent and motion, and aligns the policy's representations of the target demonstrations with the pre-trained representations of the retrieved demonstrations. Second, RARA fine-tunes the policy on the target demonstrations while limiting how much its representations of the retrieved demonstrations change during fine-tuning. By reducing the representation mismatch before action learning and preserving pre-trained knowledge during adaptation, RARA limits unnecessary representation drift during fine-tuning. Experiments with Diffusion Policy and GR00T on 13 tasks in RoboTwin, RoboCasa, and on a real robot show that RARA improves the average success of fine-tuning by 9.5 points in-distribution and 13.6 points out-of-distribution, and outperforms all compared baselines. Project page: https://shua-kang.github.io/RARA_robot/

1 INTRODUCTION

Fine-tuning limited target demonstrations can distort pre-trained representations and weaken robustness to changed viewpoints, backgrounds, or objects. RARA aligns target representations with relevant pre-training demonstrations, then anchors those demonstrations during fine-tuning; it improves average success across 13 tasks.

1 INTRODUCTION

  • Limited-data fine-tuning can distort pre-trained representations, leaving policies vulnerable to camera, background, and object shifts despite in-distribution success.Preserving generalization under out-of-distribution shifts remains a challenge when adapting pre-trained robot policies.Source: §1 · p. 1
  • RARA retrieves demonstrations relevant to the target task and motion, aligns target representations to their pre-trained representations, and anchors them during fine-tuning.The anchor acts only on retrieved demonstrations, preserving task-related pre-trained structure without pulling the target representation back.Source: §1 · p. 2
  • Across 13 tasks, RARA improves fine-tuned policies’ average success by 9.5 % points in-distribution and 13.6 % points out-of-distribution.The evaluation includes Diffusion Policy on RoboTwin and GR00T on RoboCasa and real-robot tasks.Source: §1 · p. 2
  • The paper evaluates RARA across policies, simulated and real-world manipulation tasks, and studies alignment, representation constraints, source-data choice, and network layers.These studies separate the effects of initial alignment and constraints during fine-tuning.Source: §1 · p. 2

2 RELATED WORK

Prior robot-policy adaptation methods preserve pre-trained capabilities through weights, architecture, or target-data representations, while retrieval methods select relevant prior examples. RARA instead anchors retrieved pre-training demonstrations in representation space.

2 RELATED WORK

  • Prior approaches preserve capabilities through weight constraints or merging, frozen components, auxiliary knowledge, or alignment to a teacher on target data (Huang et al., 2026; Yadav et al., 2026; Kachaev et al., 2026; Shang et al., 2024).These methods do not focus alignment on pre-training demonstrations most relevant to the target task.Source: §2 · pp. 2–3
  • RARA anchors representations of retrieved pre-training demonstrations, rather than restricting architecture or which policy weights can be updated.It aligns target representations with those demonstrations and limits their representation change during fine-tuning.Source: §2 · p. 3
  • Prior data-selection methods retrieve examples by state-action or motion similarity, while RARA uses selected demonstrations as representation anchors rather than additional training data (Du et al., 2023; Lin et al., 2024; Dass et al., 2026).Other domain-adaptation lines include sim-real data mixing, distribution alignment, and representation regularization.Source: §2 · p. 3

3 PROBLEM DEFINITION

The problem is to adapt a pre-trained visuomotor policy from limited target demonstrations while maintaining success under both familiar and shifted evaluation conditions. Standard target-only fine-tuning uses behavior cloning, whereas the robust objective covers both distributions.

3 PROBLEM DEFINITION

  • The policy is pre-trained on diverse robot demonstrations and adapted to a downstream task using a limited set of observation-action demonstrations.Observations may include visual inputs and language or state information.Source: §3 · p. 3
  • Standard target-only fine-tuning updates the policy by behavior cloning on the target demonstrations.This objective directly optimizes performance on in-distribution target data.Source: §3 · p. 3
  • Robust fine-tuning aims for high expected task success under both in-distribution and out-of-distribution evaluation distributions.Success is defined over trajectories induced by the policy in each evaluation setting.Source: §3 · p. 3

4 RARA: RETRIEVAL-ANCHORED REPRESENTATION ALIGNMENT

RARA adapts a pre-trained policy in two stages: it aligns target representations with retrieved pre-training representations, then fine-tunes on target demonstrations while constraining changes to retrieved representations. Retrieval uses task relevance and motion similarity to establish anchors and frame pairs.

4 RARA: RETRIEVAL-ANCHORED REPRESENTATION ALIGNMENT

  • RARA first aligns target representations to pre-trained representations of retrieved demonstrations, then fine-tunes on the target task while anchoring the retrieved representations.Stage 1 trains the representation-producing layers and a projector; Stage 2 uses the policy’s behavior-cloning loss while keeping the anchor active, then discards the projector for inference.Source: §4 · pp. 3–4

4.1 RETRIEVING RELEVANT PRE-TRAINING DEMONSTRATIONS

  • Retrieval is task-informed: instruction-based selection returns demonstrations of the same kind of task, while action similarity alone can return chunks from a different task.For the target Open Blender Lid task, the figure contrasts instruction retrieval with action-similarity selection.Source: Figure 3 · p. 4
  • RARA pairs target frames with retrieved frames by action-chunk similarity measured with DTW, whose distance matches similar actions performed at different speeds.The closest retrieved chunks supply positives, while chunks above the median distance supply negatives; retrieved actions form pairs, not behavior-cloning targets.Source: §4.1 · p. 5

4.2 ALIGNING AN INTERMEDIATE REPRESENTATION

  • The alignment compares projected target embeddings from the adapting policy with projected retrieved-frame embeddings from a frozen pre-trained policy.The representation may come from any policy layer; the projector is used for contrastive learning and is discarded after Stage 1.Source: §4.2 · p. 5
  • The anchor loss compares adapting and pre-trained representations on the same retrieved frames, limiting their change without directly constraining target representations.It uses no action labels and remains active through both stages, leaving target representations free to change under alignment and task losses.Source: §4.2 · p. 6

5 EXPERIMENTS

The experiments test whether RARA preserves pretrained representations and improves ID/OOD success, using two policies across simulated and real-robot tasks and studies of layer choice and components.

5.1 SETUP

  • The study uses Diffusion Policy on four RoboTwin tasks and DROID-pretrained GR00T N1.7 on five RoboCasa and four real-robot tasks, evaluating ID and specified OOD shifts.Each task has a small target demonstration set, and experiments are repeated three times.Source: §5.1 · p. 6, §5 · p. 6

5.3 RQ2: DOES RARA IMPROVE ID AND OOD SUCCESS?

  • RARA has the highest task-average ID/OOD success in all three settings, gaining 8.6/11.0 points on RoboTwin, 4.8/11.2 on RoboCasa, and 16.4/19.2 on the real robot over FT.The gains are larger in OOD than ID in each setting.Source: Figure 5 · p. 7, Table 2 · p. 8, §5.3 · pp. 8–9

5.2 RQ1: HOW DOES RARA CHANGE THE REPRESENTATION?

  • RARA limits target-frame distortion to 0.23–0.33 versus 0.47–0.63 for FT and Co-FT, while retaining pretraining-frame CKA 0.91–0.92 and probe R2 0.70–0.71.These representation results are reported for two RoboTwin tasks in ID and OOD conditions.Source: §5.2 · pp. 7–8, Table 1 · p. 8
  • Stage 1 aligns target frames with retrieved frames but leaves distortion at 0.46–0.56; the Stage-2 anchor then reduces it to 0.23–0.33, and both stages are needed for the full gain.Source: §5.2 · p. 8

5.5 RQ4: WHICH COMPONENTS AND WHICH DATA PRODUCE THE GAIN?

  • Both stages are needed for the full gain; Co-FT on retrieved demonstrations lowers success in all four columns, while random pretraining data or random frame pairing can erase gains.Random Stage-1 pairing falls below FT on RoboTwin (65.3%/40.6% vs. 67.4%/47.3%), and random pretraining data underperforms FT by up to 17.1 points there.Source: §5.2 · p. 8, §5.5 · pp. 9–10, Table 3 · p. 9

5.4 RQ3: DOES THE CHOICE OF ANCHORED LAYER MATTER?

  • Every tested Diffusion Policy layer improves Move Pill Bottle success over FT, from 21.2% to 32.0%–39.7%; four of five GR00T locations improve by 4–10 points.Only an early GR00T vision-encoder layer remains at FT level.Source: Figure 6 · p. 9, §5.4 · p. 9

6 CONCLUSION

RARA retrieves same-type pre-training demonstrations, aligns target representations to them, and anchors those representations during fine-tuning. Across 13 simulation and real-robot tasks, it improves average success, while relying on related pre-training data and visual-shift settings with the same robot and action space.

  • Across 13 simulation and real-robot tasks, RARA improves fine-tuning success by 9.5 points in ID and 13.6 points in OOD.The method also keeps target-task frame representations closer to those of the pre-trained policy, although anchoring is applied only to retrieved demonstrations.Source: §6 · p. 10
  • RARA retrieves same-type pre-training demonstrations, aligns target representations to their pre-trained representations, then anchors the retrieved demonstrations during fine-tuning.Source: §6 · p. 10
  • RARA relies on pre-training data containing the same type of task and evaluates visual changes with the same robot and action space.Future work includes retrieving references when such data are scarce and extending to new embodiments and action spaces.Source: §6 · p. 10

A ANALYSIS OF VISUAL FEATURE SHIFT

The analysis identifies collapsed visual features for unseen RoboCasa tasks and motivates aligning target features to task-relevant pre-trained regions. Retrieval combines instruction similarity with action-chunk similarity to select and pair relevant demonstrations.

A ANALYSIS OF VISUAL FEATURE SHIFT

  • Unseen RoboCasa task features collapse into a compact PCA region despite representing different tasks, unlike well-separated DROID features.The visualization uses features from a DROID-pretrained π0 policy and five tasks from each domain.Source: A ANALYSIS OF VISUAL FEATURE SHIFT · p. 14, Figure 7 · p. 14
  • Moving target visual features toward a task-relevant pre-trained region is intended to make new observations easier for the pre-trained policy to interpret before action learning.Feature collapse can leave the action model with weak state information and make it prone to overfitting limited demonstrations.Source: A ANALYSIS OF VISUAL FEATURE SHIFT · p. 14

A.1 PCA VISUALIZATION DETAILS

  • For PCA, 196 patch tokens per camera are averaged into 1,152-dimensional vectors, concatenated across three views, then projected to two components.The PCA projection is fitted once on combined feature sets when comparing sources or checkpoints.Source: §A.1 · p. 14

B RETRIEVAL AND PAIR CONSTRUCTION DETAILS

  • RARA retrieves the top 100 demonstrations by cosine similarity between instruction embeddings and the average normalized embedding of target-task instructions.Instruction filtering separates different task intents even when their action chunks have similar end-effector displacements.Source: B RETRIEVAL AND PAIR CONSTRUCTION… · pp. 14–15
  • Overlapping action chunks are compared with DTW, which matches actions executed at different speeds; closest chunks are positives and farther retrieved chunks supply negatives.Pairs are drawn only from retrieved demonstrations, avoiding unrelated tasks that happen to have similar actions.Source: B RETRIEVAL AND PAIR CONSTRUCTION… · p. 15

C IMPLEMENTATION DETAILS

RARA anchors visual representations, with the representation and loss scaling adapted to the policy architecture. The implementation uses fixed retrieval and training budgets across methods.

  • RARA anchors pooled ResNet-18 camera features for Diffusion Policy and vision-tower output tokens for GR00T.Stage 1 trains the representation-producing layers and projector for 3,000 steps while keeping the rest of the policy fixed.Source: C IMPLEMENTATION DETAILS · p. 15
  • The anchor loss uses mean squared error for Diffusion Policy and batch-variance-normalized squared error for GR00T representations.Normalization allows the same weight to be used for representations with different scales.Source: C IMPLEMENTATION DETAILS · p. 15
  • All settings retrieve 100 demonstrations and use 16-step action chunks, with the same optimizer settings and training budget for all methods in a setting.Source: C IMPLEMENTATION DETAILS · p. 15

D EVALUATION ENVIRONMENTS

The appendix specifies evaluation protocols and visual-shift conditions for RoboCasa, real-robot, and RoboTwin tasks, alongside representation metrics and anchored-layer studies. It also details RoboTwin data, policy, and adaptation settings.

D.1 ROBOCASA AND REAL ROBOT

  • RoboCasa uses 50 episodes per task and three training seeds; real Franka policies use 10 trials per task, repeated three times with matched randomization settings when possible.Real-world GR00T executes the first 8 steps of each predicted 16-step chunk before replanning; RoboCasa executes all 16.Source: §D.1 · p. 16
  • RoboCasa and real-robot OOD evaluations vary camera views, backgrounds, distractors, objects, clutter, or layouts.The figure captions identify the specific visual variations for each environment.Source: Figure 9 · p. 16, §D.1 · p. 16, Figure 10 · p. 17

D.2 ROBOTWIN

  • RoboTwin targets test novel objects for pick-and-place skills and new objects or motion variants for pressing and shaking, amid widened layout, background, and distractor shifts.The pre-training data include unrelated laptop-opening and QR-code-rotation skills intended to be excluded by retrieval.Source: §D.2 · p. 17
  • RoboTwin uses a ResNet-18 Diffusion Policy trained for 120 pre-training epochs and 300 adaptation epochs; Stage 2 mixes target chunks with retrieved frames for behavior cloning and anchoring.Baselines share the pre-trained policy, 50 target demonstrations, and 300-epoch budget; evaluation averages the last three checkpoints across three seeds.Source: §D.2 · pp. 17–18

D.3 RQ1 METRICS

  • Representation metrics use held-out target frames from ID and OOD regions and pre-grasp pre-training frames; distortion measures feature movement, while CKA measures agreement up to rotation and scale.The appendix cautions that CKA can change without a functional change, so it is not used alone.Source: §D.3 · pp. 18–19
  • Position-probe R2 predicts grasp-time object position from features, excludes post-grasp frames, and treats a drop relative to the pre-trained representation as forgetting.Source: §D.3 · p. 19

E CHOICE OF ANCHORED LAYER

  • The layer study changes only the anchored location, comparing pooled visual features with intermediate or action-network representations across Diffusion Policy and GR00T.For GR00T, the fine-tuned parameters match the FT reference; the LLM is frozen in Stage 2 under the official recipe.Source: E CHOICE OF ANCHORED LAYER · p. 19
  • Layer-study points average the last three checkpoints, with GR00T checkpoints evaluated on 100 ID and 100 OOD rollouts.Source: E CHOICE OF ANCHORED LAYER · p. 19

F PROJECTOR ANALYSIS

The analysis finds that projector choice changes representation alignment, but the reported success rates are similar across projector variants. RARA therefore uses a trainable projector by default.

F PROJECTOR ANALYSIS

  • The default shared two-layer MLP projector maps representations from 1024 to 512 to 256 dimensions, is trained with the policy, and is discarded after Stage 1.Removing it leaves Stage 1 followed by unanchored fine-tuning with similar success rates on the reported tasks; the trainable projector is used by default.Source: F PROJECTOR ANALYSIS · p. 20
  • With either projector, full RARA reaches 49.3% vs. 53.3% on Move Pill Bottle and 82.0% vs. 84.7% on Place Container (ID; frozen vs. trainable).Stage 1 followed by fine-tuning reaches 32.0% ID on Move Pill Bottle and 85.3% vs. 87.3% on Place Container (frozen vs. trainable).Source: F PROJECTOR ANALYSIS · p. 20
  • A frozen projector improves held-out Move Pill Bottle frame-ranking AUC by 8.1 points over pre-training, versus 5.1 points with the trainable projector, but moves representations further from pre-training (CKA 0.86 vs. 0.95).The frozen-projector contrastive loss remains two to three times higher; the trainable projector absorbs part of the alignment.Source: F PROJECTOR ANALYSIS · p. 20

G CHOICE OF PRE-TRAINING DATA

Changing only the pre-training data shows that RARA’s gains depend on how well the source data’s skill matches the target. This motivates retrieving relevant demonstrations rather than using all pre-training data.

G CHOICE OF PRE-TRAINING DATA

H RETAIN UNDER A LARGE PRE-TRAINING–TARGET GAP

RETAIN’s performance relative to full fine-tuning differs across the two simulated settings: it retains more success on RoboTwin than RoboCasa. The authors relate this difference to the size of the pre-training–target gap and contrast weight interpolation with RARA’s representation constraint.

H RETAIN UNDER A LARGE PRE-TRAINING–TARGET GAP

  • RETAIN interpolates pre-trained and fine-tuned weights, and its effectiveness depends on how far the target task is from the pre-training data.Source: H RETAIN UNDER A LARGE… · p. 20
  • On RoboTwin, RETAIN averages 47.0% vs. 67.4% ID and 34.2% vs. 47.3% OOD success, retaining about 70% of FT’s success.On Move Pill Bottle, RETAIN reaches 13.3% OOD success vs. 9.6% for FT.Source: H RETAIN UNDER A LARGE… · pp. 20–21, Table 6 · p. 21
  • On RoboCasa, RETAIN averages 31.1% vs. 69.2% ID and 14.8% vs. 44.3% OOD success, and falls below FT on all five tasks in both conditions.Its OOD success reaches 3.3% on Pick-Place to Drawer.Source: Table 7 · p. 21, H RETAIN UNDER A LARGE… · p. 21
  • The authors suggest the larger gap between real-world DROID pre-training and simulated kitchens may require larger weight changes, making interpolation preserve neither solution’s behavior.RARA constrains retrieved-demonstration representations rather than weights, allowing policy changes as needed for the target task.Source: H RETAIN UNDER A LARGE… · p. 21

Related papers

Loading 2610.13578v1…