Distributional value gradients for stochastic environments

Deze paper introduceert Distributional Sobolev Training, een methode die de sample-efficiëntie in stochastische omgevingen verbetert door niet alleen de verdeling van waarden, maar ook die van hun gradiënten te modelleren via een wereldmodel en een gecontracteerde Bellman-operator.

De kern

De Kunst van het Voorspellen in een Onvoorspelbare Wereld

Stel je voor dat je een beginnend skateboarder bent die probeert een nieuwe truc te leren. Je wilt niet alleen weten of je de truc lukt (dat is de beloning), maar je wilt ook precies voelen hoe je je gewicht moet verplaatsen om het te doen (dat is de 'gradiënt' of de afgeleide).

In de wereld van kunstmatige intelligentie (AI) proberen robots dit ook. Ze leren door te proberen en fouten te maken. Maar er is een groot probleem: de wereld is chaotisch.

Het Probleem: De Rijdende Bus

Stel je voor dat je een robot leert om een bus te besturen.

• De oude methode (Deterministisch): De robot leert: "Als ik het stuur 5 graden naar links draai, gebeurt er X." Dit werkt prima als de weg perfect glad is. Maar als er een windvlaag komt, of een andere auto plotseling remt (stochastische omgeving), faalt de robot. Hij denkt dat zijn voorspelling perfect was, maar de realiteit was anders.
• Het nieuwe probleem: De robot probeert ook te leren hoe hij het stuur moet bewegen om beter te worden (de gradiënt). Maar als de wereld chaotisch is, is deze 'richting' ook wazig. Het is alsof je probeert een lijn te trekken op een trillend schip. De lijn wordt onzeker en de robot raakt in de war.

De auteurs van dit paper zeggen: "Waarom proberen we niet de onzeekerheid zelf te leren?"

De Oplossing: De "Waarschijnlijkheids-Compass"

Deze onderzoekers hebben een nieuwe manier bedacht om AI te trainen, genaamd Distributional Sobolev Training. Laten we dit opsplitsen in drie simpele onderdelen:

1. Niet één voorspelling, maar een hele voorspellingstabel
In plaats van dat de AI zegt: "Ik ga 10 punten scoren", zegt hij: "Ik heb 50% kans op 8 punten, 30% kans op 12 punten, en 20% kans op 5 punten."

• Vergelijking: Een oude AI is als een weerman die zegt: "Morgen is het 20 graden." Een nieuwe AI is als een weerman die zegt: "Morgen is het 20 graden, maar het kan ook 15 of 25 worden, en het kan regenen." Dit helpt de AI om beter voorbereid te zijn op verrassingen.

2. Het leren van de 'richting' in het onzekere
Dit is het slimme deel. De AI leert niet alleen de uitkomsten, maar ook de richting waarin hij moet bewegen om die uitkomsten te verbeteren, zelfs als die richting onzeker is.

• Vergelijking: Stel je voor dat je in een mistig bos loopt.
  • De oude methode zegt: "Loop rechtuit." (Als er een boom staat, val je er tegenaan).
  • De nieuwe methode zegt: "Er is een kans dat er links een boom staat, dus loop een beetje naar rechts, maar houd rekening met de kans dat rechts een moeras is." De AI leert de verdeling van de beste paden, niet alleen één pad.

3. De "Spiegel" (De Wereldmodel)
Omdat de echte wereld (zoals een fysieke robot of een video-game) niet altijd perfect te voorspellen is, bouwen de onderzoekers een virtuele spiegel van de wereld. Dit is een AI die de fysica van de wereld nabootst (een 'cVAE').

• Vergelijking: Het is alsof de robot eerst in een virtueel trainingscentrum oefent waar hij duizenden keren kan vallen zonder pijn te doen. Hij leert daar hoe de wind en de grond werken, en past die kennis toe in de echte wereld.

Waarom is dit zo belangrijk?

De onderzoekers hebben bewezen dat hun methode wiskundig stabiel is. Ze hebben een nieuwe "rekenregel" (een Bellman-operator) bedacht die garandeert dat de AI niet blijft rondlopen in cirkels, maar echt leert, zelfs als de omgeving erg onvoorspelbaar is.

Ze hebben dit getest op twee manieren:

1. Een simpel spelletje: Waar de AI moest zoeken naar een verborgen schat in een doolhof met veel mogelijke locaties. De nieuwe methode vond de schat veel sneller en was minder snel in de war door de vele opties.
2. Complexe robot-simulaties (MuJoCo): Hier moesten robots lopen en rennen. Toen ze de robots "dronken" maakten (door ruis en onzekerheid toe te voegen), bleef de nieuwe methode stabiel lopen, terwijl de oude methoden struikelden en vielen.

Samenvatting in één zin

Deze paper introduceert een slimme manier om robots te leren niet alleen wat er gaat gebeuren, maar ook hoe onzeker dat is en hoe ze zich moeten aanpassen in een chaotische wereld, waardoor ze veel robuuster en slimmer worden dan hun voorgangers.

Het is alsof je een student leert niet alleen de antwoorden op een toets te kennen, maar ook hoe hij moet redeneren als de toetsvragen veranderen of als er een storing in het licht is.

Technische samenvatting

Probleemstelling

Reinforcement Learning (RL) in continue actie-ruimtes maakt vaak gebruik van actor-critic algoritmen waarbij de critic (waardefunctie) actie-gradiënten levert voor het optimaliseren van het beleid (policy). Bestaande methoden die gebruikmaken van gradiëntinformatie, zoals MAGE (Model-based Action-Gradient Estimator), verbeteren de sample-efficiency door een wereldmodel te leren en daar doorheen te backpropageren.

Echter, deze methoden hebben twee fundamentele beperkingen in stochastische of ruizige omgevingen:

1. Deterministische aanname: Ze behandelen actie-gradiënten als deterministische waarden. In onzekere omgevingen zijn deze gradiënten echter zelf stochastisch, wat leidt tot een verlies van sample-efficiency en instabiliteit.
2. Onvoldoende onzekerheidsmodellering: Traditionele distributionele RL (DRL) modelleert de verdeling van de opbrengst (return), maar negeert de verdeling van de gradiënten van die opbrengst. Dit is problematisch omdat de onzekerheid in de omgeving ook de onzekerheid in de gradiënten beïnvloedt, wat cruciaal is voor robuust beleidsoptimalisatie.

Methodologie: Distributional Sobolev Training

De auteurs introduceren een nieuw raamwerk genaamd Distributional Sobolev Reinforcement Learning, met als kern het Distributional Sobolev Deterministic Policy Gradient (DSDPG) algoritme.

1. Sobolev Bellman Operator:
In plaats van alleen de verwachte opbrengst $Q(s,a)$ te leren, modelleert de methode de gezamenlijke verdeling van de opbrengst én de actie-gradiënt ($\nabla_a Q$). Dit wordt gedefinieerd als een "Random Action Sobolev Return":
$$Z^{Sa}(s, a) = \left[ \sum \gamma^t r_t; \nabla_a \sum \gamma^t r_t \right]$$
De auteurs leiden een nieuwe Sobolev Bellman operator af die zowel de verdeling van de opbrengst als die van de gradiënt bootstrapt. Dit vereist dat de operator de transformatie van de gradiënt onder de overgangsdynamica en het beleid correct doorgeeft (via de kettingregel).

2. Generatief Wereldmodel (cVAE):
Omdat echte omgevingen vaak niet differentieerbaar zijn, leren de auteurs een differentieerbaar wereldmodel om de overgangs- en beloningsdynamica te benaderen. Ze gebruiken een Conditionele Variational Autoencoder (cVAE).

• De cVAE leert de verdeling $P(s', r | s, a)$.
• Het model maakt gebruik van de reparameterization trick, waardoor het mogelijk is om stochastische steekproeven te trekken die differentieerbaar zijn ten opzichte van de invoer $(s, a)$. Dit is essentieel om de gradiënten van de opbrengst te kunnen berekenen via backpropagation.

3. Afstandsmaat en Contractie:
Om de verdelingen te vergelijken en te trainen, gebruiken ze de Maximum Mean Discrepancy (MMD). Omdat de standaard MMD niet altijd contractief is in deze context, introduceren ze de Max-Sliced MMD (MSMMD).

• MSMMD projecteert de multivariate verdelingen op 1D-assen en maximaliseert de MMD over alle mogelijke projectierichtingen.
• De auteurs bewijzen theoretisch dat de Sobolev Bellman operator een contractie is onder MSMMD (en ook onder de Wasserstein-metriek), mits aan bepaalde gladheidsvoorwaarden (Lipschitz-continuïteit van het beleid en de dynamica) wordt voldaan. Dit garandeert een uniek vast punt (fixed point).

4. Implementatie Details:

• Over-schatting Bias: Om het bekende probleem van over-schatting in RL aan te pakken, gebruiken ze een ensemble van twee critics en passen ze truncatie toe (TQC-methode), waarbij de hoogste $p\%$ van de geschatte waarden in de doelverdeling wordt verwijderd.
• Training: Het algoritme traint de critic en het wereldmodel gezamenlijk, waarbij de loss functie de MSMMD afstand minimaliseert tussen de voorspelde Sobolev-verdeling en de bootstrap-doelverdeling.

Belangrijkste Bijdragen

1. Distributional Sobolev Training: Een nieuw paradigma dat distributionele RL uitbreidt naar het modelleren van zowel waarden als hun gradiënten, specifiek ontworpen voor stochastische omgevingen.
2. Theoretische Grondslagen: De eerste contractiebewijzen voor een gradiënt-bewust RL-raamwerk. De auteurs tonen aan dat de Sobolev Bellman operator contractief is onder MSMMD, wat een fundamentele trade-off blootlegt tussen de gladheid van de omgeving (Jacobian-bounds) en de effectieve horizon ($\gamma$).
3. Robuustheid in Stochastische Omgevingen: Het introduceren van een generatief wereldmodel (cVAE) dat differentieerbaar is, waardoor het mogelijk wordt om gradiënten te schatten in niet-differentieerbare omgevingen zonder de stochastische aard van de gradiënten te negeren.
4. Empirische Validatie: Uitgebreide experimenten die aantonen dat de methode superieur is aan bestaande gradiënt-gebaseerde methoden (zoals MAGE) in omgevingen met hoge onzekerheid.

Resultaten

De methode werd getest op een synthetisch "toy" probleem en op zes standaard MuJoCo-omgevingen (zoals Ant-v2, Humanoid-v2, Walker2d-v2).

• Toy Environment: In een omgeving met meerdere modale uitkomsten (multimodal returns) presteerde DSDPG (met MSMMD) consistent beter dan deterministische Sobolev-methoden en andere baselines. Het bleek robuust te zijn tegen toenemende multimodaliteit.
• MuJoCo Benchmarks:
  • In ruisvrije omgevingen presteerde DSDPG gelijkwaardig aan de state-of-the-art baselines (TD3, IQN, MAGE).
  • Onder multiplicatieve waarnemingsruis en additieve Gaussische dynamische ruis overtrof DSDPG alle concurrenten significant, vooral in complexe, hoog-dimensionele taken zoals Ant-v2 en Humanoid-v2.
  • Deterministische Sobolev-methoden (MAGE) leden aan sterke prestatiedalingen en hogere variantie onder ruis, wat aantoont dat het modelleren van de gradiëntverdeling cruciaal is voor stabiliteit.
• Ablatie Studies: Experimenten bevestigden dat de correctie voor over-schatting (TQC) en de keuze van het wereldmodel (cVAE vs. Normalizing Flows) essentieel zijn voor de stabiliteit en prestaties.

Significantie

Dit paper is een belangrijke stap in de evolutie van model-based reinforcement learning voor continue controle. Het lost een fundamentele beperking op van eerdere gradiënt-gebaseerde methoden: hun onvermogen om om te gaan met de inherent stochastische aard van gradiënten in onzekere omgevingen.

Door Distributional Sobolev Training te combineren met MSMMD, bieden de auteurs een wiskundig onderbouwde en empirisch bewezen oplossing die:

• De stabiliteit van RL-algoritmen in ruige, realistische omgevingen verbetert.
• Een brug slaat tussen differentieerbaar programmeren, distributionele RL en Sobolev-training.
• Toont dat het expliciet modelleren van onzekerheid in zowel waarden als hun afgeleiden leidt tot robuustere agenten, wat essentieel is voor toepassingen in de echte wereld waar onzekerheid de norm is.

De code en implementatie in JAX zijn openbaar beschikbaar, wat de reproduceerbaarheid en verdere ontwikkeling van dit raamwerk faciliteert.