Ego-Foresight: Self-supervised Learning of Agent-Aware Representations for Improved RL

Probleemstelling

Deep Reinforcement Learning (RL) heeft aanzienlijke vooruitgang geboekt, maar het vereist nog steeds enorme hoeveelheden trainingsdata om effectieve beleidsregels te leren, zowel in gesimuleerde als in echte omgevingen. Bestaande methoden die proberen de efficiëntie te verbeteren door de agent en de omgeving apart te modelleren, zijn vaak afhankelijk van toezicht (supervisie). Dit betekent dat ze bijvoorbeeld een masker nodig hebben van de robot in de scène (verkregen via geometrische ID's, CAD-modellen of gesegmenteerde modellen) om de agent te onderscheiden van de omgeving.

In de echte wereld is het verkrijgen van dergelijke supervisiecomplex en onflexibel; het werkt niet goed als de robotvorm verandert (bijvoorbeeld door het gebruik van gereedschappen). Mensen daarentegen leren nieuwe vaardigheden met weinig trials en zonder expliciete supervisie, vaak door motorische voorspelling: het ontwikkelen van een intern model van de eigen bewegingen en de daaruit voortvloeiende zintuiglijke input. Het paper stelt de vraag of deze menselijke capaciteit voor zelfbewustzijn en voorspelling in kunstmatige systemen kan worden nagebootst zonder toezicht.

Methodologie: Ego-Foresight (EF)

De auteurs introduceren Ego-Foresight (EF), een zelftoezichtende (self-supervised) methode om agent-informatie te ontkoppelen (disentangle) van omgevingsinformatie, gebaseerd op beweging en voorspelling.

Kernarchitectuur:

1. Encoder-Decoder met Recurrente Eenheid: Het systeem gebruikt een convolutie-encoder die een reeks RGB-frames (context) en geplande acties als input ontvangt.
2. Ontkoppeling van Features: De encoder produceert een feature-vector $h$, die wordt opgesplitst in twee delen:
   • $h_s$: Scène-features (omgevingsinformatie).
   • $h_a$: Agent-features (informatie over de robot).
   • Er wordt een "bottleneck" toegepast op $h_a$ (kleine dimensie) om de recurrente eenheid te dwingen zich te focussen op de meest voorspelbare dynamiek: die van de agent zelf.
3. Voorspellingsmechanisme:
   • De recurrente eenheid neemt de huidige agent-features ($h_a$) en de volgende actie ($a_{t+1}$) als input.
   • Het voorspelt de toekomstige configuratie van de agent ($\hat{h}_a$) op een tijdstip $t_k$ binnen de voorspellingshorizon.
   • De voorspelde agent-features worden gecombineerd met de gefixeerde scène-features ($h_s$) uit het begin van de context.
   • Een decoder reconstrueert de toekomstige frame ($\hat{x}_{t_k}$) op basis van deze combinatie.
4. Verliesfunctie: Het doel is om het verschil te minimaliseren tussen de gereconstrueerde frame en het werkelijke toekomstige frame ($L_{ef} = ||\hat{x}_{t_k} - x_{t_k}||^2$). Omdat de scène-informatie "vaststaat" (fast-forwarded), moet de recurrente eenheid alleen de beweging van de agent voorspellen om de reconstructie correct te maken.

Integratie met RL:
EF wordt gebruikt als een auxiliary task (hulptaak) voor bestaande RL-algoritmen (specifiek DrQ-v2 en TD-MPC2). De totale loss-functie is een gewogen som van de RL-loss (bijv. critic loss) en de EF-reconstructieloss:
$$L_{total} = L_{RL} + \beta L_{ef}$$
Dit regulariseert de feature-learning, waardoor de encoder leert features die niet alleen waardevol zijn voor de beloning, maar ook voorspelbaar zijn in termen van de agent's eigen beweging. Om te voorkomen dat de agent tijdens het trainen alleen doelgerichte bewegingen maakt (wat de diversiteit aan bewegingen beperkt), wordt er een motor-babbling-fase ingevoerd aan het begin van het trainen, waarbij willekeurige acties worden uitgevoerd om de visuo-motorische mapping te leren.

Belangrijkste Bijdragen

1. Zelftoezichtende Ontkoppeling: Een methode om agent-informatie te ontkoppelen zonder supervisie (zoals masks of CAD-modellen), puur gebaseerd op de voorspelbaarheid van de eigen beweging.
2. Aanpassingsvermogen: Het bewijs dat het systeem zich kan aanpassen aan veranderende "body-schemas", zoals het opnemen van gereedschappen (bijv. een hamer) die dan als onderdeel van de agent worden voorspeld.
3. Verbeterde RL-prestaties: De uitbreiding van state-of-the-art algoritmen (DrQ-v2 en TD-MPC2) met EF, wat leidt tot betere sample-efficiency en prestaties.
4. Ablatiestudies: Een grondige analyse van hyperparameters (zoals de gewicht $\beta$, de voorspellingshorizon en de dimensie van de agent-features).

Resultaten

De methode werd getest op twee uitdagende benchmarks: Meta-World (robotische objectmanipulatie, inclusief taken met gereedschappen) en DMC (locomotie en manipulatie).

• Sample-Efficiency: In 21 van de 26 geteste taken verbeterde DrQv2-EF de prestaties van de originele DrQ-v2, vaak met een aanzienlijke reductie in het aantal benodigde stappen om een taak op te lossen.
• Prestaties bij Gereedschappen: Het grootste voordeel werd gezien in taken die gereedschappen vereisen. Omdat de beweging van een vastgehouden gereedschap voorspelbaar is vanuit de robotbeweging, leert het model dit effectief als onderdeel van de agent, wat de prestaties ten opzichte van baselines (zoals SEAR, CURL en Dreamer-v3) verbetert.
• Model-Based RL: Ook TD-MPC2-EF toonde verbeteringen (8 van de 10 taken beter dan de originele TD-MPC2), wat aantoont dat de methode niet beperkt is tot model-free RL.
• Efficiency Normalized Score (ENS): De auteurs definiëren een nieuwe metric (ENS) die sample-efficiency en eindprestatie combineert. EF behaalde hierin hoge scores, vaak concurrerend met of beter dan gesuperviseerde methoden.
• Visualisatie: Visualisaties van de feature-vectors tonen aan dat de agent-features ($h_a$) zich specialiseren in de robotbeweging, terwijl de scène-features ($h_s$) de statische en veranderende omgeving vastleggen.

Significantie en Conclusie

Het paper biedt een belangrijke stap naar RL in de echte wereld door de afhankelijkheid van dure en starre supervisie (masks/CAD) te elimineren. Door de agent en omgeving te ontkoppelen via zelfvoorspelling, kunnen RL-algoritmen sneller leren en zich beter aanpassen aan veranderingen in de robotvorm (zoals het gebruik van tools).

De auteurs benadrukken dat hoewel pixel-wise reconstructie computatiekosten met zich meebrengt, de winst in sample-efficiency (minder interacties met de omgeving nodig) in robotica-toepassingen vaak belangrijker is dan pure rekensnelheid. De methode is generiek toepasbaar op verschillende RL-algoritmen en opent de weg voor robuustere, adaptieve robotcontrole in onvoorspelbare omgevingen.