Discrete World Models via Regularization

Davide Bizzaro; Luciano Serafini

cs.LGMar 2, 2026

Discrete World Models via Regularization

Authors: Davide Bizzaro, Luciano Serafini

Organizations: University of Padua, Padua, Italy · Fonadazione Bruno Kessler, Trento, Italy

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.

Figures & tables

Figure 1: Overview of the model architecture and of the loss function. Encoders map successive observations into a shared Boolean latent space, and a predictor transforms the current latent state into the next, given the action. We illustrate and evaluate this setup on an 8-puzzle benchmark with MNIST digits, where actions move the blank tile and induce local state changes and new renderings of the digits. The crossed lines denote operations that stop the gradient, and we introduce a training scheme with two-steps updates: initially, the predictor is updated solely based on hard bits $b$ , followed by a joint update of the encoder and predictor using probabilities $p$ . Test-time inference relies only on the pathway involving $b$ . The parameters $\phi^{\prime}$ are an EMA copy of the parameters $\phi$ .

Figure 2: Example transition in IceSlider.

Models	MNIST 8-puzzle				IceSlider
	Clean		Noisy		Clean		Noisy
	Enc.	Im.	Enc.	Im.	Enc.	Im.	Enc.	Im.
AE	41 ±26	36 ±22	50 ±10	44 ±8	96 ±11	83 ±11	83 ±8	84 ±9
$\beta$ -VAE	74 ±2	61 ±2	70 ±1	58 ±1	93 ±10	62 ±13	70 ±14	58 ±4
DeepcubeAI	79 ±3	77 ±4	74 ±2	73 ±2	92 ±14	85 ±14	81 ±14	80 ±13
DWMR	91 ±1	84 ±2	90 ±2	80 ±3	99 ±2	95 ±2	81 ±6	85 ±5

Table 1: Mean per-cell f1-scores, for reconstructing the board configurations from encoder states ( Enc. ) and from one-step roll-outs in imagination ( Im. ), using linear probes. We report the mean and standard deviation over

10

different runs, for each model and benchmark.

Model	Scheduling	Clean		Noisy
		Encoding	Imagination	Encoding	Imagination
Fully differentiable	yes	95 ±2	80 ±2	92 ±2	74 ±2
Fully differentiable	no	93 ±5	76 ±3	85 ±19	75 ±16
Straight-through	yes	96 ±1	85 ±1	88 ±15	81 ±13
Straight-through	no	96 ±4	80 ±2	86 ±4	77 ±4
Two steps	yes	93 ±6	86 ±5	95 ±2	83 ±3

Table 2: Mean per-cell accuracies, for reconstructing the board configurations from encoder states and from one-step roll-outs in imagination, using linear probes. We report the mean and standard deviation over

10

different runs on the MNIST 8-puzzle benchmark.

Model	Clean		Noisy
	Enc.	Im.	Enc.	Im.
DWMR	93 ±6	86 ±5	95 ±2	83 ±3
No $\mathcal{L}_{\text{var}}$	12 ±0	12 ±0	11 ±1	11 ±1
No $\mathcal{L}_{\text{cor}}$	42 ±8	39 ±6	39 ±4	35 ±3
No $\mathcal{L}_{\text{cos}}$	56 ±13	53 ±11	67 ±14	60 ±11
No $\mathcal{L}_{\text{loc}}$	89 ±15	82 ±14	38 ±24	33 ±23

Table 3: Ablation study starting from the DWMR hyperparameters. We report the mean and standard deviation over 10 different runs on the MNIST 8-puzzle benchmark of the mean per-cell accuracies.

World Model Learning

World model learning builds predictive models of how environments change under actions to support planning, simulation, and decision-making.

See all 414 papers

DWM: Separating World Effects from Actions in Latent World Models
Jul 21, 2026Yi-Ge Zhang, Tianqi Du, Qi Zhang +1World Model Learning Latent Dynamics Modeling
Peking University
CausalDreamer: Learning Predictive World Models with Latent Disentanglement
Oct 8, 2026Prince Jha, Nils Lukas, Kun Zhang +1World Model Learning Causal World Models
Mohamed bin Zayed University of Artificial Intelligence (MBZUAI) · Carnegie Mellon University
Dreamer-CDP: Improving Reconstruction-free World Models Via Continuous Deterministic Representation Prediction
Mar 7, 2026Michael Hauri, Friedemann ZenkeWorld Model Learning World Models
Friedrich Miescher Instiute for Biomedical Research, Fabrikstrasse 24, 4056 Basel, Switzerland

Jul 21, 2026J/K move · Enter open · S save

DWM: Separating World Effects from Actions in Latent World Models

Yi-Ge Zhang, Tianqi Du, Qi Zhang, Yisen Wang

Peking University

World Model Learning Latent Dynamics Modeling World Model-Based Planning Action-Conditioned World Models Model-Based Planning Latent World Models

Open paper AskPDF arXiv

Latent world models underpin much of modern model-based control, yet current action-conditioned formulations supervise the next-latent transition with a single, undifferentiated target, forcing a monolithic learning signal to absorb every source of state change. In real world, however, transitions arise from two heterogeneous sources: an action-driven component induced by the agent, and an action-invariant world effect -- the change that would still occur under a null action, dictated by the environment's intrinsic dynamics (e.g., gravity-driven sliding, inertia, contact rebound, and persistent drift). Fusing them into a single target entangles the two inside the latent transition, prevents the model from attributing observed changes to their underlying causes, and undermines the transferability of the learned dynamics. We introduce DWM (Decomposed World Model), a supervision-level framework that operationalizes this decomposition. DWM augments the predictor of a latent world model with an auxiliary world head, regularized by a normalized world-contrastive objective to be action-invariant, while the original pred head is coupled to it via an orthogonality constraint; together, the two signals induce an explicit additive decomposition of the predicted transition into an action-invariant and a complementary action-driven component, without altering the underlying architecture or inference pipeline. To evaluate DWM under persistent world effects, we construct W-variants of three standard control benchmarks -- PushT-W, Reacher-W, and TwoRoom-W -- each instantiating a distinct action-invariant dynamic. DWM matches strong baselines on the flat counterparts and delivers a mean absolute improvement of 13.1% in CEM planning success across the W-variants.

Abstract

World Model Learning

Explore similar work

DWM: Separating World Effects from Actions in Latent World Models

CausalDreamer: Learning Predictive World Models with Latent Disentanglement

Dreamer-CDP: Improving Reconstruction-free World Models Via Continuous Deterministic Representation Prediction