Technische Samenvatting: TaskSense

Probleemstelling

Wereldmodellen voor visuele controle vertrouwen doorgaans op het reconstrueren van volledige observaties om compacte latente toestanden te leren. Hoewel deze aanpak effectief is in gestandaardiseerde benchmarks met statische achtergronden, stimuleert het impliciet latente representaties om informatie over de gehele visuele input te behouden, ongeacht de relevantie daarvan voor de controletaak. In omgevingen waar de taakrelevante inhoud slechts een klein deel van de observatie beslaat, of waar achtergrondruis en afleiders aanwezig zijn, verbruikt deze reconstructie-objectief kostbare representatieve capaciteit voor taakonrelevante visuele inhoud. Deze mismatch tussen visuele reconstructie en controle-objectieven vertekent de latente dynamiek, waardoor de leer-signalen voor controle-relevante kenmerken worden verdund en de prestaties onder visuele afleidingen ernstig verslechteren. Omgekeerd biedt het uitsluitend vertrouwen op taak-supervisie (bijv. beloningsvoorspelling) vaak een te zwak signaal om rijke visuele representaties te leren. Bestaande oplossingen die reconstructie volledig elimineren, offeren de generatieve capaciteiten op die nodig zijn voor verbeelding-gebaseerde planning, terwijl methoden die aandacht of masking toepassen na de extractie van kenmerken, er niet in slagen te controleren welke informatie de latente dynamica-modellen überhaupt binnenkomt.

Methodologie

De auteurs stellen TaskSense voor, een taak-centraal wereldmodel-framework dat taakrelevantie afdwingt voordat de latente codering plaatsvindt. TaskSense breidt de DreamerV3-architectuur uit met een differentieerbaar stochastisch ruimtelijk aandachtmechanisme dat geconditioneerd is op de vorige latente geloofstoestand (belief state).

Kerncomponenten

1. Latent-geconditioneerde Stochastische Ruimtelijke Aandacht:

   • In plaats van aandacht voor elke pixel te voorspellen, definieert het model aandacht over een grove rooster $\Lambda$ die de observatie dekt.
   • Een lichtgewicht MLP voorspelt aandacht-logits op basis van de vorige latente toestand $s_{t-1}$ (specifiek de verborgen toestand $h_{t-1}$ en de stochastische latente toestand $z_{t-1}$), waardoor het model kan beslissen waar het zich op moet concentreren op basis van zijn evoluerende begrip van de taak in plaats van op de instantane visuele verschijning.
   • Het aandachtveld wordt gemodelleerd als een collectie onafhankelijke stochastische variabelen met behulp van een Binary Concrete (Relaxed Bernoulli) distributie. Dit maakt het mogelijk om meerdere regio's tegelijkertijd te bezoeken of te negeren, waardoor de competitieve beperkingen van softmax-aandacht worden vermeden.
   • Tijdens de training worden differentieerbare masks gesampled; tijdens de inferentie wordt het deterministische gemiddelde gebruikt.

2. Aandacht-bewuste Codering (Attention-Aware Encoding):

   • De gesamplede aandacht-mask wordt opgehoogd (upsampled) en gebruikt om de observatie te moduleren via elementgewijze vermenigvuldiging, waardoor regio's met lage aandacht worden gedempt.
   • Om de ambiguïteit tussen pixels met lage intensiteit en onderdrukte pixels met hoge intensiteit op te lossen, wordt de opgehoogde aandacht-kaart als een extra inputkanaal aan de encoder toegevoegd.

3. Auxiliaire Inverse-Dynamics Objectief:

   • Om te voorkomen dat de sparsity-regularisator de aandacht naar een triviale oplossing drijft (het negeren van alle observaties), wordt een auxiliair objectief geïntroduceerd.
   • Het model voorspelt de actie $a_{t-1}$ die werd uitgevoerd tussen twee opeenvolgende tijdstappen met behulp van het tijdsverschil (temporal difference) van de gecodeerde kenmerken ($f_t - f_{t-1}$).
   • Dit objectief fundeert het aandachtmechanisme in controle-relevante regio's door de visuele kenmerken te behouden die voorspellend zijn voor het gedrag van de agent.

4. Aandacht-geconditioneerde Reconstructie:

   • In plaats van de volledige observatie te reconstrueren, reconstrueert de decoder alleen de beïnvloede regio's ($o^{att}_t$).
   • De decoder is geconditioneerd op de gesamplede aandacht-kaart (via stop-gradient) om ervoor te zorgen dat het weet welke regio's zijn behouden, wat voorkomt dat de latente toestand wordt bestraft voor het missen van informatie die opzettelijk is weggefilterd.

Trainingsdoelstelling

De totale verliesfunctie combineert de standaard DreamerV3 latente dynamica-verliezen ($L_{dyn}$) met de nieuwe componenten:
$$L_{WM} = L_{dyn} + L_{rec} + \lambda_{sparse}L_{sparse} + \lambda_{inv}L_{inv}$$
Waarbij $L_{rec}$ de reconstructieverlies is over de gefilterde observaties, $L_{sparse}$ compacte aandacht-kaarten aanmoedigt, en $L_{inv}$ de inverse-dynamics verlies is.

Belangrijkste Bijdragen

1. TaskSense Framework: Een op reconstructie gebaseerd wereldmodel-framework dat taakrelevantie afdwingt vóór de latente codering via latent-geconditioneerde stochastische ruimtelijke aandacht.
2. Zelfgesuperviseerd Leermechanisme: Een nieuwe aanpak die stochastische aandacht fundeert in controle-relevante regio's door inverse-dynamics supervisie te combineren met reconstructie over gefilterde observaties.
3. Empirische Validatie: Demonstratie dat TaskSense concurrerende prestaties behoudt op standaard benchmarks, terwijl het aanzienlijk beter presteert dan baselines onder visuele afleidingen.

Experimentele Resultaten

De auteurs hebben TaskSense geëvalueerd op de DeepMind Control Suite (DMC) en de Distracting Control Suite (DCS), die dynamische natuurlijke beeldachtergronden bevat.

• Standaard Prestaties: Op de standaard DMC (statische achtergronden) evenaart TaskSense de prestaties van de DreamerV3 baseline over de meeste taken, wat aantoont dat het reconstrueren van slechts een subset van de observatie voldoende is voor effectieve latente planning.
• Robuustheid tegen Afleidingen: Op de Distracting Control Suite presteert TaskSense consistent beter dan DreamerV3. Bijvoorbeeld, op de Walker Run taak behaalde TaskSense een gemiddelde return van 531.2 vergeleken met 286.2 voor DreamerV3. Dit bevestigt dat het filteren van irrelevante visuele inhoud de robuustheid verbetert.
• Kwalitatieve Analyse: Visualisaties van de geleerde aandacht-kaarten op de Cheetah Run taak laten zien dat TaskSense consequent focust op controle-relevante regio's (bijv. de poten van de cheetah) terwijl de romp en de afleidende achtergronden worden onderdrukt.
• Ablatie-studies:
  • Het verwijderen van het inverse-dynamics objectief leidt tot diffuse aandacht die er niet in slaagt om controle-relevante regio's te lokaliseren, wat een significante prestatieverslechtering veroorzaakt.
  • Het vervangen van stochastische aandacht door deterministische aandacht resulteert in een significante prestatiedaling, wat het belang benadrukt van het blootstellen van het model aan diverse gedeeltelijke observaties tijdens de training.
  • Het verwijderen van aandacht-geconditioneerde reconstructie veroorzaakt een kleinere maar merkbare degradatie, wat wijst op de noodzaak om de reconstructiedoelstelling af te stemmen op de gefilterde input.

Betekenis en Claims

Het artikel beweert dat TaskSense een fundamentele mismatch in huidige wereldmodellen aanpakt: de neiging om alle visuele informatie te bewaren in plaats van te focussen op wat belangrijk is voor controle. Door de focus te verleggen naar observatie-niveau filtering in plaats van post-hoc representatie-verfijning, stelt TaskSense wereldmodellen in staat om hun representatieve capaciteit efficiënt toe te wijzen aan controle-relevante informatie.

De auteurs stellen dat deze aanpak wereldmodellen in staat stelt om de generatieve voordelen van reconstructie te behouden (essentieel voor verbeelding-gebaseerde planning) terwijl ze robuustheid bereiken tegen visuele afleidingen die standaardmodellen doorgaans degraderen. Het werk suggereert dat het afdwingen van taakrelevantie bij de observatie-interface een veelbelovende richting is voor het verbeteren van visuele wereldmodellen en representatieleer in model-gebaseerde reinforcement learning.

Het artikel erkent een beperking: het gaat ervan uit dat visuele kenmerken die voorspellend zijn voor directe acties (gevangen door inverse dynamics) synoniem zijn aan de kenmerken die relevant zijn voor langetermijnsucces van de taak. In scenario's waar deze uiteenlopen, kan het huidige supervisiesignaal de aandacht op suboptimale kenmerken richten, een richting die is geïdentificeerd voor toekomstig werk.