Cosmic Dreams: Lucid Dreaming in Latent Space via Temporal Context Windows

Generative world models allow robots to simulate possible futures for planning, but modeling rigid-body dynamics in compressed latent spaces is difficult. Standard deterministic methods produce blurry predictions, while probabilistic ones add noise. By systematically testing 11 architectures, the authors show that stable and sharp predictions depend on temporal context, not model size. Their lightweight ResNet-Rollout model, with a 4-frame context window, achieves a Root Mean Square Error of 71.08 pixels—a 57% reduction compared to the best 84-million-parameter baseline—while training in 15 minutes. The results indicate that addressing partial observability of velocity via history windows is more effective than scaling model size for deterministic physics modeling.

Study at a glance

Abstract

Generative world models enable robotic agents to "dream" potential futures for planning. However, modeling rigid-body dynamics in high-dimensional compressed latent spaces (such as NVIDIA Cosmos) presents significant challenges. Standard deterministic approaches often suffer from conditional mean averaging (often visually mistaken for posterior collapse), resulting in blurry predictions, while probabilistic approaches introduce stochastic sampling noise. In this work, we investigate the trade-off between model capacity and temporal observability. Through a systematic ablation of 11 architectures, we demonstrate that "lucid" (stable and sharp) dreaming is a function of temporal context, not parameter count. We introduce ResNet-Rollout, a lightweight (∼1M parameter) architecture augmented with a 4-frame context window. This model achieves a Root Mean Square Error (RMSE) of 71.08 pixels, achieving a 57% reduction in error compared to the best-performing 84M-parameter Rollout U-Net baseline (165.20 px), while training in just 15 minutes on a single NVIDIA RTX 5090. Our results suggest that addressing the partial observability of velocity via history windows is more effective than scaling model size for deterministic physics modeling in this domain. Code is available at: github.com/Pouria-007/Cosmos-Lucid-Dreamer.