Robota-Embed: Joint Action and Visual Encoding for Robot In-Context Learning

To reuse past operational experience, embodied AI robots typically retrieve historical demonstration videos and imitate previously successful actions.

However, traditional retrieval methods often rely solely on visual appearance: if a table or object looks similar, the system assumes a match, ignoring the robot’s current motion direction. For instance, sealing a box with tape and tearing it off produce nearly identical visual frames for the same cardboard box and tape roll, yet the subsequent actions are completely opposite. Relying only on visual similarity can easily lead robots to select incorrect reference trajectories.

On October 14, a joint team from NVIDIA, Massachusetts Institute of Technology, University of California, Berkeley, University of Oxford, and Carnegie Mellon University published the paper "Robota-Embed: Unified VLA Embeddings for Robot In-Context Learning and Evolution." The paper features 19 authors, with Yunze Liu and Jing Wang as co-first authors, and NVIDIA Research Scientist Han Cai serving as corresponding author.

The author list includes several prominent scholars: Yuke Zhu, co-lead of NVIDIA’s GEAR team and a robotics learning expert; Trevor Darrell, founding co-director of Berkeley’s Artificial Intelligence Research (BAIR) lab; Jim Fan, co-lead of NVIDIA’s GEAR team; and Song Han, Associate Professor at MIT, known for his research on model compression and efficient AI computing.

Robota-Embed jointly encodes visual, textual, and action information into a unified vector space, incorporating action dimensions for both query samples and candidate samples. When retrieving historical experience, the robot finds references that match the logic of subsequent behaviors, rather than just visually similar clips.

Enabling robots to access past experience is an emerging technical route in robot learning. Robot policies no longer rely solely on knowledge hardcoded within model weights; instead, they retrieve historical demonstrations and replay trajectories as context during inference to complete decisions based on existing experience. Studies such as GEN-1.5 and S1 have achieved on-site learning based on sensor-action samples and video, while strategies like GPT-6 Astra can generate actions based on observation and execution history. Once experience is incorporated into the context, selecting the most effective clips becomes a critical challenge in robot control.

The core logic of traditional experience retrieval is visual similarity. While general embedding models align images with text, and action-aware metric learning associates perceptual information and text with behavior, most retrieval schemes only compare scene appearance without including the action itself in the similarity calculation. This creates a misalignment: the robot’s current state describes the "now," while the candidate recording captures the behavior that will occur after that interaction. For example, taping and untaping a box look nearly identical visually but involve opposite subsequent actions; relying solely on scene matching often retrieves the wrong clip. Solutions such as ROSER, RAEA, and REGENT either retrieve action sequences or use indices without action data to fetch actions, failing to incorporate action information on both the query and candidate sides.

In contrast, Robota-Embed achieves bidirectional embedding of actions on both ends of the retrieval process. The query side includes task instructions, cropped interaction-region images, and current actions; the candidate side consists of the same instruction paired with a strictly post-current-moment 4-second video and aligned actions. Queries and candidates each output normalized vectors, which are sorted using cosine similarity. To accommodate actions from different data sources, the research team equipped each source with an independent action projector and established three training objectives responsible for alignment prediction, multimodal localization discrimination, and representation regularization. The team also constructed a multi-source dataset and a retrieval benchmark, Robota-Embed-Bench. This unified vector interface supports two downstream tasks: supplementing context for policies and continuously accumulating task memory.

01 Equipping Experience Retrieval with Actions

Robota-Embed addresses the challenge of robot experience retrieval: given a robot’s current state, it matches subsequent action segments from a historical trajectory library, rather than merely finding visually similar frames. The team decomposes both queries and candidate samples into three types of information: instructions, visual frames, and actions. The query combines task instructions, cropped images of interaction regions, and current native actions; candidates use the same instructions paired with 4-second videos and aligned actions recorded at times later than the current moment. Information from both ends is encoded into normalized vectors, with retrieval ranking performed via cosine similarity.

The query and future candidate branches share a backbone network, video tokenizer, and projection head, each attached to independent adapters; actions from different data sources are mapped into the backbone space through dedicated projectors. This design ensures that action information directly participates in similarity calculations, rather than being output as an additional result only after retrieval. Incorporating actions into the retrieval key means the model evaluates the robot’s ongoing behavior during comparison, moving beyond judgments based solely on visual appearance.

The model is trained jointly on three objectives. Predictive alignment uses bidirectional InfoNCE loss, treating query samples and their corresponding future segments as positive pairs, while samples with conflicting tasks, objects, or actions serve as negatives; semantically ambiguous, overlapping, or repetitive candidates are not hard-labeled as negatives. Multimodal localization distinguishes targets, ensuring similarity scores remain sensitive to input changes: replacing any modality or introducing semantic conflict results in score degradation. Representation regularization constrains the centers of two-batch vectors near zero to prevent representation collapse into low-rank spaces. These three objectives work together, enabling the vector space to retrieve relevant experiences while also determining whether a segment can logically follow the current action.

02 Accumulating 1.69 Million Experience Pairs

Training data was sourced from four datasets. After filtering, the final collection comprises 1,695,684 pairs of "current state – future segment" samples.

• EgoDex: First-person human hand motion data, contributing over 1.34 million entries;

• DROID: Franka dual-arm real-robot trajectories, more than 26,000;

• RoboMIND: Franka dual-arm real-robot data, over 180,000 entries;

• LIBERO: Simulation environment data, over 130,000 entries.

During the training phase, equal sampling was applied to four data sources, and pairing was performed exclusively within each source. Although EgoDex, the largest dataset, accounts for nearly eighty percent of the overall corpus, the team adopted an equal sampling strategy to prevent massive amounts of human data from suppressing the learning performance of the three types of robot data.

Each sample undergoes a four-step processing pipeline. Researchers first split the complete trajectory into non-overlapping segments, densely selecting time anchors within each segment and reserving the subsequent 4 seconds of content after each anchor. Once an anchor is determined, SAM 3 is used to locate the interactive object based on the instruction, cropping a 320×192 local zoomed-in image. The cropped image also undergoes quality verification by Qwen3-VL, requiring that the target object be visible, match the task instruction, be distinguishable from similar objects, and be unobstructed; samples failing these criteria are returned for re-selection with an earlier timestamp. A verified anchor, paired with the subsequent 4-second action sequence and 16 frames of video, constitutes one training sample. The entire filtering process uses only instructions and current observation information, ensuring no future data leakage.

The training backbone uses Cosmos3-Edge 4B, a 28-layer Transformer. The backbone and shared projection head use a lower learning rate, while the adapter and action projection head use a higher learning rate; training runs for 26,275 steps with a global batch size of 512, and the video tokenizer is frozen throughout. Evaluation uses the model checkpoint from the final iteration, sampling four data sources in a 1:1:1:1 ratio.

03 Missing Actions Hurt More Than Missing Frames

The evaluation benchmark, Robota-Embed-Bench, integrates human, simulation, and real-robot experiences into a unified retrieval pool. The setup uses 4,000 query samples to search within a shared candidate library containing 8,000 video segments. Each of the four data sources provides 1,000 queries derived from 800 complete task trajectories. For each query set, there are 1,000 strictly matched future segments and 1,000 neighboring segments from the same trajectory. Two statistical metrics are used: at the clip level, success requires retrieving the exact matching target future segment; at the episode level, any result originating from the same task trajectory is considered relevant. Ranking relies solely on cosine similarity of normalized vectors, with no additional reranking steps.

The experiment compares eight baseline models of similar scale, including text-to-video retrieval models like CLIP and unified vision-language embedding models. Across all four data sources, Robota-Embed achieves an average clip-level R@1 of 47.33% and an episode-level R@1 of 81.77%. The two best-performing baselines without action information—Gemini Embedding 2 and Omni-Embed-Nemotron-3B—achieve clip-level R@1 scores of 25.24% and 48.59%, respectively. Robota-Embed leads these by 22.09 and 33.18 percentage points, respectively. Performance varies by dataset: DROID is the easiest, reaching a clip-level R@1 of 74.10%, while EgoDex is the most difficult at only 36.70%. The team explains that EgoDex contains many visually similar segments within single trajectories, making it relatively easy to identify the correct trajectory but very hard to pinpoint the specific target pair. This distinction highlights the difference between clip-level and episode-level metrics.

The performance gain is not inherent to the backbone network. Using the same backbone to extract vectors directly yields a clip-level R@1 of just 0.65%; this rises to 47.33% after specialized retrieval training. In controlled split experiments on DROID, training only the read head resulted in a score of 13.8%, whereas jointly training the backbone and read head achieved 78.7%. The study also quantified the contribution of three objective functions: using only the prediction alignment objective reached 75.6%; adding the regularization objective improved it to 77.0%; and incorporating the multimodal localization discrimination objective brought the final score to 78.7%. This indicates that the latter two objectives primarily optimize retrieval ranking precision rather than building recall capability from scratch.

Controlled retraining experiments provide core evidence supporting the claim that "action must be included in retrieval": removing the action modality caused the clip-level R@1 to drop by 18.9 percentage points, while removing the visual modality resulted in only a 1.9 percentage point drop. However, input masking experiments on the fully trained model showed different results: removing action information from the query caused a 27.25 percentage point drop, while removing visual information caused a 24.33 percentage point drop, narrowing the gap. The team emphasizes that these experiments measure different dimensions: the retraining experiment observes whether remaining modalities can compensate for missing information, while the masking experiment measures the trained model's dependency on various inputs. Therefore, the smaller decline in visual metrics during retraining does not imply the final model can discard visual input.

This section also presents engineering performance data. When the candidate library size expanded from 1 million to 15 million segments, the total latency for query encoding and retrieval on a single NVIDIA A100 GPU increased from 160.36 milliseconds to 195.17 milliseconds. Encoding time remained fixed at 157.70 milliseconds, with the incremental delay coming primarily from the retrieval step. These timings exclude preprocessing of interaction regions, reference segment reranking, and downstream policy inference; actual deployment would require adding overhead from these stages.

04 How to Use Retrieved Experiences

The retrieved fragments serve only as intermediate outputs. The team has designed two technical pathways to integrate retrieval experience with robot strategies. The first pathway, Embed2ICL, offers two schemes: a frozen model and a fine-tunable model. The frozen scheme utilizes retrieved real-robot videos to initialize the fixed DreamZero policy visual cache. The fine-tunable scheme converts retrieved and filtered human demonstrations into context tokens, which are then input into the fine-tuned Fast-WAM model.

The freezing scheme was tested across three tasks in DROID-Sim, with 100 rounds of testing set for each group. Without reference pre-warming, the average success rate was 47.0%; when reference clips were randomly selected, the success rate actually dropped to 36.0%; using Robota-Embed to retrieve references could increase it to 50.3%, and after adding a reranking step, it reached 56.3%. These results confirm that choosing the wrong reference sample can yield even worse results than not using any reference at all.

The fine-tuning approach was validated on five high-difficulty task subsets of LIBERO. The model's baseline version achieved an average success rate of 86.4%, which increased to 92.8% after incorporating retrieved context tokens; replacing these with random context tokens resulted in only 88.0%. The performance improvement stems from the effective samples filtered by retrieval, rather than gains from simply increasing the number of input tokens.

In real-robot experiments, the team adapted DreamZero to the YAM robot, completing basic training using only 55 demonstration trajectories. In tests of the "place banana in bowl" task with the frozen model, average task progress increased from 21.67 to 40.00, and task success rate rose from 0 to 16.7%. With fine-tuning, task progress improved from 75.00 to 79.17 and from 66.67 to 87.50 across two tasks. The frozen approach nearly doubled task progress, but the absolute success rate remained low, indicating that retrieval-based references provide positive assistance but there is still a gap toward reliably completing tasks.

The second path is Embed2Evolve, which applies retrieval mechanisms to task memory management. Leveraging the GPT-6 Astra policy with frozen weights, the robot stores the trajectory of each attempt—regardless of success or failure—into its task memory after every round. The retrieval module dynamically adjusts its configuration based on execution feedback: controlling the number of recalled segments, the frequency of memory refreshes, and the scale of historical retention. The memory bank and retrieval configurations iterate in synergy, while the policy weights remain fixed throughout.

In the five ManiSkill tasks, each setting was repeated with twenty new initializations. The native GPT-6 Astra baseline achieved an average success rate of 24%; after incorporating fixed random historical samples, it dropped to 20%. With Embed2Evolve enabled, the success rate reached 68%. Under the same experimental setup, replacing the retrieval model with Omni-Embed-Nemotron-3B for iteration yielded a success rate of only 33%. Meanwhile, Embed2Evolve reduced policy token consumption by approximately 79.5% compared to that baseline; relative to the native strategy, token usage decreased by 43.7%, and reference interface costs fell from USD 9.96 to USD 5.80, with cost savings primarily driven by input-side optimization.

Both approaches share the same set of fused action-aware retrieval vectors. On one hand, the policy can be plug-and-play, reusing external demonstration experience; on the other hand, the robot can accumulate exclusive memories that grow continuously as tasks progress. The team proposes future research directions focusing on action spaces, robot bodies, and variable-length retrieval, aiming to make the continuation duration of retrieval segments adaptively adjust according to task types and execution phases, rather than remaining fixed at four seconds.