Toward Compute-Optimal and Sample-Efficient World Models via Spectral Filtering

Abstract

Deep reinforcement learning agents remain notoriously sample inefficient, and learning within the imagination of a world model is among the most promising remedies. However, imagined experience is only as good as the world model that generates it: prediction errors compound over multi-step rollouts, degrading the trajectories that the policy learns from. The temporal backbone—the component responsible for modeling how observations evolve across time— is the architectural bottleneck. Motivated by recent advances in spectral state-space models, we introduce SWIFT (Spectral World model with Inner Filtering for Temporal dynamics), a data-efficient agent that replaces the causal Transformer in the IRIS world model with a spectral filtering module built from the fixed eigenvectors of a Hankel matrix. These filters carry provable regret guarantees on cumulative prediction error and operate at subquadratic cost. On the Atari 100k benchmark, SWIFT achieves lower long-horizon prediction error than its Transformer counterpart, scales favorably in memory and training time to longer temporal contexts, and delivers competitive or superior policy returns—all within a controlled comparison that isolates the effect of the backbone. Analysis of the decoder’s attention dynamics further reveals that the spectral backbone acts as a stable temporal prior, absorbing increasing predictive responsibility as spatial predictions degrade during extended rollouts. For reproducibility, we release our code and models at https://github.com/cbrownaz24/SWIFT .