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

Authors: Ruijie Zheng, Yongyuan Liang, Shuaiyi Huang, Jianfeng Gao, Hal Daumé III, Andrey Kolobov, Furong Huang, Jianwei Yang
Primary area: foundation or frontier models, including LLMs
Keywords: Vision Language Model, Robot Learning

Abstract:
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.

Decision: accept (poster)
Reviewer scores: 8, 8, 6, 6

Reviews and discussion:
<thread>
creation date:
1730327697010

thread_type:
official_review

<thread_object>
creator: reviewer_qxt2

<content>
summary:
This paper proposes overlaying the history trace on the RGB image as an indicator of spatial and temporal information to help large foundation models better manipulate objects. Experiments are conducted in both simulation and the real world.

strengths:
Although trace has been implemented in various aspects, from my knowledge, this is the first time it serves as a history spatial-temporal information indicator. Comprehensive experiments and analyses are conducted to demonstrate its effectiveness, and the authors also consider real-time inference by optimizing the inference process.

weaknesses:
1. The choice of active points depends on a threshold. Does this threshold need to be chosen differently for each robot or dataset?

2. In the analysis of trace length, the authors mention that a longer trace length can obscure key information. However, in the method, they state that they mitigate this by inputting both the original and overlay images. In that case, a longer trace length should not lead to worse results. This explanation seems self-contradictory."

3. For fine-tuning OpenVLA with history images, six history frames work best for TraceVLA, but this might be too redundant when using image. Would using three or two history frames be better when using history images as input?

questions:
Does the dataset used for fine-tuning OpenVLA contain scenes similar to those in simplerEnv? That is, are all the testing scenes entirely unseen in the training dataset?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730534708992

thread_type:
official_review

<thread_object>
creator: reviewer_4xpz

<content>
summary:
This paper introduces TraceVLA, a novel method which leverages visual traces to enhance the spatial-temporal awareness of VLA models. 
Specifically, the VLA model takes as inputs two images: a plain image and an image overlaid with traces of a set of active points.
The traces are generated by CoTracker and the active points are selected from a grid of dense points based on the changes in pixel locations.
The state-action history can be effectively communicated to the model through the provided visual traces.
Experiment results show that the proposed method outperforms comparing baseline methods in the SimplerEnv simulation environment and four real-robot tasks.
Ablation studies show that using visual traces is a more effective method for conveying historical information than prompting with text traces or appending historical image observation.

strengths:
The proposed method, TraceVLA, is a neat and effective approach for providing large VLA models with historical state-action information.
Visual traces help VLA models to maintain awareness of the robot movements and do not introduce redundant information compared to appending historical image observations.
Experiments in SimplerEnv show that TraceVLA outperforms multiple comparing baseline methods in overall performance.
The paper also evaluates on real-robot tasks to validate the proposed method in the real world.
Ablation studies offer good insights on how leveraging visual trace prompts compares to other prompting strategies.

weaknesses:
1. The experiments in the real world contains only one task for generalization assessment. It would be better to evaluate the proposed method in more generalization settings, including novel backgrounds and more novel tasks. This would provide a more comprehensive insight on the generalization capabilities of the proposed method.

2. RT1-X, despite having a much small number of parameters, outperforms TraceVLA-Phi3 with 3B parameters in SimplerEnv. The advantage of TraceVLA with 7B parameters is also not substantial. Is it possible to provide more discussions on potential reasons? Additionally, it would be beneficial to include RT1-X in real-robot experiments for comparison. In particular, observing how these two methods perform in generalization settings would be informative.

questions:
1. The proposed method leverages CoTracker to provide visual traces. Were any failure cases due to inaccurate tracking observed during experiments? How robust is TraceVLA in handling these inaccuracies?

2. As the active points are selected based on changes in pixel locations, will points on moving background objects be chosen as active points? And how does the proposed method perform when the active points are located on irrelevant background?

3. How many iterations were used for training during fine-tuning on the data in real-robot experiments? Is the data from the pre-training stage also included in fine-tuning as well?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730567802951

thread_type:
official_review

<thread_object>
creator: reviewer_gs1b

<content>
summary:
This work leverages visual trace prompting to enhance VLA models' spatial-temporal awareness for action prediction by visually encoding state-action trajectories. The authors curated a dataset of 150K robot manipulation trajectories using visual trace prompting and fine-tuned it on OpenVLA.

strengths:
- The paper introduces visual trace prompting, a novel technique that enhances spatial-temporal reasoning in VLA models for manipulation tasks by representing spatial-temporal relationships in robotic contexts.

- A curated visual trace prompting dataset was developed, with state-of-the-art 7B and 4B VLA models fine-tuned using this approach, providing an efficient method to improve VLA model performance.

- Their approach was rigorously validated through extensive evaluations across diverse simulated and real-world robot tasks, demonstrating superior generalization by leveraging spatial-temporal information.

weaknesses:
- While the author presented results highlighting the importance of historical information from the trace, there were no results demonstrating spatial reasoning capabilities as claimed in the "spatial-temporal" contribution. One potential approach could involve fine-tuning a Vision-Language Model (VLM) for line-drawing tasks, connecting it to low-level policy, and leveraging the VLM independently for spatial reasoning, similar to the RoboPoint framework.
  
- The author relied entirely on real-world data to generate the trace, which may have introduced variability from environmental factors such as lighting. Generating trace data in simulation could mitigate these issues, potentially minimizing the sim-to-real gap. Additionally, evaluating performance would be easier with ground truth data generated from a platform like RLBench across diverse tasks.

- Generalization in the Vision-Language Alignment (VLA) could be a strong aspect of this approach, as the data captures a general representation of actions conditioned on language. It would be beneficial for the author to test generalization on simulation benchmarks, such as the Colosseum, to showcase its robustness.

questions:
Please refer to weakness section for my questions. I would consider raising my rating if the weaknesses mentioned in my questions can be addressed. Overall, I find this to be a strong paper that offers a valuable perspective on training VLA models beyond simply adding more data.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730611979981

thread_type:
official_review

<thread_object>
creator: reviewer_xxsr

<content>
summary:
This paper proposed a visual trace prompting method called TraceVLA to enhance VLA models in manipulation tasks. 
By incorporating Co-Tracker to visually prompt keypoint trajectories into existing VLA frameworks, TraceVLA achieves better performance than baselines in simulated and real robot experiments. It also shows better generalziation across environmental variations than vanilla VLA model and text prompted ones, further proves the advantage of the proposed visual trace method.

strengths:
1. Clear description of the method. Comprehensive ablation on prompt setting and trace length, as well as analysis of memory and speed. 
2. Regarding the drawback in VLA practice, this work creatively applies tracking as visual prompting for VLAs to learn the invariance across different scene settings in manipulation.

weaknesses:
1. Sec 4.1 is a bit hard to read. Mixing figures, tables, and your main context together with reduced blanks is not a good idea to present your experiment results clearly. 
2. Lack of details in experiments. For example, how do you define success, especially in real robot experiments? When you add noise to the environment, do you have a limit over related parameters? Is there any case study that qualitatively shows the advantage of applying visual tracking prompts? Please add them to your revision, which is also good for reproducibility.

questions:
1. How does this method generalize under camera orientation changes? By nature it helps the model against illumination/background changes, but how can VLA remain effective when the camera moves, given only 2D tracking result? If the change is small, then to what degree do you vary? You mentioned in the real robot section that finetuning is needed due to domain shift, but not for the simulation. Is that also due to the domain gap being small in simulated experiments?
2. Regarding the active point selection, did you choose the $\kappa$ so that $M=5$ takes a reasonable ratio in the active points, and therefore similar traces are learned in the training set? If not, does it have the potential to generalize across visually different robots?
3. Will the color, thickness, or even transparency of the visual prompt affect the performance? 
4. Is K somehow related to the size of the image patch size? 
5. Will the performance drop if the action is moving into/out of the camera frame, or rotating in/out of it?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732353420969

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
## Review Highlights

We sincerely thank all reviewers for their thorough feedback and constructive suggestions. We particularly appreciate the acknowledgment of:

* Novelty of our proposed visual trace prompting method (**Reviewer Gs1B, 4XPZ**)
* Clear presentation and methodology (**Reviewers xXSr, cSN9, vJK6**)
* Technical contributions including efficient methodology and our proposed dataset (**Reviewers xXSr, Gs1B, 4XPZ, qxt2**)
* Comprehensive empirical validation demonstrating strong generalization (**Reviewers Gs1B, 4XPZ, qxt2**)

## Main Concerns Addressed

* Generalization Capacity (Reviewers xXSr, Gs1B, 4XPZ):
We extended evaluation on LIBERO benchmark and additional validation on out-of-domain real-robot tasks with novel objects, goals, language instructions, and motion scenarios, further highlighting the substantial generalization improvement from visual trace prompting across diverse simulation and real-world environments
* Trace Quality and Tracking Robustness (Reviewers xXSr, 4XPZ):
We conducted additional ablation studies on visualization parameters show TraceVLA maintains consistent performance under various trace conditions and illustrated its robustness in imperfect tracking scenarios.

These responses address key concerns while providing new experimental results that strengthen our paper's contributions. We incorporate our clarifications and additional analyses in our revised manuscript.

## Summary of paper updates
Main Text:
* Reorganized **Section 4.1** for better readability
* Enhanced related work discussion with newly suggested references (**Section 6**)

Appendix:
* Detailed real-robot experimental setup and protocols (**Appendix A**)
* Out-of-domain real-robot experiments to further demonstrate generalization (**Appendix B**)
* Qualitative analysis of real-robot rollouts (**Appendix C**)
* Ablation studies on trace visualization parameters (**Appendix D**)
* Comparative study with OpenVLA using varying historical observations (**Appendix E**)

These updates are highlighted in red for easy reference. We believe our work makes significant contributions to the ICLR robotics community, further strengthened by the reviewers' constructive feedback.

We sincerely thank the reviewers and AC for their thoughtful comments during the discussion period. We hope our responses have adequately addressed all concerns and welcome any further discussion.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper introduces TraceVLA, a vision-language-action model enhanced with visual trace prompting to improve spatial-temporal awareness in robotic manipulation, achieving state-of-the-art performance and robust generalization across simulated and real-world tasks with improved inference efficiency.

All reviewers acknowledge the contributions of this work, highlighting its (1) novelty, (2) technical advancements, including the efficient methodology and proposed dataset, (3) clear presentation, and (4) comprehensive empirical evaluation.

The authors have also effectively addressed the reviewers' concerns during the discussion period. All reviewers are in unanimous agreement to accept this paper.