Technische Samenvatting: Distributionally Robust Adaptive Task Sampling (DRATS)

Probleemstelling

Multi-task Reinforcement Learning (MTRL) heeft tot doel een enkele agent te trainen om prestaties over meerdere taken tegelijkertijd te optimaliseren, waarbij kennisoverdracht wordt benut om de data-efficiëntie te verbeteren. Echter, standaard MTRL-benaderingen lijden vaak aan ongebalanceerd leren: agenten lossen gemakkelijke taken snel op terwijl ze slechts langzame vooruitgang boeken op moeilijkere taken.

Vorig werk heeft deze onbalans grotendeels toegeschreven aan conflicterende taakgradiënten, met oplossingen die gradiëntmanipulatie (bijv. gradiëntchirurgie) of gespecialiseerde netwerkarchitecturen omvatten. Dit artikel identificeert een onderscheidende, onderbelichte oorzaak van onbalans: ongebalanceerde data-allokatie. Standaard MTRL kent een gelijke interactiebegroting toe aan alle taken. Aangezien moeilijke taken aanzienlijk meer ervaring vereisen om op te lossen dan gemakkelijke taken, resulteert uniforme toewijzing in een overtoewijzing van data aan gemakkelijke taken (die weinig extra data nodig hebben) en een undertoewijzing van data aan moeilijke taken. Hoewel Curriculum Learning (CL)-methodes steekproefverdelingen aanpassen, prioriteren ze doorgaans eerst gemakkelijke taken om overdracht te faciliteren, wat de data-ongebalanceerdheid kan verergeren in situaties waar positieve overdracht zwak is of taken conflicteren.

Methodologie: Distributionally Robust Adaptive Task Sampling (DRATS)

De auteurs stellen DRATS voor, een algoritme dat adaptief prioriteit geeft aan het steekproeven van taken die het verst verwijderd zijn van opgelost. De methode is afgeleid via de volgende stappen:

1. Formalisering als een Haalbaarheidsprobleem

In plaats van de gemiddelde opbrengst over taken te maximaliseren, wordt MTRL geformaliseerd als een haalbaarheidsprobleem: het vinden van een beleid $\theta$ zodat de opbrengst $J_i(\theta)$ in elke taak $i$ een referentiedoel $J^{ref}_i$ bereikt.
Het doel wordt getransformeerd tot het minimaliseren van de opbrengstgaping $g_i(\theta) = \max(J^{ref}_i - J_i(\theta), 0)$. Om gevallen te behandelen waarin beperkingen niet gelijktijdig kunnen worden voldaan (vanwege capaciteitsbeperkingen of conflicten), wordt het doel het minimaliseren van de maximale opbrengstgaping over alle taken (een minimax-doel):
$$\min_{\theta} \max_{i \in [k]} g_i(\theta)$$

2. Minimax-Optimalisatie en Distributionele Robuustheid

De innerlijke maximalisatie wordt herschreven als een maximalisatie over een taakverdeling $q \in \Delta_k$ (het waarschijnlijkheidssimplex). Het direct oplossen hiervan leidt tot een niet-gladde, hebzuchtige oplossing die alle waarschijnlijkheid concentreert op de enkelste slechtst presterende taak, wat instabiliteit en data-inefficiëntie veroorzaakt.

Om dit aan te pakken, hanteert DRATS Distributionally Robust Optimization (DRO) door de taakverdeling $q$ te beperken tot een Kullback-Leibler (KL)-bol met straal $\epsilon$ gecentreerd rond een basisverdeling $p_0$ (uniform). Dit levert de beperkte optimalisatie op:
$$\min_{\theta} \max_{q \in \mathcal{Q}} \mathbb{E}_{i \sim q}[g_i(\theta)]$$
waarbij $\mathcal{Q} = \{q \in \Delta_k \mid \text{KL}(q \| p_0) \leq \epsilon\}$.

3. Gesloten-vorm Oplossing en Algoritme

Met behulp van Lagrange-relaxatie wordt de optimale taakverdeling $q^*$ in gesloten vorm afgeleid als een softmax van de opbrengstgapingen:
$$q^*_i = \frac{\exp(\eta g_i(\theta))}{\sum_j \exp(\eta g_j(\theta))}$$
Hierbij fungeert $\eta$ als een inverse temperatuurparameter die de scherpte van de prioritering regelt.

DRATS Algoritme (Algoritme 1):

1. Datacollectie: Steekproeven trekken van taken uit de huidige verdeling $q_t$ en trajecten verzamelen.
2. Gap-schatting: Monte Carlo-schattingen van opbrengsten $\hat{J}_i$ berekenen en referentie-opbrengsten $J^{ref}_i$ updaten (online geschat als de maximaal waargenomen opbrengst).
3. Gap-normalisatie: Genormaliseerde opbrengstgapingen $\hat{g}_i$ berekenen.
4. Verdeling-updates: $q_t$ updaten via mirror ascent (geëxponentieerde gradiënt) op het doel $L(\theta, q) = \mathbb{E}_{i \sim q}[g_i(\theta)] - \frac{1}{\eta}\text{KL}(q \| p_0)$.
5. Projectie: Een minimale steekproefkans $\epsilon$ afdwingen om ervoor te zorgen dat alle taken voldoende worden gestreept voor stabiele gap-schatting.
6. Beleid-update: Het beleid $\theta$ updaten met een on-policy algoritme (bijv. PPO) om de gewogen opbrengst te maximaliseren.

4. Convergentiegaranties

De auteurs bewijzen dat DRATS onder standaard aannames (inclusief een laag-regret garantie voor de beleid-update) convergeert naar de optimale worst-case opbrengstgaping met een snelheid van $O(1/\sqrt{T})$, waarbij $T$ het aantal iteraties is.

Belangrijkste Bijdragen

1. DRATS Algoritme: Introductie van een multi-task RL-algoritme dat taken adaptief prioriteert op basis van de opbrengstgaping, afgeleid uit een minimax-formulering van MTRL.
2. Theoretische Convergentie: Bewijs dat DRATS convergeert tot binnen $\epsilon$ fout van de optimale worst-case taak-opbrengstgaping met een standaard $O(1/\sqrt{T})$ snelheid.
3. Empirische Validatie: Aantonen dat DRATS hogere data-efficiëntie, hogere aggregaatopbrengst en aanzienlijk hogere worst-case taakopbrengst bereikt dan bestaande curriculum learning-baselines (bijv. Uniform, Learning Progress, Learning Potential, Hard First).
4. Orthogonaliteit: Aantonen dat DRATS compatibel is met bestaande multi-task netwerkarchitecturen (bijv. Soft Modularization, MOORE) en de prestaties verbetert, ongeacht of taken positieve overdracht vertonen of conflicteren.

Experimentele Resultaten

De auteurs evalueerden DRATS op vier benchmarks: een 4-taken Gridworld, MuJoCo6 (6 continue controle-taken), MetaWorld-MT10 en MetaWorld-MT50.

• Aggregaat Prestatie: DRATS behaalde de hoogste aggregaatopbrengst over alle benchmarks.
• Worst-Task Prestatie: De verbetering werd voornamelijk gedreven door superieure prestaties op de moeilijkste taken. Bijvoorbeeld, in MuJoCo6 was DRATS de enige methode die succesvol de moeilijke Ant en Humanoid taken oploste, terwijl baselines geen vooruitgang boekten op Ant.
• Data-efficiëntie: DRATS bereikte de prestatieniveaus van baselines met aanzienlijk minder interacties. In MT10 overtrof het de uiteindelijke opbrengst van de Uniform-baseline op moeilijke taken (bijv. pick-place, push) na slechts 6 miljoen tijdstappen.
• Vergelijking met Baselines:
  • Uniform: Overtoewijzing aan gemakkelijke taken; faalt in het oplossen van moeilijke taken.
  • Learning Progress/Potential: Neigen om taken te prioriteren met hoge leersnelheden of hoge waarde-schatfouten, wat vaak correleert met gemakkelijke taken, en verwaarlozen moeilijke taken die vastzitten.
  • Hard First (SMT): Lijdt aan prestatie-instorting op moeilijke taken; zodra een taak als "opgelost" wordt verklaard op basis van een drempel, daalt de steekproefkans abrupt, waardoor de agent de taak vergeet en opnieuw moet leren.
• Ablaties:
  • Adaptieve Steekproef vs. Herweging: Adaptieve steekproef (DRATS) is aanzienlijk data-efficiënter dan het herwegen van het doel van uniform gestreepde data. Herweging leidt tot gradiënt-schattingen met hoge variantie voor moeilijke taken vanwege onvoldoende datacollectie.
  • Hyperparameters: DRATS is robuust ten opzichte van de keuze van temperatuur $\eta$ en stapgrootte $\alpha$.

Betekenis en Claims

Het artikel stelt dat DRATS een fundamentele inefficiëntie in MTRL aanpakt: de verkeerde toewijzing van dataresources. Door de focus te verschuiven van gradiëntconflict naar data-allokatie, biedt DRATS een data-efficiënt mechanisme om prestaties over taken in evenwicht te brengen.

De auteurs benadrukken dat:

• DRATS de eerste toepassing is van Distributionally Robust Optimization op multi-task RL in settings voor continue controle.
• In tegenstelling tot veel Curriculum Learning-methodes die positieve overdracht van gemakkelijke naar moeilijke taken aannemen, richt DRATS zich expliciet op de worst-case opbrengstgaping, waardoor het effectief is zelfs wanneer taken conflicteren of overdracht zwak is.
• De methode modulair is; het wijzigt alleen het datacollectieproces en kan worden gecombineerd met gespecialiseerde architecturen of gradiëntmanipulatie-technieken.

Het artikel concludeert dat DRATS een principiële, theoretisch onderbouwde aanpak biedt voor het bereiken van gebalanceerd multi-task leren, waardoor zowel de efficiëntie van het leren als de robuustheid van het uiteindelijke beleid over alle taken worden verbeterd.