PRTS: A Primitive Reasoning and Tasking System via Contrastive Representations

1. Probleemstelling

Huidige Vision-Language-Action (VLA) modellen hebben robotbesturing verbeterd door Vision-Language Models (VLM's) te gebruiken als perceptuele ruggengraat, doorgaans met een dual-systeem architectuur (Systeem 2 voor redenering, Systeem 1 voor besturing). Echter, bestaande benaderingen lijden aan een fundamentele beperking:

• Statisch vs. Temporaal: Ze frameren pre-training voornamelijk als supervised behavior cloning (BC), met focus op statische visueel-linguïstische semantische correlaties. Ze missen een intrinsiek begrip van doelbereikbaarheid—de kwantitatieve waarschijnlijkheid dat een huidige staat-actie-paar over tijd zal leiden tot een taak gespecificeerd in taal.
• Knelpunten in Waarde-estimatie: Pogingen om waardelearning te integreren vertrouwen vaak op aparte waardenetwerken, vereisen handmatige beloningsannotaties, of gebruiken Temporal-Difference (TD) learning wat lijdt aan bootstrapping-fouten en hoge variantie bij taken met een lange horizon.
• Generalisatiekloof: Zonder expliciet temporaal bewustzijn worstelen deze modellen met planning op lange horizon, zero-shot nieuwe instructies, en herstel na menselijke ingrepen, omdat ze niet kunnen onderscheiden tussen "makkelijke-maar-foutieve" staten en "moeilijke-maar-bereikbare" staten.

2. Methodologie: PRTS

De auteurs stellen PRTS voor, een VLA fundamenteel model dat pre-training herformuleert via Goal-Conditioned Reinforcement Learning (GCRL) met gebruik van Contrastive Reinforcement Learning (CRL).

A. Kernconcept: Contrastive RL voor VLA

In plaats van direct acties te voorspellen of een apart waardenetwerk te gebruiken, leert PRTS een uniforme embedding-ruimte waar het inwendig product van een staat-actie embedding ($\phi(s, a)$) en een doel embedding ($\psi(l)$) de log-gedisconteerde doelbezettingskans benadert.

• Theoretische Basis: Gebaseerd op de equivalentie tussen de Q-functie voor doel-bereikingsbeloningen en de conditionele bezettingskans (Lemma 1).
• Doel: Het model leert te discrimineren tussen staat-actie paren die een specifiek taakdoel kunnen bereiken (positieve samples) en die welke dat niet kunnen (negatieve samples).

B. Belangrijkste Technische Innovaties

1. Taal-geconditioneerde Contrastive RL:

   • In tegenstelling tot standaard CRL dat toekomstige staten geometrisch samplet, behandelt PRTS de taalinstructie als het vaste doel voor een hele traject.
   • Multi-Positief Probleem: Aangezien een enkele taakinstructie geldt voor alle tijdstappen in een traject, bevat een batch veel positieve paren voor één doel.
   • Temporele Weging: Om de geometrische sampleverdeling van standaard CRL te herstellen, introduceren de auteurs een temporele wegingsschema. Staat-actie paren dichter bij het einde van een traject (dichter bij het doel) worden zwaarder gewogen ($\gamma^{T-t}$). Dit stelt het model in staat de temporele structuur van doelbereiking te leren zonder expliciete beloningslabels.
   • Bidirectionele Doelen:
     • Staat-Actie $\to$ Taal: Koppelt $(s, a)$ aan zijn specifieke instructie.
     • Taal $\to$ Staat-Actie: Dwingt het model om staat-acties binnen een traject te rangschikken op basis van hun nabijheid tot het doel, waardoor temporele vooruitgang wordt gecodeerd.

2. Single-Forward-Pass Architectuur (Verwaarloosbare Overhead):

   • PRTS voegt twee auxiliaire token-blokken toe aan de invoersequentie: <CRL_action> en <CRL_goal>.
   • Rol-bewuste Causale Maskering: Een aangepast attentiemasker (geïmplementeerd via een gefuseerde CuTe-FlashAttention kernel) zorgt voor informatiescheiding:
     • <CRL_action> tokens attenderen alleen op visie/proprioceptie (niet op taal).
     • <CRL_goal> tokens attenderen alleen op de instructie.
     • Standaard autoregressieve actietokens behouden standaard causale attentie.
   • Resultaat: Het model produceert discrete actie-logits, staat-actie embeddings en doel-embeddings in een enkele forward pass, met verwaarloosbare rekenkundige overhead vergeleken met pure behavior cloning.

3. Twee-fasen Trainingspijplijn:

   • Pre-training: Optimaliseert gezamenlijk de VLM-ruggengraat op behavior cloning (next-token prediction), auxiliaire taken (Robot-QA) en de bidirectionele contrastieve verliezen. Er is nog geen actie-expert aangesloten.
   • Post-training: Een Flow-Matching Action Expert (gebaseerd op Diffusion Transformer) wordt aangesloten en fijn afgestemd op downstream data. De voorgeöefende ruggengraat levert de priors voor doelbereikbaarheid.

3. Belangrijkste Bijdragen

1. Herformulering van VLA Pre-training: Eerste toepassing om beloningsvrije, taal-geconditioneerde Contrastive RL direct te schalen naar de VLA pre-fase, waardoor de noodzaak voor handmatige beloningsannotaties of aparte waardenetwerken wordt geëlimineerd.
2. Intrinsiek Bewustzijn van Doelbereikbaarheid: De geleerde representaties coderen kwantitatief de waarschijnlijkheid van het bereiken van een doel, waardoor de kloof tussen semantisch redeneren en temporele taakvoortgang wordt overbrugd.
3. Efficiënte Architectuur: De "Rol-bewuste Causale Maskering" en de aangepaste CuTe-kernel maken dichte supervisie voor doelbereikbaarheid mogelijk met bijna nul rekenkundige overhead, wat schaalbaar maakt naar grote modellen (4B parameters) en enorme datasets (167B tokens).
4. Unificerend Kader: Unificeert semantisch redeneren, discrete actiepredictie en waarde-estimatie binnen een enkele VLM-architectuur.

4. Experimentele Resultaten

PRTS werd voorgeöefend op 167,8 miljard tokens uit diverse datasets (AgiBotWorld, RoboMind, Open X-Embodiment, etc.) en geëvalueerd op simulatie- en real-world benchmarks.

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

  • SOTA Prestaties: PRTS bereikt state-of-the-art resultaten, die sterke baselines (bijv. $\pi^0.5$, ABot-M0) evenaren of overtreffen, terwijl het aanzienlijk minder post-training rekenkracht gebruikt (bijv. 1/8e van het budget van $\pi^0.5$ op LIBERO).
  • Taken met Lange Horizon: Aanzienlijke winst op de "Long" suite van LIBERO (96,6% vs. 94,8% voor baselines met aangepast budget), wat de effectiviteit van temporeel bewustzijn bewijst.
  • Robuustheid: Overtreft baselines in zero-shot generalisatie onder distributieveranderingen (camera, robotstatus, verlichting, objectindeling) en Taak/Positie verstoringen in LIBERO-Pro.
  • Zero-Shot Instructie: Enorme verbetering op nieuwe instructies (31,5% vs. 0,8% voor $\pi^0.5$ op de "Taak" as van LIBERO-Pro).

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

  • Complexe Manipulatie: Bereikte een 95,9% gemiddeld slagingspercentage op 11 dual-arm taken en 90,0% op 3 fijnmazige single-arm taken.
  • Herstel na Menselijke Ingrijpen: PRTS toonde superieure robuustheid wanneer mensen taken onderbraken (bijv. objecten verwijderen tijdens uitvoering). In tegenstelling tot baselines die faalden of verouderde trajecten afspeelden, evalueerde PRTS de staat succesvol opnieuw en plande het opnieuw om het doel te bereiken.
  • Generalisatie: Hield hoge slagingspercentages aan onder veranderingen in verlichting, objectidentiteit en ruimtelijke indeling.

• Visualisatie van Waarde:

  • Het inwendig product $\phi(s, a)^\top \psi(l)$ bleek sterk te correleren met taakvoortgang, toenemend naarmate de robot het doel nadert en laag blijvend voor niet-overeenkomende instructies, wat bevestigt dat het model een geldige doel-geconditioneerde waardenfunctie leert.

5. Betekenis

PRTS vertegenwoordigt een paradigma-verschuiving in robotische fundamentele modellen:

• Van Imitatie naar Redenering: Het gaat verder dan simpele behavior cloning (imitatie) naar het leren van doelbereikbaarheidsbewustzijn, waardoor de robot begrijpt hoe dicht het bij een doel is en hoe haalbaar een pad is.
• Schaalbaarheid: Het toont aan dat contrastieve doelen, die eerder beperkt waren tot kleinere besluitvormingsdomeinen, kunnen schalen naar grote VLM-ruggengraten (4B+ parameters) en enorme datasets zonder prohibitieve kosten.
• Algemeen Toepasbare Intelligentie: Door temporele en doel-bewuste priors tijdens pre-training in te brengen, bereikt PRTS opmerkelijke zero-shot generalisatie en robuustheid, wat wijst op een weg naar echt algemeen toepasbare, doelgerichte robotische intelligentie die complexe, lange-horizon en dynamische real-world omgevingen aankan.