World models using latent diffusion model

Abstract

World models model the external world from limited information and can be used to predict future external states and observations for learning. In spatio-temporal prediction, reinforcement learning methods using deep generative models have attracted attention. In generative models, Imagen and Stable-diffusion based on diffusion models are known for their high image generation capability. In this study, we propose a method to generate a better latent representation from the hidden states of LSTM by changing the vision part of World Models from conventional β-VAE to latent diffusion model, and compare these methods.