Discrete World Models via Regularization

TL;DR

This paper introduces DWMR, a novel unsupervised method for learning discrete, Boolean world models without reconstruction, using regularizers to maximize entropy and independence, improving accuracy over existing approaches.

Contribution

DWMR presents a reconstruction-free, contrastive-free approach with regularizers for learning Boolean world models, enhancing representation accuracy and robustness.

Findings

DWMR outperforms reconstruction-based methods on benchmarks.

Regularizers improve the independence and entropy of representations.

Combining DWMR with a reconstruction decoder yields further gains.

Abstract

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.