Structural-Test Geometry as an Inductive Bias for Partially Observable World Models

Todd Y. Zhou

cs.LGOct 28, 2026

Structural-Test Geometry as an Inductive Bias for Partially Observable World Models

Authors: Todd Y. Zhou

Organizations: Harvard University (Undergraduate Student), Cambridge, MA, USA

Abstract

Predictive world-model objectives need not organize latent histories according to distinctions required for later control. With privileged simulator access during training, should response-derived structural information be predicted, decoded, or imposed directly on latent geometry--and how do these choices compare with structured belief distillation? We define posterior-marginalized structural tests whose response distributions induce a pseudometric over histories, then train quotient-preserving (QP) world models by aligning shared latent Euclidean distances to that metric. Across a fixed six-task diagnostic suite, the largest differences appear under held-out structural variation: QP improves planning success by 26.1 points $[21.5,30.7]$ over metric decoding and 6.4 points $[3.8,9.0]$ over structured belief distillation. On GridWorld, metric decoding accurately recovers held-out structural distances ( $ρ=.91$ ) while its native latent geometry remains more distorted than QP's (paired IQD difference $.22$ $[.20,.24]$ ), showing that relational recoverability and direct geometric alignment can diverge.

Figures & tables

Figure 1

Figure 2

Domain	TR	MD	BD-S	QP (ours)
GridWorld	$.814/.574$	$.792/.514$	$.842/.740$	$\mathbf{.874/.798}$
Key–Box	$.772/.521$	$.753/.492$	$.812/.701$	$\mathbf{.854/.783}$
Hidden-Mass CartPole	$.741/.492$	$.723/.473$	$.772/.653$	$\mathbf{.823/.742}$
Visual GridWorld	$.712/.473$	$.693/.452$	$.762/.661$	$\mathbf{.794/.722}$
MiniWorld–Room	$.674/.434$	$.652/.412$	$.722/.611$	$\mathbf{.733/.652}$
MiniGrid–HiddenKey	$.723/.502$	$.701/.482$	$.762/.642$	$\mathbf{.772/.693}$

Table 3

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Domain	Method	ID [95% CI]	Shift [95% CI]
GridWorld	RNC	$.710[.680,.740]$	$.480[.440,.520]$
	TR	$.814[.782,.846]$	$.574[.534,.614]$
	MD	$.792[.758,.826]$	$.514[.472,.556]$
	BD-S	$.842[.812,.872]$	$.740[.702,.778]$
	QP	$\mathbf{.874[.842,.906]}$	$\mathbf{.798[.756,.840]}$
Key–Box	RNC	$.680[.650,.710]$	$.450[.410,.490]$

Appendix

Table 4

Domain	Method	ID [95% CI]	Shift [95% CI]
Visual GridWorld	RNC	$.620[.590,.650]$	$.370[.330,.410]$
	TR	$.712[.678,.746]$	$.473[.429,.517]$
	MD	$.693[.659,.727]$	$.452[.410,.494]$
	BD-S	$.762[.730,.794]$	$.661[.623,.699]$
	QP	$\mathbf{.794[.762,.826]}$	$\mathbf{.722[.682,.762]}$
MiniWorld–Room	RNC	$.580[.550,.610]$	$.340[.300,.380]$

Appendix

Table 5

Method	Intended target metric	Value	Comparison reference
TR	response-distribution TV	$.18\ [.15,.21]$	split-half TV $.06$
MD	$d_{\mathcal{T}}$ decoding	RMSE $.12\ [.10,.14]$ ; $\rho=.91\ [.88,.93]$	constant RMSE $.38$
MD	native IQD	$.41\ [.38,.44]$	QP IQD $.19$
Paired MD–QP	IQD difference	$\mathbf{.22\ [.20,.24]}$	—
BD-S	posterior NLL (nats)	$1.45\ [1.33,1.57]$	prior NLL $2.10$
QP	native IQD	$.19\ [.17,.21]$	best-found SMACOF $.17$

Appendix

Table 6

Condition	ID [95% CI]	Shift [95% CI]
BD-S (bivariate Gaussian)	$.842[.812,.872]$	$.740[.702,.778]$
BD-S (factorized categorical $x/y$ )	$.838[.806,.870]$	$.736[.696,.776]$
QP	$.874[.842,.906]$	$.798[.756,.840]$
QP, family-balanced weights	$.866[.832,.900]$	$.787[.745,.829]$
QP, uniform pair sampling	$.862[.830,.894]$	$.782[.742,.822]$
QP, 500 spurious nuisance coordinates	$.844[.812,.876]$	$.754[.714,.794]$

Appendix

Table 7

Diagnostic/subset	$\boldsymbol{\rho}$	95% CI
IQD, all models (unstratified)	$.83$	$[.79,.87]$
IQD, pooled within-cell ranks (non-QP)	$.70$	$[.63,.77]$
High action-value range	$.75$	$[.68,.81]$

Appendix

Table 8

Architecture	Method	IQD [95% CI]	ID [95% CI]	Shift [95% CI]
TD-MPC style	Value only	$.58\ [.54,.62]$	$.52\ [.49,.55]$	$.31\ [.27,.35]$
TD-MPC style	RNC	$.40\ [.36,.44]$	$.67\ [.64,.70]$	$.40\ [.36,.44]$
TD-MPC style	QP	$\mathbf{.21\ [.17,.25]}$	$\mathbf{.74\ [.71,.77]}$	$\mathbf{.62\ [.58,.66]}$
RSSM reference	QP	$.19\ [.17,.21]$	$.77\ [.74,.80]$	$.69\ [.65,.73]$

Appendix

Table 9

Domain	Horizon	$\|\mathcal{A}\|$	Hidden factor	Train/val episodes	Observation
Hidden-Goal GridWorld	50	5	4 goal corners	20k/2k	coordinates
Key–Box	50	4	key room/possession + door	20k/2k	symbolic
Hidden-Mass CartPole	200	2	mass in $[.5,1.5]$	10k/1k	continuous
Visual GridWorld	50	5	4 goal corners	20k/2k	$64\!\times\!64$ RGB
MiniWorld–Room	100	4	4 object locations	10k/1k	first-person RGB
MiniGrid–HiddenKey	100	3	3 key regions	10k/1k	$7\!\times\!7$ egocentric

Appendix

Table 10

Domain	Probe settings	Continuations	Outcome	$\|\mathcal{T}\|$
GridWorld	corner teleport (4), centre teleport (1), no-op (1)	5 each (3 random walks, 2 heuristics), plus 70 random-walk variants	revealed corner (4-way)	100
Key–Box	grab key (1), open door (2), door-state probe (2), no-op (1)	5 each, plus 70 random-walk variants	binary	100
CartPole	set mass $\{.70,.85,1.0,1.15,1.30\}$ , no-op (1)	4 each, plus 76 random continuations	scalar return	100
Visual GridWorld	rendered corner teleport (4), centre teleport (1), no-op (1)	GridWorld construction with rendered observations	revealed corner (4-way)	100
MiniWorld	location teleport (4), look-at location (4), no-op (1)	5 each, plus 55 random-walk variants	object visible	100
MiniGrid	reveal key (3), quadrant teleport (3), no-op (1)	5 each, plus 65 random continuations	key quadrant	100

Appendix

Table 11

Domain	Training values/compositions	Held-out shift
GridWorld	goals $(1,1),(1,5),(5,1)$	goal $(5,5)$
Key–Box	key room A; door A–B	key room B; door B–C
CartPole	mass $\{.8,1.0,1.2\}$	mass $\{.6,1.4\}$
Visual GridWorld	goals $(1,1),(1,4),(4,1)$	goal $(4,4)$
MiniWorld–Room	desk, shelf, table	cabinet; new room configuration
MiniGrid–HiddenKey	top-left, top-right	bottom-centre

Appendix

Table 12

Domain	BD-S target	Prediction family	Auxiliary loss
GridWorld / Visual GridWorld	goal coordinate $(x_{g},y_{g})$	bivariate Gaussian	posterior-coordinate NLL
Key–Box	key room, possession, door state	factorized categoricals	summed cross-entropy
CartPole	pole mass on $\{.8,1.0,1.2\}$	five-component Gaussian mixture	posterior-mass NLL
MiniWorld–Room	object coordinate	bivariate Gaussian	posterior-coordinate NLL
MiniGrid–HiddenKey	key-region coordinates	factorized categorical	summed cross-entropy

Appendix

Table 13

Domain	$\boldsymbol{\widehat{L}_{.95}}$	Miss	Grid maximum
GridWorld	$1.6$	$1.1\%$	$1.1\%$
Key–Box	$1.4$	$1.2\%$	$3.8\%$
CartPole	$2.0$	$1.5\%$	$4.1\%$

Appendix

Table 14

Method	Candidate coefficients	Selected coefficient
TR	$\{.1,.5,1,2,5\}$	$.5$
MD	$\{.1,.5,1,2,5\}$	$1$
BD-S	$\{.1,.5,1,2,5\}$	$.5$
QP	$\{.1,.5,1,2,5\}$	$1$

Appendix

Table 15

Method	RNC	TR	MD	BD-S	QP
GPU-hours/seed	$4.2$	$4.1$	$4.4$	$4.0$	$4.5$
Total GPU-hours	$252$	$246$	$264$	$240$	$270$

Appendix

Table 16

World Model Learning

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

See all 419 papers

ATLAS: Aligned Transport of Latent Structure for Reliable World Model Planning
Sep 28, 2026Ke Fang, Yupu Yao, Lu ChengWorld Models World Model-Based Planning
Pennsylvania State University
Supervise What Decides Success: Criterion-Aligned Auxiliary Losses for Latent World-Model Planning
Oct 1, 2026Takumi Hara, Kanata SuzukiWorld Model Learning World Models for Robotics
Kyoto University Kyoto, Japan · Fujitsu Limited Kanagawa, Japan
Latent State Design for World Models under Sufficiency Constraints
May 3, 2026Keon Woo KimRepresentation Learning Predictive Representation Learning
Om Labs

Sep 28, 2026J/K move · Enter open · S save

ATLAS: Aligned Transport of Latent Structure for Reliable World Model Planning

Ke Fang, Yupu Yao, Lu Cheng

Pennsylvania State University

World Models World Model-Based Planning Geometric Representation Learning Model-Based Planning Latent World Models

Open paper AskPDF arXiv

Latent world models rely on representation geometry for planning, yet regularizing the latent marginal alone does not determine the state-to-state relationships used for action selection. We show that this can cause planning-relevant novelty structure to be weakened as representations are transformed into the final latent used by the planner. We introduce Aligned Transport of Latent Structure (ATLAS), a training objective that explicitly preserves relational geometry while calibrating the global latent distribution. ATLAS transfers normalized pairwise structure from an informative encoder representation to the planning latent and uses Wasserstein embedding matching (WEMReg) to calibrate its marginal through one-dimensional Wasserstein-2 transport. Our analysis shows that relational preservation and marginal calibration impose non-redundant constraints, and connects finite-candidate planning stability to relational distortion, latent-scale mismatch, and prediction error. Instantiated in LeWM, ATLAS improves mean goal-reaching success across PushT, TwoRoom, and OGBench-Cube on both lower- and higher-novelty evaluation subsets, with the largest gain on higher-novelty TwoRoom episodes. Representation and rollout diagnostics further show stronger novelty-related structure in the planning latent, improved marginal calibration, and lower multi-step prediction error. Together, these results highlight preservation of planning-relevant latent geometry as an important ingredient for reliable world-model planning. Code is available at https://anonymous.4open.science/r/atlas-world-model-72C4/.

Abstract

World Model Learning

Explore similar work

ATLAS: Aligned Transport of Latent Structure for Reliable World Model Planning

Supervise What Decides Success: Criterion-Aligned Auxiliary Losses for Latent World-Model Planning

Latent State Design for World Models under Sufficiency Constraints