PRTS: A Primitive Reasoning and Tasking System via Contrastive Representations

1 Introduction

• Existing VLAs emphasize static semantic understanding and immediate action prediction, leaving System 2 without quantitative goal-reachability awareness for assessing execution difficulty.The gap includes distinguishing difficult-but-reachable states from easy-but-erroneous ones.Source: §1 · pp. 1–2
• PRTS treats language instructions as goals and uses contrastive reinforcement learning to make state-action and goal embedding similarity encode discounted goal occupancy.Offline trajectory structure supplies dense goal-conditioned supervision without manual reward annotations.Source: §1 · p. 2
• A role-aware causal mask isolates CRL token streams within one VLM forward pass while preserving the behavior-cloning loss and adding negligible pretraining cost.PRTS appends <CRL_action> and <CRL_goal> token blocks and integrates the mask into a custom FlashAttention kernel.Source: §1 · p. 2
• PRTS unifies semantic reasoning, discrete action prediction, and goal-conditioned representation learning in one VLM architecture, trained on over 167B tokens of diverse robotic trajectories with language annotations.The learned reachability signal also requires task-relevant visual-language representations of functional affordances for task completion.Source: §1 · p. 2
• PRTS reports state-of-the-art performance across LIBERO, LIBERO-Pro, LIBERO-Plus, SimplerEnv, and a 14-task real-world suite, with gains in long-horizon execution, zero-shot instruction generalization, and recovery from human interventions.The real-world suite spans dual-arm RealMan and single-arm Flexiv platforms.Source: §1 · p. 3

2 Preliminaries

• In goal-conditioned reinforcement learning, an agent acts in an MDP to reach a desired goal state, with state and action spaces, transition dynamics, and an initial-state distribution defining the environment.The goal state belongs to the state space.Source: §2 · p. 3
• The discounted state occupancy measure describes the probability of visiting a state under a goal-conditioned policy, and geometric time sampling yields states from this occupancy distribution.For goal-reaching rewards, the Q-function equals the conditional occupancy probability and thus the likelihood of reaching the goal from a state-action pair.Source: §2 · p. 3
• Contrastive RL estimates the Q-function by scoring state-action and goal representations, distinguishing achievable future goals sampled from discounted occupancy from random goals sampled from the marginal data distribution.The critic uses an inner product of state-action and goal embeddings, and the contrastive objective avoids temporal-difference bootstrapping errors.Source: §2 · p. 4

3.1 Problem Setup

• PRTS replaces behavior-cloning-only pretraining with language-conditioned contrastive reinforcement learning, using each instruction as a goal and learning goal reachability from the shared VLM backbone.The formulation avoids reward labels, return discretization, pairwise progress-label curation, and auxiliary value networks.Source: §3.1 · pp. 4–5

3.2 Language-Conditioned Contrastive Reinforcement Learning

• Because one language goal is shared across trajectory timesteps, a batch contains multiple positive state-action examples for that goal, violating standard CRL’s single-positive assumption.PRTS reformulates the objective to preserve temporal structure despite this multi-positive setting.Source: §3.2 · p. 5
• Temporal weights based on discounted remaining trajectory length emulate geometric sampling, assigning exponentially larger weights to samples closer to task completion.The resulting soft targets recover the temporal structure of standard CRL.Source: §3.2 · pp. 5–6
• Bidirectional contrastive losses encode task identity in state-action-to-language matching and task progress in language-to-state-action matching.Together, the embedding inner product estimates log-discounted occupancy; sample-wise weighting alone would not encode temporal progression.Source: §3.2 · p. 6
• The temporal-weighting theorem establishes that learned representations encode log-discounted goal occupancy rather than simply the number of steps remaining.The formulation also extends to action chunks without changing its objectives or theoretical results.Source: §3.2 · p. 6

3.3 PRTS Architecture

• A role-aware causal mask separates vision/state, instruction, autoregressive action, CRL_action, and CRL_goal tokens to prevent information leakage between contrastive representations.The action representation attends to vision and proprioception, while the goal representation attends within its own block.Source: §3.3 · pp. 7–8, Figure 1 · p. 7
• Pretraining jointly optimizes bidirectional contrastive losses and next-token prediction; post-training attaches a flow-matching action expert for continuous action chunks.The backbone also receives auxiliary robot-QA and sub-task prediction losses, while the action expert is absent during pretraining.Source: §3.3 · pp. 8–9
• The learned embedding geometry supplies dense goal-conditioned action values from the same backbone used for action generation, without a separate value network.The representations align observations, robot states, and instructions by task progress, which the authors associate with generalization to long-horizon and fine-grained manipulation.Source: §3.3 · p. 9

4 Datasets

• The 404 million-sample, 167.8 billion-token pretraining corpus has two complementary domains: action-labeled data for cross-embodiment execution priors and visual-reasoning data for grounding and planning.The visual-reasoning domain also provides affordance-perception supervision.Source: §4 · p. 9, Figure 2 · p. 10

4.1 Action-Labeled Data

• Action-labeled trajectories combine AgiBotWorld, RoboMind, Open X-Embodiment, and PRTS self-collected data across diverse tasks, embodiments, and complex manipulation settings.The sources include cross-embodiment demonstrations and contact-rich, multiscene scenarios.Source: §4.1 · p. 9

4.2 Visual-Reasoning Data

• The visual-reasoning corpus includes spatial-grounding and affordance datasets for precise localization, spatial relations, referential disambiguation, and operable-object-part perception.Examples include pixel-coordinate prediction and distinguishing an object's functional region from other parts.Source: §4.2 · p. 10
• For planning and high-level reasoning, the corpus adds physical-causality reasoning, instruction decomposition into actionable sub-goals, and temporal reasoning over embodied video.These data sources target long-horizon execution and context awareness.Source: §4.2 · p. 11
• A curated LLaVA-Instruct subset is included to retain general vision-language competence, zero-shot language following, visual world knowledge, and conversational fluency.Source: §4.2 · p. 11

Visual-Language-Action Models

• VLA pretraining commonly uses next-token prediction over discretized actions, while auxiliary VQA and reasoning objectives add static grounding without encoding trajectory-level goal progress.The paper identifies this temporal representation gap as the bottleneck PRTS addresses.Source: §5 · p. 11
• Value-augmented approaches such as π0.6 and VLAC use separate value networks, preventing value-learning gradients from directly shaping the shared VLM representation.PRTS instead integrates contrastive reinforcement learning into VLM pretraining.Source: §5 · p. 11

Contrastive Reinforcement Learning

• Contrastive RL formulates value learning as InfoNCE classification without Bellman regression, and its cross-entropy objective aligns with VLM token prediction.The paper presents this classification-based formulation as a natural fit for jointly learning language reasoning and goal-aware planning.Source: §5 · pp. 11–12

6.1 Experimental Setup

• The evaluation spans four simulation benchmarks and 14 real-world tasks, with LIBERO-finetuned checkpoints tested zero-shot on LIBERO-Plus and LIBERO-Pro.The setup describes perturbations in object layout, viewpoint, robot state, language, lighting, background, and sensor noise across the benchmark suite.Source: §6 · p. 12, §6.1 · p. 13, §6.2 · p. 14

6.2 Overall Performance

• The experiments test standard LIBERO execution, zero-shot robustness and generalization, WidowX cross-embodiment performance, and real-world manipulation across two robot platforms.The real-world suite has 11 dual-arm RealMan tasks and three single-arm Flexiv tasks; simulations use benchmark-specific success-rate protocols.Source: §6.1 · pp. 12–14, §6.2 · p. 14

6.2.1 LIBERO

• 98.4% average SR on LIBERO matches the strongest reported result with batch size 32 and 30K steps, while matched-budget ABot-M0 reaches 97.9%.PRTS reaches 96.6% on Long; matched-budget ABot-M0 reaches 94.8% there.Source: §6.2.1 · p. 14, Table 2 · p. 15

6.2.2 LIBERO-Plus

• 81.4% average SR on zero-shot LIBERO-Plus exceeds π0.5’s 80.7% after the checkpoint is post-trained only on standard LIBERO.The largest stated gains are Robot (+14.4) and Background (+15.3).Source: §6.2.2 · p. 14, Table 3 · p. 15

6.2.3 LIBERO-Pro

• 58.8% average SR on zero-shot LIBERO-Pro improves over π0.5 by 5.5 points, with PRTS reaching 31.5% on Task versus π0.5’s 0.8%.The Task axis gives a novel instruction to manipulate a different target object; Position swaps the relative locations of pick and place objects.Source: §6.2.3 · p. 16, Table 4 · p. 16

6.2.4 SimplerEnv with Visual Matching

• 77.1% average SR on SimplerEnv improves over GR00T-N1.5 by 15.2 points under the same batch size and gradient-step budget.PRTS reaches 100% on Put Eggplant in Yellow Basket and leads on two other listed tasks.Source: §6.2.4 · p. 16, Table 5 · p. 17

6.2.5 Real-World Evaluation

• 95.9% average SR on the 11-task RealMan suite and 90.0% average success on Flexiv’s three tasks show strong real-world results; PRTS also leads on fine-grained and long-horizon tasks.RealMan evaluation uses one shared checkpoint across tasks; Flexiv uses per-task fine-tuning.Source: §6.2.5 · pp. 16–17, Figure 4 · p. 17

6.3 Effect of Contrastive RL

• With the same architecture, data, and post-training budget, CRL raises LIBERO average SR from 97.8% to 98.4% and improves zero-shot average SR by 4.9% on LIBERO-Plus and 5.0% on LIBERO-Pro.On LIBERO-Pro, Object-Position rises from 4.8% to 36.0%; Spatial-Task rises from 34.4% to 62.2%, and Goal-Task from 20.2% to 33.8%.Source: §6.3 · pp. 18–19, Table 6 · p. 18, Table 7 · p. 19

6.4 Controlled Real-World Generalization

• Paper Rubbish and Place Block remain at 100% success across illumination, spatial, and object shifts, while Pick Shoes and Stack Cups stay at least 80%.The four axes range from lighting and layout changes to same-category object replacements and recombined language goals.Source: §6.4 · p. 20, Figure 6 · p. 21, Figure 7 · p. 22
• 73.8% success on recombined task instructions exceeds π0.5’s 35.0% and π0’s 13.8%, though instruction-level recombination remains the hardest real-world shift.PRTS reaches 80% on the mouse-grasp task, 55% on Place Block, and 85% on Pick Shoes; Stack Cups reaches 75%.Source: §6.4 · pp. 20, 23

6.5 Robustness and Recovery under Human Interventions

• Across roughly five intervention trials per task, PRTS recovered and completed tasks by tracking changed states, revisiting invalidated subgoals, and redoing affected stages.The tests are qualitative stress tests, not a fully instrumented benchmark; π0 largely replays fixed phases, while π0.5 can fail to re-enter earlier subgoals.Source: §6.5 · p. 24

6.6 CRL Value Visualization

• 0.57 was the CRL value near shoebox closure, rising from 0.28 initially; the mismatched instruction remained at 0.04–0.20 throughout the same trajectory.Correct-instruction value peaks align with task milestones, while the wrong-instruction curve never crosses it, without post-training adaptation.Source: §6.6 · pp. 24–25, Figure 9 · p. 25

6.7 Pre-training Efficiency

• 1.18× the FA3 reference was the fused CuTe kernel’s per-layer attention time while optimizing the full CRL objective.The role-aware fused kernel keeps the added attention cost close to BC training, unlike the general FlexAttention path.Source: §6.7 · p. 26, Figure 10 · p. 26
• 85.1% scaling efficiency at 64 H100 GPUs followed almost 100% through 8 GPUs and 95.2% at 32 GPUs.The throughput comparison uses a linear reference anchored at 2 GPUs; the larger-scale drop is attributed to ZeRO-2 communication overhead.Source: §6.7 · p. 26, Figure 10 · p. 26

7 Conclusion

• PRTS scales reward-label-free contrastive reinforcement learning into VLA pretraining, deriving goal-reachability supervision from offline trajectories without manual rewards or a separate value network.A role-aware causal mask fused into FlashAttention, sequence packing, and sharded contrastive similarity computation fold the contrastive head into the behavior-cloning forward pass at near-BC throughput.Source: §7 · p. 26
• Across LIBERO, LIBERO-Plus, LIBERO-Pro, SimplerEnv, RealMan, and Flexiv, PRTS matches or exceeds strong prior VLAs at substantially smaller post-training compute.The reported gap widens on novel instructions, long-horizon execution, and human-intervention recovery.Source: §7 · pp. 26–27
• The pre-trained backbone produces goal-reachability-aware and goal-discriminative values on out-of-distribution rollouts without post-training adaptation.The conclusion frames these results as extending the scaling story for goal-conditioned RL to VLA pretraining and as a step toward goal-directed robotic intelligence.Source: §7 · p. 27

A Contributions

• Jiangyuan Zhao and Yang Zhang led data collection and standardization, with annotation and supplemental data by the same two contributors.Source: A Contributions · p. 34
• Yang Zhang led model architecture and algorithm; Yang Zhang and Jiangyuan Zhao handled pretraining, while seven contributors worked on training infrastructure.Source: A Contributions · p. 34
• Yang Zhang, Jiangyuan Zhao, and Chenyou Fan handled post-training and evaluation and writing; Zhang led the project, and Chenjia Bai advised.Source: A Contributions · p. 34

B Proof of Theorem 1

• The language-to-state-action objective defines, for each goal in a batch, the subset of state-action samples from trajectories assigned that language task.The batch contains samples with state-action pairs, timestep, and trajectory length; the objective is written with explicit expectation notation.Source: B Proof of Theorem 1 · p. 34
• Cross-entropy minimization matches predicted probabilities to temporal soft labels, yielding a representation inner product tied to remaining trajectory length.Taking expectation over batch sampling makes the batch-dependent constant converge to a goal-dependent function.Source: B Proof of Theorem 1 · pp. 34–35
• Under deterministic expert demonstrations, the goal is reached exactly at step T, linking the representation score to the log goal-conditioned action value.The resulting representation inner product targets log discounted goal occupancy.Source: B Proof of Theorem 1 · p. 35
• Unlike standard CRL’s geometric sampling of future-state goals, PRTS weights all batch positives and supports discrete, fixed language goals.The temporal weighting reaches the same optimal solution; using multiple positives reduces variance relative to a single geometric sample.Source: B Proof of Theorem 1 · p. 35

• The appendix details real-world task instructions, scene arrangements, and required behaviors; platform configurations and trial-level protocol are described in Section 6.1.Source: C Detailed real-world task… · p. 35

• Paper Rubbish requires one gripper to hold a trash can while the other drops crumpled paper inside; Place Block puts a block at a specified tea-tray location.Source: §C.1 · pp. 35–36
• Eleven RealMan tasks span ten short-horizon desktop manipulations and one long-horizon office task, with representative rollouts shown in Figure 11.Source: §C.1 · p. 35, Figure 11 · p. 36
• Serve Tea and Tea One place cups on designated surfaces, with Tea One also returning the gripper to an initial-like pose; Push Button tasks press a target using the named gripper.Source: §C.1 · p. 36
• Flip Tennis Tube and Hand Over require transferring an object between grippers before placement at a designated tray target.Source: §C.1 · p. 36
• Stack Cups alternates grippers to build a four-cup vertical stack, while Pick Shoes places two shoes in a shoebox before the other gripper closes it.Source: §C.1 · p. 36
• Office(LongTerm) combines mouse placement, vertical pen insertion, power-strip and switch presses, and book placement in a multi-object sequence.Source: §C.1 · p. 36

• The Flexiv suite comprises three single-arm tasks requiring accurate geometric placement and contact-rich manipulation.Source: §C.2 · p. 37
• Block Combination aligns a four-hole block with the posts of a fixture.Source: §C.2 · p. 37
• Gear Assembly mounts small, medium, and large gears on corresponding fixture posts while aligning their center holes.Source: §C.2 · p. 37
• Flower Arrangement retrieves blue and red flowers by their stems and places them in a glass receptacle.Source: §C.2 · p. 37

• LIBERO-Pro's Position and Task variations test generalization under shifts that conventional VLAs handle poorly, while PRTS shows superior instruction-following generalization.Source: Table 10 · p. 38