R2-Dreamer: Redundancy-Reduced World Models without Decoders or Augmentation

1. Het Probleem

In Model-Based Reinforcement Learning (MBRL) op basis van beelden is het leren van effectieve latente representaties een fundamentele uitdaging. Bestaande methoden kampen met twee hoofdproblemen:

• Decoder-gebaseerde methoden (bijv. DreamerV3): Deze leren representaties door pixels te reconstrueren. Dit leidt tot een verspilling van representatieve capaciteit op taak-irrelevante details (zoals de achtergrond), omdat het reconstructieverlies dominant wordt door grote, statische gebieden in de afbeelding. Hierdoor worden kleine, maar cruciale objecten genegeerd.
• Decoder-vrije methoden: Deze vermijden reconstructie en gebruiken zelftoezicht (self-supervised learning), maar zijn vaak afhankelijk van Data Augmentation (DA) (zoals willekeurige verschuivingen of kleurveranderingen) als externe regularisatie om "representation collapse" (het instorten van de representatie naar een triviaal punt) te voorkomen.
  • Het nadeel: DA is taakafhankelijk. Bijvoorbeeld, een verschuiving kan een klein, belangrijk object uit het beeld snijden, en kleurverandering kan schadelijk zijn als kleur een cruciale taakfeature is. Dit beperkt de veelzijdigheid en robuustheid van deze agenten.

2. Methodologie: R2-Dreamer

R2-Dreamer is een decoder-vrij en DA-vrij MBRL-framework dat de bestaande RSSM-architectuur (Recurrent State-Space Model) van DreamerV3 behoudt, maar de leerdoelstelling fundamenteel herontwerpt.

Kerncomponenten:

• Architectuur: De decoder wordt volledig verwijderd. In plaats daarvan wordt een lichtgewicht lineaire projector (projector head) toegevoegd die de latente toestand $s_t$ (bestaande uit een deterministische $h_t$ en een stochastische $z_t$) afbeeldt op de feature-ruimte van de beeldembedding $e_t$.
• Nieuwe Leerdoelstelling (Redundancy Reduction): Het reconstructieverlies wordt vervangen door een zelftoezicht-objectief gebaseerd op Barlow Twins.
  • Het doel is om de correlatie tussen de projectie van de latente toestand ($k_t$) en de beeldembedding ($e_t$) te maximaliseren (invariantie), terwijl de redundantie tussen de verschillende dimensies van de projectie wordt geminimaliseerd.
  • De loss-functie $L_{BT}$ bestaat uit twee termen:
    1. Invariance Term: Straft afwijkingen van de diagonaalelementen van de kruiscorrelatiematrix $C$ van 1 af (zorgt dat $k_t$ en $e_t$ overeenkomen).
    2. Redundancy Term: Straft de off-diagonale elementen van $C$ af (zorgt dat de dimensies van de representatie ongecorreleerd zijn, wat leidt tot een gefactoriseerde en compacte representatie).
• Interne Regularisatie: In tegenstelling tot andere methoden die DA gebruiken om positieve paren te creëren, gebruikt R2-Dreamer de natuurlijke paren van de interne signalen van het model ($e_t$ en $k_t$) als regularisatie. Dit elimineert de behoefte aan externe augmentaties.

3. Belangrijkste Bijdragen

1. Nieuw Paradigma: Een decoder-vrij MBRL-systeem dat DA vervangt door een interne, op redundantiereductie gebaseerde regularisatie. Dit voorkomt dat taak-kritieke informatie wordt vervormd door augmentaties.
2. Efficiëntie en Prestaties: Het framework bereikt concurrerende prestaties op standaard benchmarks (DeepMind Control Suite en Meta-World) terwijl het 1,59x sneller traint dan DreamerV3 door het verwijderen van de computationally expensive reconstructiestap.
3. Robuustheid op Subtiele Taken: Het introduceert en presteert superieur op de nieuwe DMC-Subtle benchmark, waar taak-relevante objecten extreem klein zijn en makkelijk verloren gaan bij methoden die afhankelijk zijn van reconstructie of ruwe augmentatie.
4. Open Source: De release van een verenigde PyTorch-codebasis en de DMC-Subtle benchmark om toekomstig onderzoek te faciliteren.

4. Resultaten

De experimenten werden uitgevoerd op drie benchmarks: DeepMind Control Suite (DMC), Meta-World en de nieuwe DMC-Subtle.

• Standaard Benchmarks (DMC & Meta-World): R2-Dreamer presteert gemiddeld even goed als sterke baselines zoals DreamerV3 (decoder-gebaseerd) en TD-MPC2 (decoder-vrij met DA), zowel op het gemiddelde als op de mediaan van de taken.
• DMC-Subtle (Kritieke Test): Op taken waar de doelobjecten zijn verkleind (bijv. een doelwit dat 1/3 van de originele grootte is), overtreft R2-Dreamer alle baselines met een aanzienlijke marge.
  • Analyse: Saliency-maps tonen aan dat R2-Dreamer zich scherp richt op het kleine doelwit, terwijl baselines (zoals Dreamer en DreamerPro) verspreide aandacht hebben die vaak de achtergrond of irrelevante delen omvat.
• Ablatie Studies:
  • Het toevoegen van DA aan R2-Dreamer levert slechts marginale winst op en kan zelfs schadelijk zijn op de subtiele taken.
  • DreamerPro (afhankelijk van DA) crasht zonder augmentatie, wat de kritieke afhankelijkheid van externe regularisatie bij andere methoden bevestigt.
  • De methode is robuust tegen veranderingen in batchgrootte.
• Berekeningskosten: R2-Dreamer is aanzienlijk sneller (1,59x t.o.v. DreamerV3 en 2,36x t.o.v. DreamerPro) omdat het geen decoder hoeft te trainen en geen meerdere augmented views hoeft te verwerken.

5. Betekenis en Conclusie

R2-Dreamer bewijst dat een principieel intern regularisatie-objectief (reductie van redundantie) de noodzaak van beeldreconstructie en handmatige data-augmentatie in MBRL kan vervangen.

• Verschil met de staat van de kunst: Het verschuift de focus van "visuele fideliteit" (hoe goed kan ik het beeld reconstrueren?) naar "informatieve efficiëntie" (hoe goed kan ik de essentiële taakfeatures isoleren?).
• Toekomstperspectief: De methode biedt een schaalbare basis voor agenten in omgevingen met dynamische of irrelevante achtergronden, waar traditionele augmentaties vaak falen omdat ze taak-kritieke informatie kunnen verstoren. Het opent de weg voor meer universele en capabele leeragenten die niet afhankelijk zijn van taak-specifieke heuristieken voor data-augmentatie.

Kortom, R2-Dreamer biedt een efficiëntere, robuustere en theoretisch onderbouwde oplossing voor representatie-leren in beeldgebaseerde reinforcement learning.