Rethinking the Predictive Loss in Latent World Models

Latent world models rely on predictive penalties to learn representations, meaning the choice of loss function dictates how prediction errors shape the latent space. The derivative of standard squared-error loss grows linearly with the error, amplifying the influence of large prediction errors. To address this, we investigate the element-wise log-cosh loss, which behaves quadratically for small errors but smoothly bounds the derivative for large ones. Through residual-space analysis on models with explicit representation regularization, we show that this attenuation improves local gradient interaction where it reduces gradient conflict more than cooperation. We applied this training-time loss replacement across multiple latent world model families without changing their underlying architectures or inference methods. On navigation, manipulation, and continuous control benchmarks, log-cosh consistently improved planning success over squared error and outperformed a locally curvature-matched quadratic control. On real-world driving data NAVSIM, log-cosh also improved perception and planning. We further study the learned representations through physical probing and action-space analyses. We observe that log-cosh makes task-relevant state information more accessible, and action-space analyses reveal closer alignment with the true-system cost landscape. These findings establish the predictive penalty as a fundamental component of representation design.