TraceVLA: Enhancing Robotic Spatial-Temporal Skills

• TraceVLA is a vision-language-action model that introduces visual trace prompting, encoding a robot’s state–action history to enhance spatial-temporal reasoning for action prediction.
• It leverages dual visual streams and active-point selection on trace overlays, leading to significant performance improvements in both simulated and real-world robotic tasks.
• TraceVLA demonstrates robust generalization and efficiency, with compact variants reducing computational overhead while outperforming baseline models.

TraceVLA is a vision-language-action (VLA) model architecture and methodology designed for generalist robotic learning, with a focus on enhancing spatial-temporal awareness in action prediction by introducing visual trace prompting. This approach involves encoding a robot’s recent state–action trajectory as an overlaid trace image, which is processed jointly with raw RGB observation to achieve state-of-the-art action policy performance and robust generalization across diverse robotic embodiments. TraceVLA builds upon the OpenVLA backbone, applies advanced trajectory tracking and rendering, and is validated on extensive simulation and real-world robot benchmarks. A compact variant using the Phi-3-Vision VLM demonstrates that spatial-temporal gains are achievable with efficient, lower-parameter models (Zheng et al., 13 Dec 2024).

1. Visual Trace Prompting Concept

At the core of TraceVLA is visual trace prompting: the representation of robot state–action history as a visual overlay. For a trajectory $\tau = \{ (s_t, a_t) \}_{t=1}^T$, where $s_t \in \mathbb{R}^{H \times W \times 3}$ is the RGB observation and $a_t$ the continuous action at time $t$, a sliding window of $N$ frames is processed with Co-Tracker to yield a dense grid of $K \times K$ ($1600$) point trajectories $P = \{ p^{(i)} \}$. Active-point selection retains traces with total $\ell_1$ displacement over threshold $\kappa$, and $M$ traces are sampled to form $\tilde{P}$. The overlay function $\mathrm{Overlay}(s_t, \tilde{P})$ renders each sampled trace as colored polylines on the latest frame, producing an augmented input $I_{trace}$.

At each timestep, the model receives both the raw RGB image $I_{rgb}$ and $I_{trace}$, which encodes a concise, history-anchored cue for spatial-temporal reasoning in manipulation tasks.

2. Architecture and Training Objectives

TraceVLA adopts the OpenVLA architecture, extending it to accommodate dual visual streams. The vision encoder $\mathcal{F}_\phi$ yields patch embeddings; a linear projector $W_{proj}$ aligns these with the token dimension for the LLM (Prismatic-7B). Images $I_{rgb}$ and $I_{trace}$ are tokenized into $V^{(0)}$ and $V^{(1)}$, separated by a special $\langle \mathrm{IMG\_SEP}\rangle$ token, and input with a guiding text prompt.

The training objective is strictly next-token cross-entropy for discrete action prediction, with each action axis quantized into $256$ bins: $$L_{\text{action}} = - \sum_{t=1}^T \log P(a_t \mid I_{rgb,1:t}, I_{trace,1:t}, \text{TextPrompt})$$ No auxiliary losses were employed; empirical results indicate the model leverages $I_{trace}$ for improved prediction via $L_{\text{action}}$ alone. Robustness is enforced by trace dropout: with probability $\alpha$ during training, $I_{trace}$ is replaced by $I_{rgb}$ and the prompt is modified, promoting fallback to RGB-only cues when tracking fails.

3. Dataset Construction and Preprocessing

The fine-tuning dataset comprises $150,000$ robot manipulation trajectories from BridgeData-v2 ($\sim$80K), Google RT-1 ($\sim$50K), and WidowX-250 real-robot demos ($120$ demos across four physical tasks). Each trajectory is cut into overlapping $2N=12$-frame segments and processed with Co-Tracker to extract $K=40$ grid points over $N=6$ frames. After active trace selection (typically $10$–$20\%$ of points), $M=5$ traces are randomly sampled and rendered for $I_{trace}$. No further class balancing, augmentation beyond active-point randomness, or additional regularization was applied.

4. Evaluation Methodology and Results

TraceVLA is validated in both simulated and physical settings.

SimplerEnv (Simulation)

A total of $137$ configurations across Move Near, Pick Coke Can, and Open/Close Drawer tasks, under numerous domain shifts (lighting, backgrounds, distractors), provide a comprehensive testbed for spatial-temporal policy generalization. TraceVLA-7B achieves $47.7\%$ overall success (+7.5 over OpenVLA-7B), with per-task gains up to $+12.7\%$. Ablation confirms that visual trace overlays, rather than simple frame histories or text prompts, yield the full improvement (+6.4\% for visual traces vs +2.4\% for text-only).

Model	Overall (%)	Notes
OpenVLA-7B	40.2	Baseline
TraceVLA-7B	47.7	+7.5 with trace prompting
OpenVLA-Phi3-4B	39.9	—
TraceVLA-Phi3-4B	44.0	+4.1 with trace prompting

WidowX-250 Real-Robot Evaluation

Four tasks (Fold Cloth, Swipe Corn into Sink, Pick/Place Corn into Pot, Pickup Knife→Plate) and generalization tests on four unseen tasks demonstrate a $3.5\times$ improvement in average success rate versus OpenVLA (e.g., Pick/Place Corn: $9/10$ TraceVLA vs $3/10$ baseline). Generalization to unseen objects reaches $>50\%$ for TraceVLA ($<20\%$ baseline).

Ablation Insights

• TraceVLA fine-tuned without trace overlays yields marginal improvements (+1.1\% to $-0.3\%$).
• Replacing traces with raw frame histories reduces performance by $6\%$.
• Optimal trace length $N=6$; longer traces ($N=12$) slightly degrade results.

5. Compact 4B-Parameter Variant

To demonstrate deployment efficiency, OpenVLA-Phi3 (4B) is pretrained on the $970$K Open-X-Embodiment trajectories and finetuned on the $150$K trace-augmented dataset to produce TraceVLA-Phi3. This variant achieves $44.0\%$ on SimplerEnv and matches the real-robot gains of the larger TraceVLA-7B, outperforming the original 7B baseline.

Inferential efficiency increases: 4B TraceVLA-Phi3 uses approximately half the GPU memory ($-8$ GB at batch $32$) and runs $20\%$ faster per inference step. Overhead introduced by visual trace prompting remains minimal ($+0.002$ s for tokens, $+0.03$ s for trace extraction, $+0.004$ s for periodic reinitialization).

6. Significance, Limitations, and Prospective Directions

TraceVLA introduces an efficient, empirically validated mechanism for boosting spatial-temporal awareness in VLA-based robotic policies. Encoding state–action history as an overlay offers semantic compression, aiding scene understanding and prediction—particularly in manipulation domains requiring nuanced spatial reasoning. Visual trace prompting consistently outperforms baseline temporal encoding approaches.

The approach’s simplicity facilitates practical integration: minimal computational overhead, compatibility with low-parameter models, and direct improvement in generalist robot performance. A plausible implication is that trajectory overlays may serve as a generic enhancement in broader VLM-based sequential decision-making.

Limitations include necessary dependence on reliable trajectory extraction (Co-Tracker), modest performance ceilings (e.g., diminishing returns with longer traces), and inference speed constraints on resource-limited deployments.

Future exploration may address: (1) improved active-point selection strategies; (2) integration with continous control/action domains; (3) application to non-manipulation tasks requiring long-term spatial reasoning; and (4) further reduction in resource requirements for on-device use. The methodology provides a candidate blueprint for augmenting spatial-temporal perception in generalist robot agents and multi-modal sequential prediction frameworks (Zheng et al., 13 Dec 2024).