Improving Zero-Shot Offline RL via Behavioral Task Sampling

1. Probleemstelling

Offline Zero-Shot Reinforcement Learning (ZSRL) heeft tot doel agenten te trainen die onbekende beloningsfuncties kunnen optimaliseren zonder verdere interactie met de omgeving, uitsluitend vertrouwend op een vooraf verzameld offline dataset.

• Standaardbenadering: Huidige state-of-the-art methoden (bijv. Successor Features, Forward-Backward representaties) leren een state-inbedding $\phi(s)$ en trainen een beleid dat is geconditioneerd op een taakvector $z$. De beloningsfunctie wordt lineair gedefinieerd als $R_z(s) = \phi(s)^\top z$.
• De Tekortkoming: Om deze beleidslijnen te trainen, bemonsteren bestaande algoritmen taakvectoren $z$ doorgaans uniform vanuit een hoogdimensionale eenheidshypersfeer ($S^{d-1}$).
• Het Kernprobleem: De auteurs betogen dat deze uniforme bemonstering een mismatch creëert tussen de geometrie van de taakruimte en de fysica van de omgeving. In hoogdimensionale ruimtes veroorzaakt het fenomeen van "concentratie van maat" dat uniform bemonsterde vectoren bijna orthogonaal zijn ten opzichte van de feitelijke gedragingen die de agent kan produceren (de gedragsruimte). Dit leidt tot "signaalverdunning", waarbij het beloningssignaal voor bijna al het gedrag verwaarloosbaar klein wordt, waardoor het voor de agent onmogelijk wordt om optimale acties te onderscheiden van suboptimale acties. Dit resulteert in een slechte zero-shot generalisatie.

2. Methodologie: Behavioral Task Distribution (BTD)

De auteurs stellen voor om uniforme taakbemonstering te vervangen door een datagedreven aanpak genaamd Behavioral Task Sampling.

A. Theoretisch Inzicht

• Feature Occupancy: Zij definiëren de feature occupancy $\psi_\pi$ van een beleid als de afgetopte som van state-features langs zijn traject.
• Verdwijnende Variantie: Zij bewijzen (Propositie 4.1) dat naarmate de latente dimensie $d$ toeneemt, de verwachte variantie van opbrengsten over beleidslijnen onder uniforme taakbemonstering convergeert naar nul. Dit betekent dat het leersignaal instort, waardoor effectieve beleidsoptimalisatie wordt verhinderd.

B. Het Voorgestelde Algorithmus

In plaats van $z$ te bemonsteren uit een uniforme verdeling, extrahert de methode taakvectoren direct uit het offline dataset:

1. Extrahering: Uit het offline dataset $D$ extraheren de auteurs sub-trajecten $\tau$. Voor elk sub-traject berekenen zij de empirische feature occupancy $\tilde{\psi}_\tau = \sum \gamma^t \phi(s_t)$.
2. Definitie Taakvector: De taakvector voor een traject wordt gedefinieerd als de genormaliseerde occupancy: $z_\tau = \tilde{\psi}_\tau / \|\tilde{\psi}_\tau\|_2$. Deze vector vertegenwoordigt de richting in de feature-ruimte die de beloning maximaliseert voor dat specifieke waargenomen gedrag.
3. Distributiemodellering: De verzameling van alle geëxtraheerde vectoren vormt een empirische verdeling $p_{data}$. De auteurs passen een parametrisch dichtheidsmodel (specifiek een Gaussian Mixture Model of GMM) aan op deze data, waardoor de Behavioral Task Distribution (BTD) ontstaat, aangeduid als $p_\theta(z)$.
4. Training: Tijdens de beleidstraining worden taakvectoren bemonsterd uit de geleerde BTD ($z \sim p_\theta(z)$) in plaats van de eenheidshypersfeer. Dit zorgt ervoor dat de agent wordt getraind op taken die fysiek haalbaar zijn en geworteld zijn in de dynamica van de omgeving.

3. Belangrijkste Bijdragen

1. Identificatie van Signaalverdunning: Het paper toont theoretisch en empirisch aan dat uniforme taakbemonstering in hoogdimensionale ruimtes leidt tot een instorting van het leersignaal als gevolg van orthogonaliteit met de gedragsruimte.
2. Behavioral Task Distribution (BTD): Een nieuw, methode-onafhankelijk raamwerk dat de verdeling van haalbare taken direct uit offline data extrahert en modelleert.
3. Integratie: De methode is een drop-in vervanging voor de stap van taakbemonstering in bestaande ZSRL-algoritmen (Successor Features en Forward-Backward) zonder de onderliggende representatieleerdoelstellingen te wijzigen.
4. Uitgebreide Evaluatie: Uitgebreide experimenten over meerdere benchmarks, methoden voor representatieleren en latente dimensies.

4. Experimentele Resultaten

De auteurs evalueerden hun methode op standaard ZSRL-benchmarks (Cheetah, Walker, Quadruped) met behulp van datasets die zijn verzameld via Random Network Distillation (RND). Zij vergeleken met zeven baselines (waaronder Autoencoders, Transition Models, LRA en BYOL) over verschillende latente dimensies.

• Prestatiewinst: BTD-bemonstering verbeterde de zero-shot prestaties met gemiddeld 20% over alle omgevingen en baselines.
• Robuustheid ten opzichte van Dimensie:
  • Baseline Falen: Naarmate de latente dimensie toenam (tot 1000), leden standaard uniforme bemonsteringsmethodes aan een scherpe daling in prestaties, vaak instortend tot bijna-willekeurige prestaties.
  • BTD Stabiliteit: BTD-bemonstering behield hoge prestaties zelfs bij hoge dimensies (bijv. 1000), wat bewijst dat het het probleem van signaalverdunning effectief mitigeert.
• Verminderde Variantie: BTD-bemonstering verminderde de variantie in testprestaties over verschillende methoden voor representatieleren aanzienlijk, wat wijst op consistenter en betrouwbaarder leren.
• Ablatiestudies:
  • Mengstrategieën: Het mengen van uniforme bemonstering met BTD-bemonstering (interpoleren tussen de twee) verslechterde de prestaties, wat bevestigt dat het introduceren van willekeurige, onhaalbare taken het leersignaal schaadt.
  • GMM-componenten: De prestaties platte af na ongeveer 20 GMM-componenten, wat suggereert dat de gedrags-taakruimte kan worden vastgelegd met een relatief eenvoudige verdeling.
  • Visualisatie: UMAP-projecties toonden aan dat uniforme bemonstering de hele hypersfeer bestrijkt, terwijl de feitelijke gedrags-taken een klein, specifiek manifold bezetten, wat de noodzaak van datagedreven bemonstering bevestigt.

5. Betekenis en Conclusie

Dit werk verschuift fundamenteel het paradigma van offline ZSRL van willekeurige taakgeneratie naar gedragsverankering.

• Praktische Impact: Het benadrukt dat in offline settings welke taken worden gekozen voor training even kritisch is als hoe het beleid ze leert.
• Schaalbaarheid: Door de taakverdeling te ontkoppelen van de vloek van dimensionaliteit die inherent is aan uniforme bemonstering, maakt de methode het mogelijk om generalist-agenten te trainen in hoogdimensionale latente ruimtes.
• Generaliseerbaarheid: De aanpak is compatibel met diverse technieken voor representatieleren (SF, FB, Autoencoders, etc.), waardoor het een schaalbare weg biedt naar robuustere en generaliseerbaardere reinforcement learning-agenten voor implementatie in de echte wereld, waar taken fysiek beperkt zijn.

Kortom, het paper betoogt dat om echte zero-shot generalisatie in offline RL te bereiken, agenten moeten worden getraind op taken die de feitelijke capaciteiten en dynamica van de omgeving weerspiegelen, in plaats van een abstracte, uniforme verdeling van wiskundige mogelijkheden.