Multivariate Distributional Reinforcement Learning Using Sliced Divergences

Technisches Resümee: Multivariate Distributional Reinforcement Learning mittels Sliced Divergences

Problemstellung

Distributional Reinforcement Learning (DRL) modelliert die vollständige Rückkehrverteilung (Return Distribution) anstatt nur deren Erwartungswert, was theoretische und empirische Vorteile gegenüber dem klassischen erwartungswertbasierten RL bietet. Die Erweiterung von DRL auf multivariate Settings (bei denen Belohnungen $d$-dimensionale Vektoren sind, $d > 1$) bleibt jedoch eine erhebliche Herausforderung.

Bestehende Ansätze stehen vor drei primären Limitationen:

1. Komputative Intraktabilität: Gängige Metriken wie der Wasserstein-Abstand leiden unter hohen Rechenkosten in hohen Dimensionen (typischerweise $O(n^3 \log n)$ für Optimal Transport Solver) und schlechten statistischen Raten, die von der Umgebungshöhendimension abhängen.
2. Mangel an Kontraktionsgarantien: Während das skalare DRL über etablierte Kontraktionsergebnisse unter dem Bellman-Operator verfügt, fehlen für multivariate Settings mit allgemeiner Matrix-Diskontierung (wobei der Diskontfaktor eine zustands- und aktionsabhängige Matrix $\Gamma(s, a) \in \mathbb{R}^{d \times d}$ ist) rigorose Kontraktionsgarantien. Bestehende Methoden versagen oft bei der Kontraktion unter anisotroper (nicht-uniformer) Diskontierung.
3. Gradienten-Bias in stochastischem Training: Viele Divergenzen, insbesondere die Wasserstein-Divergenz, erfüllen nicht die unverzerrte Stichproben-Gradienteneigenschaft (U), wenn das Bellman-Ziel aus einem einzelnen sampelten Nachfolger instanziiert wird (Standard-TD-Learning). Dies führt zu verzerrten Gradienten und suboptimalem distributional Matching.

Methodik: Sliced Distributional RL (SDRL)

Die Autoren führen Sliced Distributional Reinforcement Learning (SDRL) ein, ein Framework, das handhabbare eindimensionale Divergenzen via zufälliger Projektionen (Slicing) auf multivariate Rückkehrverteilungen hebt.

Kernmechanismus

Anstatt eine hochdimensionale Divergenz direkt zu berechnen, projiziert SDRL die multivariaten Verteilungen auf zufällige Richtungen $\theta \in S^{d-1}$, berechnet eine Basis-Divergenz $\Delta$ in einer Dimension und aggregiert die Ergebnisse.

• Uniform Slicing: Mittelt die Basis-Divergenz über $L$ zufällige Richtungen:
  $$S\Delta_p^p(\mu, \nu) = \int_{S^{d-1}} \Delta^p((P_\theta)_\#\mu, (P_\theta)_\#\nu) d\sigma(\theta)$$
• Max-Slicing (MSDRL): Optimiert über die diskriminativste Richtung, um stärkere Kontraktionsgarantien für allgemeine Matrix-Diskontierungen zu erhalten:
  $$MS\Delta(\mu, \nu) = \sup_{\theta \in S^{d-1}} \Delta((P_\theta)_\#\mu, (P_\theta)_\#\nu)$$

Das Framework unterstützt eine breite Klasse von Basis-Divergenzen, spezifisch analysiert werden Wasserstein ($W_p$), Cramér-Distanz ($C_2$) und Maximum Mean Discrepancy (MMD).

Theoretische Beiträge

Die Arbeit etabliert rigorose theoretische Grundlagen für SDRL:

1. Metrische Eigenschaften: Es wird bewiesen, dass sowohl Uniform- als auch Max-Sliced-Lifts die Metrik-Axiome bewahren, sofern die Basis-Divergenz eine Metrik ist.
2. Kontraktionsgarantien:
   • Uniform Slicing: Beweist Bellman-Kontraktion unter gemeinsamer skalarer Diskontierung ($\Gamma = \gamma I$) für jede Basis-Divergenz, die Translation Non-Expansion, Scale Contraction und Mixture $p$-Konvexität erfüllt.
   • Max-Slicing: Führt ein neuartiges Kontraktionsresultat für allgemeine dichte Matrix-Diskontierung (anisotrope Updates) ein. Es wird bewiesen, dass Max-Slicing eine Kontraktion unter der Operatornorm der Diskontierungsmatrix liefert – ein Setting, in dem Uniform Slicing und Standard-MMD oft keine Kontraktion gewährleisten.
3. Stichprobenkomplexität: Es wird demonstriert, dass Uniform Slicing die dimensionsfreie Konvergenzrate der eindimensionalen Basis-Divergenz erbt und somit den Fluch der Dimensionalität vermeidet. Die Raten von Max-Slicing skalieren polynomiell mit der Dimension ($O(d \log n / n)$), was signifikant besser ist als exakter hochdimensionaler Optimal Transport.
4. Analyse des Gradienten-Bias:
   • Uniform Slicing: Bewahrt die unverzerrte Stichproben-Gradienteneigenschaft (U), falls die Basis-Divergenz diese erfüllt. Dies macht Sliced Cramér und Sliced MMD kompatibel mit dem Standard-Ein-Stichproben-TD-Bootstrapping.
   • Max-Slicing: Verletzt die Eigenschaft (U) aufgrund des durch den Maximierungsschritt induzierten Selektionsbias. Die gewählte Richtung hängt vom spezifischen Stichproben-Batch ab, was zu verzerrten Gradienten in Standard-TD-Settings führt.

Experimentelle Ergebnisse

Die Autoren evaluieren SDRL anhand von drei Benchmarks: einer tabellarischen Chain-MDP, einer Pixel-basierten Maze-Umgebung und einer Teilmenge von Atari-Spielen.

1. Chain-Umgebung (Policy Evaluation):

   • Vergleich von One-Sample TD (Standard-Bootstrapping) vs. Near-Exact TD (explizite Mischungskonstruktion).
   • Ergebnis: Zielfunktionen, die die Eigenschaft (U) erfüllen (z. B. Sliced Cramer, Sliced MMD), erreichen ein akkurates distributional Matching im One-Sample-Regime. Zielfunktionen, die (U) verletzen (z. B. Sliced Wasserstein, Max-Sliced Varianten), degradieren signifikant, sofern nicht die volle Mischung als Target verwendet wird.
   • Max-Slicing: Performt gut im degenerierten eindimensionalen Fall (wo der Selektionsbias verschwindet), versagt jedoch in allgemeinen multivariaten Settings unter One-Sample TD.

2. Pixel-basierte Maze (Policy Evaluation):

   • Evaluierung der distributional Genauigkeit mittels des empirischen Wasserstein-2-Abstands zwischen Monte-Carlo-Returns und der Vorhersage des Critics.
   • Ergebnis: Sliced Cramér und Sliced MMD performen auf Augenhöhe mit ihren nicht-geslicedten Gegenstücken und übertreffen Sliced Wasserstein im One-Sample-TD-Regime deutlich, was die Wichtigkeit der Eigenschaft (U) bestätigt.

3. Atari Games (Control):

   • Evaluierung der Control-Performance auf dekomponierten mehrdimensionalen Belohnungen.
   • Ergebnis: Während Zielfunktionen, die (U) erfüllen, akkurate distributional Schätzungen liefern, erreicht Sliced Wasserstein-2 überraschenderweise eine starke Control-Performance trotz schlechter distributional Matching und verzerrter Gradienten. Dies deutet darauf hin, dass für die Steuerung (Control) die akkurate Schätzung des Return-Erwartungswerts kritischer sein kann als die Treue der vollständigen Verteilung, selbst wenn das distributional Objective verzerrt ist.

Bedeutung und Ansprüche

Das Paper beansprucht, einen prinzipienbasierten Rahmen zur Erweiterung von DRL auf multivariate Settings bereitzustellen, während gleichzeitig die komputative Handhabbarkeit und theoretische Garantien gewahrt bleiben.

• Praktische Empfehlung: Die Autoren identifizieren Sliced Cramér als starken Standard für multivariates Distributional Learning. Es ist recheneffizient ($O(n \log n)$), erfüllt die unverzerrte Gradienteneigenschaft (was Standard-TD-Training ermöglicht) und bietet Kontraktionsgarantien im Standard-isotropen Setting.
• Theoretische Einsicht: Die Arbeit klärt die Trade-offs zwischen verschiedenen Divergenz-Wahlen auf. Sie verdeutlicht, dass während Max-Slicing das Kontraktionsproblem für allgemeine Matrix-Diskontierungen löst, es einen Gradienten-Bias einführt, der es ohne Modifikation ungeeignet für das Standard-One-Sample-TD macht.
• Identifikation von Forschungslücken: Das Paper schließt mit der Feststellung, dass eine Divergenz, die gleichzeitig unter allgemeinen anisotropen Diskontierungen nachweislich kontraktierend ist (zumindest im Sinne einer Norm), die unverzerrte Stichproben-Gradienteneigenschaft besitzt und eine handhabbare Schätzung bietet, weiterhin eine offene Frage bleibt.

Zusammenfassend bietet SDRL einen skalierbaren und theoretisch fundierten Ansatz für multivariates DRL, wobei sich Sliced Cramér als die robusteste praktische Wahl für Standard-Reinforcement-Learning-Pipelines erweist.