DAWM: Diffusion Action World Models for Offline Reinforcement Learning via Action-Inferred Transitions

Technische Samenvatting: DAWM – Diffusie Actiewereldmodellen voor Offline Versterkend Leren via Actie-geïnferreerde Overgangen

1. Probleemstelling

Diffusie-gebaseerde wereldmodellen hebben sterke capaciteiten aangetoond in het synthetiseren van realistische trajecten met een lange horizon voor offline versterkend leren (RL). Echter, bestaande methoden staan voor een kritiek compatibiliteitsprobleem met standaard op waarde gebaseerde offline RL-algoritmen (zoals TD3BC, IQL) die vertrouwen op één-staps tijdsverschil (TD) leren.

• Beperking van actievrije modellen: Methoden zoals het Diffusie Wereldmodel (DWM) genereren toekomstige toestanden en beloningen, maar genereren geen acties expliciet. Bijgevolg kunnen ze geen volledige overgangen $(s, a, r, s')$ construeren, waardoor ze incompatibel zijn met één-staps TD-leren, dat consistentie vereist tussen toestanden, acties en beloningen voor effectieve regularisatie van de criticus.
• Beperking van gezamenlijke modellering: Benaderingen die proberen toestanden, beloningen en acties gezamenlijk te modelleren, lijden vaak onder verhoogde trainingscomplexiteit en instabiliteit, wat leidt tot verminderde praktische prestaties.
• Beperking van op planning gebaseerde methoden: Diffusie-gebaseerde planners (zoals Diffuser, Decision Diffuser) die acties herstellen via Inverse Dynamische Modellen (IDM) of gezamenlijke modellering, vertrouwen doorgaans op iteratief bemonsteren en na-terugkeuze van beloningen. Dit resulteert in hoge inferentiekosten en omzeilt de efficiëntie van één-staps TD-leren.

De kernuitdaging is het genereren van synthetische data die stabiel, efficiënt één-staps TD-leren ondersteunt, zonder de trainingsinstabiliteit van gezamenlijke modellering of de inferentiekosten van iteratieve planning.

2. Methodologie: DAWM

De auteurs stellen DAWM (Diffusie Actiewereldmodel) voor, een modulair kader dat trajectgeneratie ontkoppelt van actie-inferentie om volledige, gelabelde overgangen te produceren die geschikt zijn voor offline RL. Het kader bestaat uit twee primaire fasen:

Fase 1: Conditionele Diffusie Trajectgeneratie

• Een diffusiemodel $p_\theta$ wordt getraind om toekomstige sequenties van toestanden en beloningen te genereren, geconditioneerd op de huidige toestand $s_t$, actie $a_t$, en een doel terugkeer-naar-gaan ($R_{tg}$).
• Het model definieert de conditionele verdeling $p_\theta(\hat{s}_{t+1:t+H+1}, \hat{r}_{t:t+H} \mid s_t, a_t, R_{tg})$.
• Het trainingsdoel minimaliseert de verwachte fout tussen voorspelde en werkelijke ruis met behulp van een score-ontruisdoelstelling met classifier-vrije begeleiding.
• Cruciaal genereert dit model actievrije trajecten (alleen toestanden en beloningen), wat de generatieve taak vereenvoudigt door de complexiteit van het gezamenlijk modelleren van acties binnen het diffusieproces te vermijden.

Fase 2: Actie-aanvulling via Inverse Dynamisch Model (IDM)

• Om de actievrije trajecten om te zetten in volledige overgangen, wordt een apart getraind Inverse Dynamisch Model (IDM), $f_\phi$, ingezet.
• Het IDM reconstrueert de ontbrekende acties door de verdeling $f_\phi(a_t \mid s_{t-m:t-1}, s_t, s_{t+1})$ te modelleren als een diagonale Gaussische verdeling.
• Het IDM wordt getraind op echte offline data om de waarschijnlijkheid van waargenomen acties gegeven toestands-overgangen te maximaliseren.
• Tijdens inferentie wordt het IDM toegepast op de door diffusie gegenereerde trajecten om acties te infereren, wat volledige tuples $(s, a, r, s')$ oplevert.

Offline Beleidstraining

• De gegenereerde volledige overgangen worden gebruikt om offline RL-agenten te trainen met standaard op waarde gebaseerde algoritmen, specifiek TD3BC (dat een behavior cloning regularizer toevoegt aan TD3) en IQL (dat voordeel-gewogen supervisie-regressie uitvoert).
• De Q-functies worden bijgewerkt met Bellman-doelen berekend uit de gereconstrueerde overgangen, waardoor efficiënte één-staps TD-updates mogelijk zijn.

3. Belangrijkste Bijdragen

1. Modulair Kader: De auteurs stellen een efficiënt kader voor dat diffusie-gebaseerde generatie van toestanden en beloningen combineert met een apart getraind IDM voor actie-inferentie. Dit ontwerp vermijdt de trainingsinstabiliteit van gezamenlijke modellering terwijl het volledige overgangen produceert.
2. Compatibiliteit met Één-staps TD-leren: Door volledige $(s, a, r, s')$-overgangen te genereren, maakt DAWM het gebruik mogelijk van conservatieve offline RL-algoritmen (TD3BC, IQL) die vertrouwen op één-staps TD-leren, en biedt een computatie-efficiënt alternatief voor iteratieve diffusieplanners.
3. Empirische Prestaties: Het artikel toont aan dat conservatieve offline RL-algoritmen significant profiteren van training op deze aangevulde trajecten, en consistent betere prestaties leveren dan eerdere diffusie-gebaseerde baselines (DWM, DD) over meerdere D4RL-locomotietaken.
4. Levensvatbaarheid van Synthetische Data: Resultaten geven aan dat agenten getraind op door DAWM gegenereerde synthetische data prestaties kunnen behalen die vergelijkbaar zijn met, of zelfs beter zijn dan, agenten getraind op echte offline datasets, wat het potentieel van synthetische trajecten als praktisch alternatief benadrukt.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op 9 standaard D4RL-locomotie-benchmarks (Hopper, Walker2d, HalfCheetah) met variërende datakwaliteit (Medium, Medium-Replay, Medium-Expert).

• Vergelijking met DWM: DAWM presteerde consistent beter dan de actievrije DWM-baseline. Met TD3BC bereikte DAWM een gemiddelde verbetering van 9,3%, en met IQL 9,5%. De winsten waren bijzonder opvallend in complexe controletaken zoals Walker2d.
• Vergelijking met Op Planning Gebaseerde Methodes (DD): DAWM presteerde beter dan de Decision Diffuser (DD) wat betreft genormaliseerde terugkeer. Bovendien toonde DAWM aanzienlijk hogere efficiëntie, met ongeveer 4,0x snellere inferentie vergeleken met DD, aangezien het iteratief bemonsteren tijdens inferentie vermijdt.
• Vergelijking met Echte Data: Agenten getraind op door DAWM gegenereerde data behaalden genormaliseerde terugkeer die vergelijkbaar was met agenten getraind op echte data met standaard TD3BC/IQL, en in sommige gevallen (bijvoorbeeld Hopper-MR met TD3BC) presteerden ze beter. Hoewel SSORL (een baseline die IDM gebruikt op echte data) goed presteerde, merken de auteurs op dat SSORL profiteert van filtering voor de langste trajecten, een factor die niet aanwezig is in de vast-lengte generatie van DAWM.
• Ablatiestudies:
  • Datavolume (DAWM-T): Het verminderen van het aantal conditionerende paren met een factor 8 (DAWM-T) resulteerde slechts in bescheiden prestatiedalingen, wat suggereert dat het mechanisme van actie-aanvulling een kritieker factor is dan het pure datavolume.
  • Generatiehorizon: DAWM behield robuuste prestaties over verschillende generatiehorizons ($H=1, 3, 7$), wat aangeeft dat het kader voor één-staps TD-leren stabiel is ongeacht de roll-outlengte.

5. Betekenis en Claims

Het artikel claimt dat DAWM de kloof overbrugt tussen efficiënte generatieve modellering en op TD gebaseerde beleidsleren. De primaire betekenis ligt in:

• Aanbieden van Één-staps TD-leren: Het biedt een mechanisme om door diffusie gegenereerde data te gebruiken voor standaard, stabiele offline RL-algoritmen die eerder incompatibel waren met actievrije wereldmodellen.
• Computatie-efficiëntie: Het biedt een praktische oplossing voor grootschalig generatief offline RL door inferentiekosten te verminderen in vergelijking met iteratieve planningsmethoden.
• Praktisch Alternatief: De resultaten suggereren dat synthetische trajecten gegenereerd via deze modulaire aanpak kunnen dienen als een levensvatbaar, hoogwaardig alternatief voor echte offline datasets, wat potentieel datatekorten kan mitigeren in real-world toepassingen waar exploratie duur of onveilig is.

De auteurs concluderen dat hoewel de methode veelbelovend is, toekomstig werk zal onderzoeken uitbreiding naar visuele controle, robotica met hoge dimensionaliteit, en robuustheid onder imperfecte dynamiek.