PRTS: A Primitive Reasoning and Tasking System via Contrastive Representations

1. Problem Statement

Current Vision-Language-Action (VLA) models have advanced robotic control by leveraging Vision-Language Models (VLMs) as perceptual backbones, typically using a dual-system architecture (System 2 for reasoning, System 1 for control). However, existing approaches suffer from a fundamental limitation:

• Static vs. Temporal: They predominantly frame pre-training as supervised behavior cloning (BC), focusing on static visual-linguistic semantic correlations. They lack an intrinsic understanding of goal reachability—the quantitative probability that a current state-action pair will lead to a language-specified goal over time.
• Value Estimation Bottlenecks: Attempts to incorporate value learning often rely on separate value networks, require manual reward annotations, or use Temporal-Difference (TD) learning which suffers from bootstrapping errors and high variance in long-horizon tasks.
• Generalization Gap: Without explicit temporal awareness, these models struggle with long-horizon planning, zero-shot novel instructions, and recovering from human interventions because they cannot distinguish between "easy-yet-erroneous" states and "difficult-yet-reachable" ones.

2. Methodology: PRTS

The authors propose PRTS, a VLA foundation model that reformulates pre-training through Goal-Conditioned Reinforcement Learning (GCRL) using Contrastive Reinforcement Learning (CRL).

A. Core Concept: Contrastive RL for VLA

Instead of predicting actions directly or using a separate value network, PRTS learns a unified embedding space where the inner product of a state-action embedding ($\phi(s, a)$) and a goal embedding ($\psi(l)$) approximates the log-discounted goal occupancy.

• Theoretical Basis: Based on the equivalence between the Q-function for goal-reaching rewards and the conditional occupancy probability (Lemma 1).
• Objective: The model learns to discriminate between state-action pairs that can reach a specific language goal (positive samples) and those that cannot (negative samples).

B. Key Technical Innovations

1. Language-Conditioned Contrastive RL:

   • Unlike standard CRL which samples future states geometrically, PRTS treats the language instruction as the fixed goal for an entire trajectory.
   • Multi-Positive Problem: Since a single language goal applies to all timesteps in a trajectory, a batch contains many positive pairs for one goal.
   • Temporal Weighting: To recover the geometric sampling distribution of standard CRL, the authors introduce a temporal weighting scheme. State-action pairs closer to the end of a trajectory (closer to the goal) are weighted higher ($\gamma^{T-t}$). This allows the model to learn the temporal structure of goal-reaching without explicit reward labels.
   • Bidirectional Objectives:
     • State-Action $\to$ Language: Aligns $(s, a)$ with its specific instruction.
     • Language $\to$ State-Action: Forces the model to rank state-actions within a trajectory by their proximity to the goal, encoding temporal progress.

2. Single-Forward-Pass Architecture (Negligible Overhead):

   • PRTS appends two auxiliary token blocks to the input sequence: <CRL_action> and <CRL_goal>.
   • Role-Aware Causal Mask: A custom attention mask (implemented via a fused CuTe-FlashAttention kernel) ensures information isolation:
     • <CRL_action> tokens attend only to vision/proprioception (not language).
     • <CRL_goal> tokens attend only to the instruction.
     • Standard autoregressive action tokens maintain standard causal attention.
   • Result: The model produces discrete action logits, state-action embeddings, and goal embeddings in a single forward pass, incurring negligible computational overhead compared to vanilla behavior cloning.

3. Two-Stage Training Pipeline:

   • Pre-training: Jointly optimizes the VLM backbone on behavior cloning (next-token prediction), auxiliary tasks (Robot-QA), and the bidirectional contrastive losses. No action expert is attached yet.
   • Post-training: A Flow-Matching Action Expert (based on Diffusion Transformer) is attached and fine-tuned on downstream data. The pre-trained backbone provides the goal-reachability priors.

3. Key Contributions

1. Reformulation of VLA Pre-training: First to scale reward-free, language-conditioned Contrastive RL directly into the VLA pre-training phase, eliminating the need for manual reward annotations or separate value networks.
2. Intrinsic Goal-Reachability Awareness: The learned representations quantitatively encode the probability of reaching a goal, bridging the gap between semantic reasoning and temporal task progress.
3. Efficient Architecture: The "Role-Aware Causal Mask" and custom CuTe kernel enable dense goal-reachability supervision with near-zero computational overhead, making it scalable to large models (4B parameters) and massive datasets (167B tokens).
4. Unified Framework: Unifies semantic reasoning, discrete action prediction, and value estimation within a single VLM architecture.

4. Experimental Results

PRTS was pre-trained on 167.8 billion tokens from diverse datasets (AgiBotWorld, RoboMind, Open X-Embodiment, etc.) and evaluated on simulation and real-world benchmarks.

• Simulation Benchmarks (LIBERO, LIBERO-Plus, LIBERO-Pro, SimplerEnv):

  • SOTA Performance: PRTS achieves state-of-the-art results, matching or exceeding strong baselines (e.g., $\pi^0.5$, ABot-M0) while using significantly less post-training compute (e.g., 1/8th the budget of $\pi^0.5$ on LIBERO).
  • Long-Horizon Tasks: Significant gains on the "Long" suite of LIBERO (96.6% vs. 94.8% for matched budget baselines), proving the efficacy of temporal awareness.
  • Robustness: Outperforms baselines in zero-shot generalization under distribution shifts (camera, robot state, lighting, object layout) and Task/Position perturbations in LIBERO-Pro.
  • Zero-Shot Instruction: Massive improvement on novel instructions (31.5% vs. 0.8% for $\pi^0.5$ on the "Task" axis of LIBERO-Pro).

• Real-World Evaluation (RealMan Dual-Arm & Flexiv Single-Arm):

  • Complex Manipulation: Achieved 95.9% average success rate on 11 dual-arm tasks and 90.0% on 3 fine-grained single-arm tasks.
  • Human Intervention Recovery: PRTS demonstrated superior robustness when humans interrupted tasks (e.g., removing objects mid-execution). Unlike baselines that failed or replayed stale trajectories, PRTS successfully re-evaluated the state and re-planned to reach the goal.
  • Generalization: Maintained high success rates under changes in lighting, object identity, and spatial layout.

• Value Visualization:

  • The inner product $\phi(s, a)^\top \psi(l)$ was shown to correlate strongly with task progress, increasing as the robot approaches the goal and remaining low for mismatched instructions, confirming the model learns a valid goal-conditioned value function.

5. Significance

PRTS represents a paradigm shift in robotic foundation models:

• From Imitation to Reasoning: It moves beyond simple behavior cloning (imitation) to learning goal-reachability awareness, enabling the robot to understand how close it is to a goal and how feasible a path is.
• Scalability: It demonstrates that contrastive objectives, previously limited to smaller decision-making domains, can scale to large VLA backbones (4B+ parameters) and massive datasets without prohibitive cost.
• General-Purpose Intelligence: By injecting temporal and goal-aware priors during pre-training, PRTS achieves remarkable zero-shot generalization and robustness, suggesting a path toward truly general-purpose, goal-directed robotic intelligence capable of handling complex, long-horizon, and dynamic real-world environments.