PaperlayerView exact revision on arXiv

Source-linked AI summary

TraceVLA: Visual Trace Prompting Enhances Spatial-Temporal Awareness for Generalist Robotic Policies

Ruijie Zheng, Yongyuan Liang, Shuaiyi Huang, Jianfeng Gao, Hal Daumé, Andrey Kolobov, Furong Huang, Jianwei Yang

arXiv:2412.10345v3cs.ROcs.AI

TL;DR

VLA models struggle to use spatial-temporal history for interactive manipulation. The paper introduces visual trace prompting, which overlays tracked movement trajectories on image inputs, and reports stronger generalization across simulated and real-robot evaluations while offering a more efficient compact model.

  • Problem

    VLA-powered robots often make decisions reactive to current inputs because they struggle to maintain awareness of past movements.

  • Method

    Visual trace prompting encodes historical robot movement as multi-point trajectories overlaid on image inputs for VLA action prediction.

  • Results

    TraceVLA consistently outperforms existing VLA models across diverse simulated environments, physical WidowX tasks, embodiments, and environmental variations.

  • Takeaways & Limitations

    Visual traces provide a spatial-temporal input that improves VLA adaptation and generalization across manipulation-task variations.

  • Takeaways & Limitations

    Trace length requires dataset-specific tuning because traces that are too short provide limited history, while overly long traces can obscure scene elements and distract the model.

Abstract

from arXiv · show

Although large vision-language-action (VLA) models pretrained on extensive robot datasets offer promising generalist policies for robotic learning, they still struggle with spatial-temporal dynamics in interactive robotics, making them less effective in handling complex tasks, such as manipulation. In this work, we introduce visual trace prompting, a simple yet effective approach to facilitate VLA models' spatial-temporal awareness for action prediction by encoding state-action trajectories visually. We develop a new TraceVLA model by finetuning OpenVLA on our own collected dataset of 150K robot manipulation trajectories using visual trace prompting. Evaluations of TraceVLA across 137 configurations in SimplerEnv and 4 tasks on a physical WidowX robot demonstrate state-of-the-art performance, outperforming OpenVLA by 10% on SimplerEnv and 3.5x on real-robot tasks and exhibiting robust generalization across diverse embodiments and scenarios. To further validate the effectiveness and generality of our method, we present a compact VLA model based on 4B Phi-3-Vision, pretrained on the Open-X-Embodiment and finetuned on our dataset, rivals the 7B OpenVLA baseline while significantly improving inference efficiency.

1 INTRODUCTION

The paper identifies limited spatial-temporal awareness as a weakness in VLA-powered manipulation and proposes visual trace prompting to provide movement history. TraceVLA and its compact counterpart are evaluated across simulated and real-world settings for generalization.

  • VLA-powered robots often react to current inputs because they struggle to maintain awareness of past movements.
  • Visual trace prompting overlays multi-point trajectories of past robot movements onto image inputs to enhance spatial-temporal reasoning.
  • TraceVLA is a 7B-parameter model fine-tuned from OpenVLA on 150K robot manipulation trajectories annotated with visual traces.
  • The compact TraceVLA-Phi3 uses a 4B-parameter Phi-3-Vision backbone pretrained on Open X-Embodiment data and offers improved inference efficiency with robust performance.
  • Evaluations across 131 SimplerEnv settings and physical WidowX experiments show consistent improvement over existing VLA models across embodiments and environments.
  • The paper contributes visual trace prompting, a trace-annotated dataset, and 7B and 4B VLA models for improving manipulation performance and generalization.

2 PRELIMINARIES

Vision-language-action models extend vision-language models to robot control by combining visual encoding, projection, language-model processing, and discretized action prediction. Generalist policies aim to support diverse sensors, action spaces, and robotic platforms.

  • VLAs extend VLMs by predicting discretized robot actions from visual and language inputs.
  • A typical VLA architecture contains a visual encoder, a projector, and a large language model backbone.
  • Action discretization maps continuous robot actions into discrete tokens, typically by dividing each action dimension into bins.
  • Generalist policies target diverse sensors, action spaces, robotic platforms, and scenarios rather than a single task setting.

3 TRACEVLA

TraceVLA supplies VLA models with visual memory by extracting active point trajectories from historical observations and overlaying sampled traces on the current image. The architecture preserves the original observation and uses training and inference procedures designed for efficiency and missing traces.

  • 3.1 VISUAL TRACE PROMPTING: Visual trace prompting replaces redundant historical-frame concatenation with tracked trajectories overlaid on the robot’s original observation.
  • 3.1 VISUAL TRACE PROMPTING: Co-Tracker extracts dense point trajectories from a historical image window, while active trajectories are selected by significant movement over N timesteps.
  • 3.1 VISUAL TRACE PROMPTING: A random sample of M active trajectories is overlaid on the original observation to provide spatial information about historical states and actions.
  • 3.2 MODEL ARCHITECTURE: TraceVLA feeds both the trace-overlaid image and original image, separated by a special token, so traces do not obscure key objects or the end-effector.
  • 3.2 MODEL ARCHITECTURE: Training uses trace dropout to prepare the model for test-time cases where visual traces are unavailable.
  • 3.3 IMPLEMENTATION DETAILS: The training dataset contains approximately 150,000 trace-annotated robot trajectories collected from BridgeData-v2, Google RT1, and WidowX demonstrations.
  • 3.3 IMPLEMENTATION DETAILS: During inference, TraceVLA reduces tracking cost by identifying active points densely at t=0 and sparsely tracking them thereafter.
  • 3.3 IMPLEMENTATION DETAILS: The models include OpenVLA-based TraceVLA and a pretrained 4B Phi-3-Vision VLA used to test a lightweight alternative.

4 EXPERIMENT

TraceVLA is evaluated against generalist policy baselines in SimplerEnv and real-world WidowX manipulation tasks, including environmental variations and unseen tasks. Ablations examine whether gains arise from visual trace prompting rather than finetuning alone, historical image sequences, text traces, or trace length.

  • Evaluation setup: TraceVLA and TraceVLA-Phi3 are compared with OpenVLA, OpenVLA-Phi3, Octo-Base, and RT1-X across simulation and real-robot settings.The evaluation covers three SimplerEnv tasks with visual matching and variant aggregation metrics, plus eight WidowX tasks including four unseen tasks.
  • Simulation evaluation: 2.4%–12.7% improvements are reported for TraceVLA over OpenVLA across all three SimplerEnv tasks and evaluation metrics.Table 1 reports performance under visual matching and variant aggregation, with overall performance averaged across results.
  • Environmental variants: Over 20% average improvement is reported under camera-orientation, distractor, and background variations.The visual trace is described as helping preserve spatial trajectory information across viewpoint and scene changes, including lighting and table-texture variations.
  • Real-robot evaluation: TraceVLA achieved 8/10 successful trials on the unseen pick-place corn task, compared with 1/10 for OpenVLA.The result is presented as generalization from related training tasks such as picking and placing eggplant in a pot.
  • Ablation studies: Finetuning alone produced a 1.1% gain for OpenVLA and a 0.3% decrease for OpenVLA-Phi3, whereas visual trace prompting raised OpenVLA success to 47.7%.These comparisons use the same finetuning dataset with and without visual trace prompting.
  • Ablation studies: Finetuning OpenVLA with six historical images reduced performance by 6%, while text traces gained 2.4% and TraceVLA gained an additional 6.4% over that alternative.The paper attributes the historical-image drop to redundant visual information and notes that text traces add approximately 150 tokens compared with visual traces.
  • Ablation studies: A three-step visual trace improved performance by 3.2%, while traces that are too short or too long can provide less useful context.The paper describes N = 3 as often too brief and N = 12 as potentially obscuring scene elements or overlapping previous motion traces.

5 LIMITATION ANALYSIS: TRAINING MEMORY COST AND INFERENCE SPEED

TraceVLA adds visual-trace processing and CoTracker overhead while keeping transformer-token inference costs negligible; dense tracking is cheaper per step than five-point tracking.

  • TraceVLA adds an image input and CoTracker tracking, so the evaluation measures both training GPU memory and inference speed.Tests use eight H100 GPUs for training-memory measurements and one H100 GPU for inference-speed measurements.
  • Approximately 300 additional image and text tokens add around 0.002 seconds per timestep, producing negligible transformer inference overhead.
  • M = 5 CoTracker point tracking adds 0.03 seconds per step, whereas dense 40 × 40 tracking has an amortized cost of 0.004 seconds per step.
  • Figure 10 compares GPU memory costs and inference times for 7B TraceVLA, OpenVLA, 4B TraceVLA-Phi3, and OpenVLA-Phi3.

6 RELATED WORK

The paper positions visual trace prompting as a distinct generalist-policy approach that injects historical visual trajectories into VLA models to improve spatial-temporal context.

  • Unlike LLARVA’s textual future 2D traces, TraceVLA integrates historical visual traces directly into VLA models as visual prompts.
  • The method targets spatial and temporal dynamics, an aspect described as previously underexplored in VLA models.
  • The work contributes a spatial-temporal dataset across diverse embodiments and state-of-the-art 7B and 4B VLA models.
  • Future directions include multi-point spatial trajectory prediction and 3D point-cloud training for richer spatial representations.

A REAL ROBOT TASKS SETUP

The real-robot setup evaluates trained and unseen tasks involving folding, sweeping, pick-and-place, lifting, and pushing under randomized object arrangements and distractors.

  • Each trial randomizes the target object’s initial location, and most tasks add 2–3 random distracting objects to the scene.
  • Tasks included in the finetuning dataset: The finetuning task set includes folding cloth, swiping corn into a sink, picking corn into a pot, and placing a knife on a plate.
  • Unseen tasks for generalization: The unseen-task set includes placing a banana right of a plate, placing an eggplant on a plate, lifting an AAA battery, and pushing cloth rightward.
  • Unseen tasks for generalization: The unseen tasks test generalization to unseen objects, goals, and motion patterns relative to the finetuning data.

B QUALITATIVE RESULTS ON REAL ROBOT ROLLOUTS

Qualitative WidowX rollouts compare OpenVLA and TraceVLA on banana placement, cloth folding, and eggplant placement, with TraceVLA showing smoother task completion and stronger spatial reasoning.

  • Figures 11–13 visualize OpenVLA and TraceVLA rollouts for pickplace banana, fold cloth, and pickplace eggplant.
  • TraceVLA accurately places the banana and eggplant and grasps the cloth edge, while OpenVLA often overfits to the finetuning distribution.
  • OpenVLA places the banana on the plate instead of to its right, illustrating a failure to follow the spatial instruction.

C.1 THICKNESS, TRANSPARENCY, AND COLOR OF VISUAL PROMPTING

TraceVLA remains stable across reasonable visual-trace thickness, transparency, and color choices, reducing the need for extensive parameter tuning. Additional comparisons also examine historical observation lengths and related visual-prompting work.

  • C.1 THICKNESS, TRANSPARENCY, AND COLOR OF VISUAL PROMPTING: Performance differences from varying trace line thickness are minor on SimplerEnv Average Success Rate.The effect is reported in Table 2.
  • C.1 THICKNESS, TRANSPARENCY, AND COLOR OF VISUAL PROMPTING: TraceVLA performance is robust to transparency adjustments controlled by the α parameter.Lower α values produce more transparent traces, with findings summarized in Table 3.
  • C.1 THICKNESS, TRANSPARENCY, AND COLOR OF VISUAL PROMPTING: Color-scheme changes produce negligible differences in TraceVLA success rates.The default RYPBG scheme was compared with the alternative POBGG scheme in Table 4.
  • C.1 THICKNESS, TRANSPARENCY, AND COLOR OF VISUAL PROMPTING: Thickness, transparency, and color have negligible performance impact when selected within reasonable ranges.This reduces the need for extensive hyperparameter tuning and supports reliability across visualization settings.
  • C.2 BASELINE WITH DIFFERENT STEPS OF HISTORICAL OBSERVATIONS: TraceVLA consistently and significantly outperforms OpenVLA when the models are fine-tuned with 2 or 3 observation-history steps.The comparison is reported on SimplerEnv; OpenVLA shows a slight improvement with 2-step history.
  • D ADDITIONAL RELATED WORK: Visual prompting methods provide pixel-level instructions on multimodal inputs and have been applied broadly in robotics.The related-work discussion cites several visual-prompting approaches and robotics applications.

E MORE IMPLEMENTATION DETAILS

TraceVLA reduces inference overhead by tracking a small set of active points after initial dense initialization, while periodically recalibrating when tracking degrades. Its inference loop overlays traces before action prediction and uses untraced prompts during initial timesteps.

  • Efficient trace extraction: Inference begins by querying Co-Tracker on a K × K grid, then samples M active points for efficient subsequent tracking.This avoids querying the full grid at every timestep.
  • Track maintenance: TraceVLA periodically recomputes dense K × K tracking to recalibrate after Co-Tracker loses track over long intervals.Track loss is observed after roughly 30 to 40 steps, depending on action magnitude.
  • Inference algorithm: Algorithm 1 provides the Python-style pseudocode for TraceVLA inference.The implementation details specify parameters including grid size, tracked-point count, trace length, and redraw interval.
  • Inference loop: The inference algorithm maintains a historical observation queue with maximum length N while iterating through up to T timesteps.The pseudocode initializes the queue after environment reset and loops over t from 0 to T.
  • Efficient trace extraction: After initialization, Co-Tracker follows the existing tracked points and updates them using their latest-frame positions.The tracked trace has shape 2 x N x M in the pseudocode.
  • Action prediction: Once traces are available, they are overlaid on the image and supplied with the original image to TraceVLA for action computation.The pseudocode passes [image, image_overlaid] together with the trace prompt.
  • Inference loop: For initial timesteps, the model receives the same image twice without a trace hint before actions advance the environment and populate history.Trace prompting is introduced after sufficient historical observations are available.

F ADDITIONAL EXPERIMENTAL RESULTS ON LIBERO SIMULATION

The LIBERO evaluation measures multitask success across four suites that vary spatial layouts, objects, goals, and task horizons. OpenVLA and TraceVLA-7B are fine-tuned on combined suite data and evaluated per suite.

  • Benchmark setup: LIBERO evaluation covers four suites—LIBERO-long, LIBERO-Spatial, LIBERO-Object, and LIBERO-Goal—with 10 tasks and 50 demonstrations per task.The study evaluates multitask performance of pretrained VLA policies across these suites.
  • LIBERO-Spatial: LIBERO-Spatial keeps objects fixed while varying layouts to test spatial-relationship understanding.Its example asks the model to place a bowl relative to other objects.
  • LIBERO-Object: LIBERO-Object keeps scene layouts consistent while changing objects to evaluate object-type understanding.Its example involves placing alphabet soup in a basket.
  • LIBERO-Goal: LIBERO-Goal keeps objects and layouts fixed while varying task goals to assess diverse task-oriented behaviors.Its example requires placing two food items in a basket.
  • LIBERO-Long: LIBERO-Long, also called LIBERO-10, uses long-horizon tasks with diverse objects, layouts, and goals.The suite challenges extended planning and execution, such as opening a cabinet drawer.
  • Training and evaluation: Both OpenVLA and TraceVLA-7B are fine-tuned on the combined data from the four LIBERO suites after unsuccessful trajectories and inactive action steps are filtered.TraceVLA additionally receives visual-trace annotations following the described procedure.
Loading 2412.10345v3…