Technische Samenvatting: In-Context Reinforcement Learning via Communicative World Models (CORAL)

1. Probleemstelling

Reinforcement Learning (RL) agenten hebben vaak moeite om te generaliseren naar nieuwe taken en contexten zonder hun parameters bij te werken. Deze beperking komt voort uit geleerde representaties en beleidsregels die overfit zijn op de specificiteiten van trainingsomgevingen. Hoewel In-Context Reinforcement Learning (ICRL) een pad biedt naar contextgestuurde adaptatie door beleidsregels te conditioneren op eerdere interacties, leunen bestaande benaderingen vaak op zelfgesuperviseerde transformers die een begrip missen van onderliggende taakdynamiek. Gevolgelijk is hun generalisatie fragiel en sterk afhankelijk van de kwaliteit van offline datasets.

Daartegenover bieden World Models (WMs) agenten een gestructureerd begrip van omgevingsdynamiek en beloningsfeedback, wat de voorspelling van toekomstige toestanden mogelijk maakt. Echter, in standaard WMs zijn de geleerde representaties vaak verstrengeld met taakspecifieke beleidslering. Deze verstrengeling leidt tot representaties die te gespecialiseerd zijn, minder overdraagbaar, en lijden aan een mismatch tussen het doel van representatieleer en beleidsoptimalisatie.

Het kernprobleem dat wordt geadresseerd is hoe men het dynamische begrip van World Models kan benutten voor ICRL zonder de verstrengeling van beleidslering, om zo een overdraagbare, communicatieve context te creëren die adaptatie in nieuwe omgevingen versnelt.

2. Methodologie: Het CORAL-framework

De auteurs stellen CORAL (Communicative Representation for Adaptive RL) voor, een framework dat ICRL formuleert als een twee-agent emergent communicatieprobleem. CORAL scheidt functioneel de leerprocessen van de latente representatie en de controle, waardoor de rollen van wereldmodellering en beloningsmaximalisatie worden ontkoppeld.

2.1 Agent Architectuur

Het systeem bestaat uit twee verschillende agenten die op elk tijdstip $t$ met elkaar interageren:

• Information Agent (IA): Geparametriseerd door $\phi$ (een Transformer-gebaseerde architectuur), fungeert de IA als een world model. Het verwerkt een context van recente ingebedde observaties en streeft naar een begrip van de omgevingsdynamiek en taakbeloningen. Cruciaal is dat de IA niet wordt getraind op directe opbrengstmaximalisatie. In plaats daarvan comprimeert de IA zijn begrip in beknopte latente berichten ($m_t$).
• Control Agent (CA): Geparametriseerd door $\theta$ (doorgaans een PPO-beleid), ontvangt de CA de observatie $o_t$ en het bericht $m_t$ van de IA. Het enige doel van de CA is het maximaliseren van de taakbeloningen door acties $a_t$ te selecteren op basis van de communicatieve context.

2.2 Trainingsdoelen

De agenten worden gezamenlijk voorgetraind op een diverse distributie van taken ($\mathcal{T}$) met behulp van een samengestelde verliesfunctie voor de IA, terwijl de CA wordt geoptimaliseerd via Proximal Policy Optimization (PPO).

De verliesfunctie van de Information Agent ($\mathcal{L}(\phi)$):
De IA wordt getraind via een zelfgesuperviseerde samengestelde verliesfunctie bestaande uit drie componenten:

1. Dynamics Awareness Loss ($\mathcal{L}_{Dyn}$): Verankert de representaties van de IA door te vereisen dat het de volgende observatie ($\hat{o}_{t+1}$), beloning ($\hat{r}_{t+1}$) en terminatiekans ($\hat{d}_{t+1}$) voorspelt op basis van het huidige bericht en de actie van de CA. Dit zorgt ervoor dat de interne staat fysieke waarheden over de omgeving vastlegt.
2. Temporal Coherence Loss ($\mathcal{L}_{Coh}$): Stimuleert consistentie in de communicatie door de IA te trainen om het volgende bericht ($\hat{m}_{t+1}$) te voorspellen vanuit het huidige bericht ($m_t$) en de actie. Dit bevordert berichten als compacte representaties van de toestand.
3. Causal Influence Loss ($\mathcal{L}_{Causal}$): Een nieuw doel dat het communicatieprotocol vormgeeft. Het meet de Instantaneous Causal Effect (ICE), gedefinieerd als de Reverse KL Divergentie tussen het beleid van de CA met het bericht versus een zero-vector bericht. Deze verliesfunctie wordt gewogen door een hybride nutssignaal (dat Generalized Advantage Estimates en directe waarde-veranderingen combineert) om de IA aan te sporen berichten te sturen die de CA aanzetten tot voordelige, beslissende verschuivingen in zijn beleid.

Het doel van de Control Agent:
De CA maximaliseert de verwachte gedisconteerde opbrengst $E[\sum \gamma^t r_t]$ met behulp van standaard PPO, geconditioneerd op de observatie en het bericht van de IA.

2.3 Trainingsregime

Om generalisatie te waarborgen, maakt het voortrainingsregime gebruik van multi-environment training. De agenten worden getraind op een diverse, discrete set taken die gemeenschappelijke entiteiten en dynamiek delen maar verschillende oplossingen vereisen. Dit wordt geïmplementeerd met behulp van vectorized training (via JAX) om parallelle omgevingen te draaien, wat catastrofaal vergeten voorkomt en de IA dwingt om abstracte dynamiek te leren in plaats van specifieke instanties te memoriseren.

2.4 Implementatie (Deployment)

Tijdens de implementatie worden de gewichten van de IA bevroren. Het dient als een vaste "contextualizer" of communicatieve prior voor een nieuwe, willekeurig geïnitialiseerde CA. De CA leert ongeziene taken op te lossen door de berichten te interpreteren die door de bevroren IA worden verstrekt, wat snelle in-context adaptatie mogelijk maakt zonder de wereldmodellen bij te werken.

3. Belangrijkste Bijdragen

1. Formulering: Het artikel modelleert single-agent ICRL als een twee-agent emergent communicatieprobleem waarbij een World Model (IA) leert om dynamiek te communiceren naar een Control Agent (CA).
2. Hybride Trainingsstrategie: Het framework ontkoppelt het leren van representatie en controle. De IA hanteert zelfgesuperviseerde training met drie specifieke koppen (dynamiek, coherentie, invloed), terwijl de CA gebruikmaakt van reinforcement pre-training. Deze scheiding voorkeert de mismatch in doelstellingen die wordt gevonden in monolithische WMs.
3. Causal Influence Loss: De introductie van een nieuwe verliesterm die de IA expliciet stimuleert om berichten te genereren die het beleid van de CA causaal beïnvloeden richting hoog-nuttige acties, wat garandeert dat de communicatie pragmatisch en effectief is.
4. Empirische Validatie: Uit uitgebreide experimenten blijkt dat CORAL de sample-efficiëntie in online sparse-reward omgevingen verbetert en superieure zero-shot generalisatie bereikt vergeleken met standaard PPO, World Models en state-of-the-art In-Context RL methoden (bijv. Decision Transformer, Agentic Transformer).

4. Experimentele Resultaten

De auteurs evalueerden CORAL over drie regimes:

• Online Adaptatie (MiniGrid Suite):

  • CORAL vertoonde aanzienlijk hogere sample-efficiëntie dan baselines (PPO, PPO-Transformer, World Model).
  • In taken zoals DoorKey 8x8 en Crossings S11N5 bereikte CORAL bijna optimale beloningen ongeveer twee keer zo snel als PPO.
  • In moeilijkere, grotere varianten (bijv. DynObs 16x16) slaagde CORAL erin de CA te ondersteunen bij het oplossen van taken waar baselines faalden of nabij random presteerden.
  • Time-to-Threshold (TTT): CORAL bereikte 90% van de maximale prestaties aanzienlijk sneller (bijv. 2.3x sneller dan PPO in DoorKey 8x8).

• Zero-Shot Generalisatie:

  • Wanneer geëvalueerd op ongeziene, complexere omgevingen zonder verdere training, presteerden CORAL-agenten die waren voorgetraind op eenvoudigere taken beter dan PPO, World Models en Decision Transformers.
  • Dit geeft aan dat de communicatieve prior beter generaliseert naar dynamische verschuivingen dan louter sequentie-gebaseerde klonering van demonstraties.

• Offline Generalisatie (D4RL):

  • In volledig offline settings behaalde CORAL een totale gemiddelde genormaliseerde score van 85.47, waarmee het Decision Transformer (77.49) versloeg en TD3+BC (84.83) evenaarde.
  • Opvallend genoeg excelleerde CORAL op moeilijke Medium-Replay datasets (bijv. 98.33 op Hopper-M-Replay), wat suggereert dat de communicatieve prior effectief ruis uit suboptimale demonstraties filtert.

• Ablatie-studies:

  • Het verwijderen van de Temporal Coherence Loss leidde tot hoge variantie en trager leren.
  • Het verwijderen van de Causal Influence Loss verminderde de sample-efficiëntie en de uiteindelijke prestaties, wat bevestigt dat een passief wereldmodel onvoldoende is; de IA moet worden aangemoedigd om het beleid van de CA te beïnvloeden.
  • Analyse van de Instantaneous Causal Effect (ICE) toonde aan dat het protocol sterke invloed biedt tijdens het leren (hoge onzekerheid) en stabiliseert naarmate het beleid convergeert, waardoor het fungeert als een leer-regularisator.

5. Betekenis en Claims

Het artikel claimt dat CORAL een principieel middenpad biedt tussen contextgestuurde adaptatie en world modeling. Door het leren van dynamiek (via de IA) te scheiden van de controle (via de CA), leert het framework een overdraagbare communicatieve representatie die niet overfit is op specifieke taakbeloningen.

De auteurs stellen dat deze aanpak de effectiviteit van het leren van een overdraagbare communicatieve prior valideert. De resultaten suggereren dat emergent communicatie niet alleen nuttig is voor multi-agent coördinatie, maar ook dient als een krachtige algemene methode voor snelle in-context adaptatie in single-agent RL. Het framework stelt agenten in staat om significante winsten te behalen in zowel sample-efficiëntie als succesvolle zero-shot adaptatie in diverse online en offline omgevingen, waarmee het de fragiliteit van huidige ICRL-benaderingen en de verstrengelingsproblemen van standaard World Models aanpakt.