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

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.