Aprendendo Modelos de Mundo Baseados em Características Visuais via Ação Latente Residual

Resumo

Modelos de mundo preveem transições futuras a partir de observações e ações. Trabalhos existentes focam predominantemente apenas na geração de imagens. Modelos de mundo baseados em características visuais, por outro lado, preveem características visuais futuras em vez de pixels brutos de vídeo, oferecendo uma alternativa promissora que é mais eficiente e menos propensa a alucinações. No entanto, abordagens atuais baseadas em características dependem de regressão direta, o que leva a previsões borradas ou colapsadas em interações complexas, enquanto a modelagem generativa em espaços de características de alta dimensão ainda permanece desafiadora. Neste trabalho, descobrimos que um novo tipo de representação de ação latente, que denominamos *Ação Latente Residual* (RLA), pode ser facilmente aprendida a partir de resíduos DINO. Também mostramos que RLA é preditiva, generalizável e codifica a progressão temporal. Com base em RLA, propomos o *Modelo de Mundo RLA* (RLA-WM), que prevê valores de RLA via *flow matching*. RLA-WM supera tanto os modelos de mundo baseados em características quanto os de difusão de vídeo do estado da arte em conjuntos de dados simulados e do mundo real, sendo ordens de magnitude mais rápido que a difusão de vídeo. Além disso, desenvolvemos duas técnicas de aprendizado robótico que utilizam RLA-WM para melhorar o aprendizado de políticas. A primeira é um modelo de ação mundial minimalista com RLA que aprende a partir de vídeos de demonstração sem ações. A segunda é o primeiro *framework* de RL visual treinado inteiramente dentro de um modelo de mundo aprendido apenas a partir de vídeos offline, usando uma recompensa alinhada ao vídeo e sem interações online ou recompensas artesanais. Página do projeto: https://mlzxy.github.io/rla-wm

English

World models predict future transitions from observations and actions. Existing works predominantly focus on image generation only. Visual feature-based world models, on the other hand, predict future visual features instead of raw video pixels, offering a promising alternative that is more efficient and less prone to hallucination. However, current feature-based approaches rely on direct regression, which leads to blurry or collapsed predictions in complex interactions, while generative modeling in high-dimensional feature spaces still remains challenging. In this work, we discover that a new type of latent action representation, which we refer to as *Residual Latent Action* (RLA), can be easily learned from DINO residuals. We also show that RLA is predictive, generalizable, and encodes temporal progression. Building on RLA, we propose *RLA World Model* (RLA-WM), which predicts RLA values via flow matching. RLA-WM outperforms both state-of-the-art feature-based and video-diffusion world models on simulation and real-world datasets, while being orders of magnitude faster than video diffusion. Furthermore, we develop two robot learning techniques that use RLA-WM to improve policy learning. The first one is a minimalist world action model with RLA that learns from actionless demonstration videos. The second one is the first visual RL framework trained entirely inside a world model learned from offline videos only, using a video-aligned reward and no online interactions or handcrafted rewards. Project page: https://mlzxy.github.io/rla-wm