Optimierung von Weltmodellen im Reinforcement Learning durch Next-Embedding Prediction

Verbesserte Weltmodelle im Reinforcement Learning: Die Rolle der "Next-Embedding Prediction"

Im Bereich des modellbasierten Reinforcement Learning (MBRL) stellt die Fähigkeit, temporale Abhängigkeiten in teilweise beobachtbaren, hochdimensionalen Umgebungen zu erfassen, eine zentrale Herausforderung dar. Aktuelle Forschungen zeigen nun vielversprechende Ansätze zur Stärkung von Weltmodellen durch die Einführung der "Next-Embedding Prediction". Diese Methode, die wir im Folgenden detailliert beleuchten werden, verspricht signifikante Fortschritte in der Entwicklung lernfähiger Agenten.

Die Evolution der Weltmodelle: Von Rekonstruktion zur Vorhersage

Traditionell basieren viele Weltmodelle im Reinforcement Learning auf der Rekonstruktion von Pixeln. Modelle wie die Dreamer-Familie nutzen einen Pixel-Decoder, um Beobachtungen zu rekonstruieren und so reichhaltige, steuerungsrelevante Merkmale zu erzeugen. Dieser Ansatz ist jedoch mit einem hohen Modellierungsaufwand verbunden. Die Rekonstruktion führt zu einem generativen Ziel, das die Optimierung erschwert und Kapazitäten auf visuell detaillierte, aber möglicherweise aufgabenirrelevante Aspekte lenken kann. Decoder-freie Methoden haben sich daher etabliert, um die Pipeline zu vereinfachen und die Effizienz zu steigern, indem sie Repräsentationen direkt lernen.

Fehlende explizite zeitliche Beschränkungen in vielen dieser decoder-freien Ansätze führen jedoch oft zu einer geringen prädiktiven Fähigkeit über längere Zeiträume hinweg. Dies ist insbesondere in Umgebungen mit teilweiser Beobachtbarkeit problematisch, wo ein Agent Informationen über die Zeit integrieren muss, anstatt nur auf einzelne Frames zu reagieren. Hier setzt die "Next-Embedding Prediction" an.

NE-Dreamer: Ein Paradigmenwechsel in der Repräsentationslernen

Ein kürzlich vorgestellter Ansatz, NE-Dreamer, revolutioniert dieses Feld, indem er auf die pixelbasierte Rekonstruktion verzichtet und stattdessen direkt die zeitliche prädiktive Ausrichtung in seinen latenten Repräsentationen optimiert. NE-Dreamer ersetzt die Rekonstruktion auf Pixelebene durch ein einfaches, aber leistungsstarkes Ziel: Zu jedem Zeitschritt sagt ein kausaler Temporal Transformer das nächste Encoder-Embedding in der Sequenz voraus. Diese Vorhersage wird dann mit dem tatsächlichen Embedding des nächsten Schritts abgeglichen, beispielsweise mithilfe einer Redundanzreduktionsmetrik wie Barlow Twins.

Dieser Fokus auf die Vorhersage des nächsten Embeddings ermöglicht es NE-Dreamer, kohärente, prädiktive Zustandsrepräsentationen zu lernen, ohne auf Rekonstruktionsverluste, Datenaugmentation oder zusätzliche Regularisierung angewiesen zu sein. Die Architektur von NE-Dreamer behält dabei die RSSM-Dynamik (Recurrent State-Space Model) und den auf Imagination basierenden Actor-Critic von Dreamer bei, während die Art und Weise, wie die latente Repräsentation gelernt wird, grundlegend verändert wird.

Architektonische Details und Funktionsweise

NE-Dreamer arbeitet in einer teilweise beobachtbaren Umgebung, in der ein Agent zu jedem Zeitpunkt t eine Bildbeobachtung x_t erhält, eine Aktion a_t auswählt und eine Belohnung r_t sowie einen Fortsetzungsindikator c_t erhält.

Problemstellung und Modellarchitektur

Der Kern von NE-Dreamer liegt in seinem latenten Weltmodell, das auf einem Rekurrenten Zustandsraummodell (RSSM) aufbaut. Ein Encoder bildet Beobachtungen auf Embeddings ab. Das RSSM aktualisiert seinen deterministischen Zustand basierend auf dem vorherigen latenten Zustand und der vorherigen Aktion. Anschließend werden Prior- und Posterior-Verteilungen über den stochastischen latenten Zustand definiert. Während des Trainings des Weltmodells wird aus dem Posterior abgetastet; während der Imagination aus dem Prior.

Das Weltmodell sagt Belohnungen und Fortsetzungen voraus, ähnlich wie bei Dreamer. Der entscheidende Unterschied ist jedoch das Fehlen eines Pixel-Decoders. Stattdessen wird die Vorhersage des nächsten Embeddings als Hauptziel verwendet.

Next-Embedding Predictive Alignment

NE-Dreamer trainiert die latente Dynamik so, dass sie im Repräsentationsraum prädiktiv ist. Ein kausaler Temporal Transformer T_\theta nutzt ausschließlich Informationen bis zum Zeitpunkt t, um eine Vorhersage \hat{e}_{t+1} für das Embedding des nächsten Schritts zu erzeugen. Das Ziel ist das tatsächliche Encoder-Embedding e_{t+1}^\star des nächsten Schritts, wobei ein Stop-Gradient-Operator verwendet wird, um den Gradientenfluss zu steuern.

Für die Verlustfunktion \mathcal{L}_{\text{NE}} wird eine Barlow Twins Redundanzreduktionsfunktion verwendet. Diese Funktion fördert Invarianz (hohe diagonale Korrelationen) und unterdrückt Redundanz (geringe nicht-diagonale Korrelationen), angewendet auf die Vorhersage des nächsten Schritts und nicht auf die Übereinstimmung im selben Zeitschritt.

Actor-Critic Lernen

Wie bei DreamerV3 lernt NE-Dreamer eine Policy und eine Wertfunktion im latenten Raum, indem es imaginierte Trajektorien mit dem Weltmodell generiert. Diese imaginären Trajektorien ermöglichen effiziente Batch-Actor-Critic-Updates. Die Policy und der Critic werden auf Basis dieser imaginierten Rollouts trainiert, wobei die Policy die Aktionen maximiert und der Critic die λ-Returns vorhersagt.

Experimentelle Ergebnisse und deren Implikationen

Die Evaluierung von NE-Dreamer erfolgte auf zwei weit verbreiteten RL-Benchmarks: der DeepMind Control Suite (DMC) und DeepMind Lab (DMLab). Die Ergebnisse unterstreichen die Effektivität des Ansatzes:

• DMLab Rooms: Auf den speicher- und navigationsintensiven DMLab Rooms-Aufgaben zeigt NE-Dreamer eine dramatische Verbesserung der Ergebnisse. Es lernt zuverlässiger und erreicht deutlich höhere Leistungen als alle Vergleichsmethoden, einschließlich starker decoder-basierter (DreamerV3) und decoder-freier Baselines (R2-Dreamer, DreamerPro). Dies wird insbesondere bei Aufgaben deutlich, die eine Integration von Informationen über lange Zeiträume und räumliches Denken erfordern.
• Ablationsstudien: Gezielte Ablationsstudien bestätigen, dass die Leistungssteigerungen hauptsächlich auf die prädiktive Sequenzmodellierung (kausaler Transformer und Next-Step Target Shift) zurückzuführen sind. Das Entfernen des kausalen Temporal Transformers oder das Verschieben des Vorhersageziels auf den aktuellen Zeitschritt führt zu einem erheblichen Leistungsabfall.
• DeepMind Control Suite (DMC): Auf den Standard-Benchmarks der DMC erreicht NE-Dreamer eine Leistung, die der von DreamerV3 und anderen führenden decoder-freien Agenten entspricht oder diese sogar übertrifft. Dies deutet darauf hin, dass die Vorteile im anspruchsvolleren DMLab-Bereich keine Leistungseinbußen in anderen Bereichen mit sich bringen.

Eine post-hoc Analyse der latenten Repräsentationen von NE-Dreamer mittels eines trainierten Pixel-Decoders zeigt zudem, dass diese in der Lage sind, relevante Objekte, räumliche Anordnungen und aufgabenbezogene Merkmale konsistent über die Zeit zu bewahren. Im Gegensatz dazu zeigen andere Modelle oft Inkonsistenzen.

Fazit und Ausblick

Die Forschungsergebnisse zu NE-Dreamer etablieren die "Next-Embedding Prediction" in Kombination mit einem kausalen Temporal Transformer als einen effektiven und skalierbaren Rahmen für das modellbasierte Reinforcement Learning in komplexen, teilweise beobachtbaren Umgebungen. Durch den Verzicht auf Pixelrekonstruktion und die direkte Optimierung der temporalen prädiktiven Ausrichtung in latenten Repräsentationen können Weltmodelle robuster und effizienter lernen.

Obwohl die aktuellen Experimente sich auf Umgebungen konzentrieren, in denen die Langzeitstruktur und nicht feine visuelle Details die primäre Herausforderung darstellen, eröffnen diese Erkenntnisse neue Wege für die Entwicklung noch leistungsfähigerer KI-Systeme. Zukünftige Arbeiten könnten die Anwendung dieses Ansatzes in visuell komplexeren Domänen untersuchen und alternative Alignment-Verlustfunktionen erforschen. Das Potenzial, die Effizienz und Robustheit von lernenden Agenten durch diese innovative Methode weiter zu steigern, ist beträchtlich und verspricht spannende Entwicklungen im Bereich der Künstlichen Intelligenz.

Zusätzliche Überlegungen: Modality-Agnostic Potentials und Generative Modelle

Die Prinzipien der Next-Embedding Prediction sind nicht auf visuelle Daten beschränkt. Der Grundgedanke, zukünftige Embeddings zu prognostizieren, findet Parallelen in der generativen Vorhersage in der Verarbeitung natürlicher Sprache. Dies deutet auf ein Potenzial für modalitätsagnostische Anwendungen hin, bei denen verschiedene Modalitäten (Text, Bild, Audio) unter einem einheitlichen Next-Embedding-Ziel trainiert werden könnten. Embeddings könnten dabei als gemeinsame Repräsentationswährung dienen.

Darüber hinaus eröffnet die Formulierung von NE-Dreamer auch Möglichkeiten für generative Modelle. Durch die Kombination des autoregressiven Embedding-Prädiktors mit einem geeigneten Bild-Decoder oder einem diffusionsbasierten Generator könnte der gleiche Rahmen für die Bildsynthese oder -bearbeitung genutzt werden. Dies würde eine Brücke zwischen Repräsentationslernen und Generierung innerhalb eines vereinheitlichten Modells schlagen – ein vielversprechendes Forschungsfeld für die Zukunft.

Bibliographie

- Bredis, G., Balagansky, N., Gavrilov, D., & Rakhimov, R. (2026). Next Embedding Prediction Makes World Models Stronger. arXiv preprint arXiv:2603.02765. - Xu, S., Ma, Z., Chai, W., Chen, X., Jin, W., Chai, J., Xie, S., & Yu, S. X. (2025). Next-Embedding Prediction Makes Strong Vision Learners. arXiv preprint arXiv:2512.16922. - Hafner, D., Lillicrap, T. P., Norouzi, M., & Ba, J. (2019a). Dream to control: learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603. - Hafner, D., Lillicrap, T. P., Norouzi, M., & Ba, J. (2019b). Learning latent dynamics for planning from pixels. In Proceedings of the 36th International Conference on Machine Learning (ICML), PMLR, 97, 2555-2565. - Hafner, D., Lillicrap, T. P., Ba, J., & Norouzi, M. (2021). Mastering Atari with discrete world models. In International Conference on Learning Representations (ICLR). - Hafner, D., Pasukonis, J., Ba, J., & Lillicrap, T. P. (2025). Mastering diverse control tasks through world models. Nature, 640(8059), 647-653. - Grill, J., Strub, F., Altché, F., Tallec, C., Richemond, P. H., Buchatskaya, E., ... & Valko, M. (2020). Bootstrap your own latent: A new approach to self-supervised learning. Advances in neural information processing systems, 33, 21271-21284. - Zbontar, J., Jing, L., Misra, I., LeCun, Y., & Deny, B. (2021). Barlow Twins: Self-supervised learning via redundancy reduction. In Proceedings of the 38th International Conference on Machine Learning (ICML), PMLR, 139, 12310-12320. - Li, Y., Inan, H. A., Yue, X., Chen, W. N., Wutschitz, L., Kulkarni, J., ... & Rajmohan, S. (2025). Simulating environments with reasoning models for agent training. arXiv preprint arXiv:2511.01824. - Yang, C., Wang, X., Zhang, Q., Jiang, Q., & Huang, X. (2025). Efficient integration of external knowledge to LLM-based world models via retrieval-augmented generation and reinforcement learning. In Findings of the Association for Computational Linguistics: EMNLP 2025, 9484-9501.