Boosting World Models Learning via Latent-Space Value Alignment

Model-based reinforcement learning aims to construct world models for efficient sampling. Current mainstream algorithms can be broadly categorized into two paradigms: maximum likelihood and value-aware world models. The former employs structured Recurrent/Transformer State-Space Models to capture environmental dynamics but overlooks task-relevant features. The latter prioritizes decision-critical states but suffers from sub-optimal performance. While recent efforts have sought to integrate these approaches, they typically rely on auxiliary modules or heavy external priors that significantly increase computational complexity. In this work, we propose a Value-Aligned World Model, a minimalist framework designed to synergize these two paradigms with negligible overhead. Specifically, We introduce an intrinsic latent-space value-alignment regularization that compels the world model to prioritize task-relevant features while maintaining the structural integrity of stochastic dynamics. To ensure stable optimization, we develop an adaptive weighting mechanism that acts as a self-regulating curriculum, balancing reconstruction fidelity with decision-making utility. Extensive experiments on Atari 100k and DeepMind Control benchmarks demonstrate that our algorithm consistently boosts existing methods with minimal added code and computational overhead. Code is available at supplementary material.

AI agents often learn by building an internal model of the environment to predict what may happen next before acting. However, these models can spend too much effort remembering irrelevant details in images or sensor inputs, causing the truly important information for decision making to be buried. Our work proposes a simple way to help AI focus on information that matters for completing the task. Instead of adding large models or complex modules, we guide the internal model to preserve information that affects future outcomes. This guidance is adjusted automatically: the agent first learns the basic rules of the environment, then gradually focuses more on decision-relevant details. Experiments on video games and robot-control tasks show that our method improves existing algorithms with almost no extra computational cost.