Discrete World Models via Regularization
This work addresses the need for compact, symbolic state representations in environments for search heuristics and planning, offering an incremental improvement over existing methods by eliminating the need for reconstruction or contrastive signals.
The paper tackled the problem of unsupervised Boolean world-model learning by introducing Discrete World Models via Regularization (DWMR), a reconstruction-free and contrastive-free method that uses a novel loss with specialized regularizers to maximize entropy and independence of representation bits, resulting in more accurate representations and transitions than reconstruction-based alternatives on benchmarks with combinatorial structure.
World models aim to capture the states and dynamics of an environment in a compact latent space. Moreover, using Boolean state representations is particularly useful for search heuristics and symbolic reasoning and planning. Existing approaches keep latents informative via decoder-based reconstruction, or instead via contrastive or reward signals. In this work, we introduce Discrete World Models via Regularization (DWMR): a reconstruction-free and contrastive-free method for unsupervised Boolean world-model learning. In particular, we introduce a novel world-modeling loss that couples latent prediction with specialized regularizers. Such regularizers maximize the entropy and independence of the representation bits through variance, correlation, and coskewness penalties, while simultaneously enforcing a locality prior for sparse action changes. To enable effective optimization, we also introduce a novel training scheme improving robustness to discrete roll-outs. Experiments on two benchmarks with underlying combinatorial structure show that DWMR learns more accurate representations and transitions than reconstruction-based alternatives. Finally, DWMR can also be paired with an auxiliary reconstruction decoder, and this combination yields additional gains.