Factored Latent Action World Models

Zizhao Wang; Chang Shi; Jiaheng Hu; Kevin Rohling; Roberto Martín-Martín; Amy Zhang; Peter Stone

cs.LGFeb 18, 2026

Factored Latent Action World Models

Authors: Zizhao Wang, Chang Shi, Jiaheng Hu, Kevin Rohling, Roberto Martín-Martín, Amy Zhang, Peter Stone

Organizations: University of Texas at Austin · Sony AI

Abstract

Learning latent actions from action-free video has emerged as a powerful paradigm for scaling up controllable world model learning. Latent actions provide a natural interface for users to iteratively generate and manipulate videos. However, most existing approaches rely on monolithic inverse and forward dynamics models that learn a single latent action to control the entire scene, and therefore struggle in complex environments where multiple entities act simultaneously. This paper introduces Factored Latent Action Model (FLAM), a factored dynamics framework that decomposes the scene into independent factors, each inferring its own latent action and predicting its own next-step factor value. This factorized structure enables more accurate modeling of complex multi-entity dynamics and improves video generation quality in action-free video settings compared to monolithic models. Based on experiments on both simulation and real-world multi-entity datasets, we find that FLAM outperforms prior work in prediction accuracy and representation quality, and facilitates downstream policy learning, demonstrating the benefits of factorized latent action models.

Figures & tables

Figure 1: In multi-entity scenarios, (left) such as an intersection with three road users: (middle) a vanilla latent action model encodes the scene change with a single latent action of dimension

d

, which makes learning challenging as this latent action space needs to model all

|\mathcal{A}|^{K}

joint action combinations. (right) In contrast, FLAM decomposes the state into

K

factors, each with its own latent action of dimension

\frac{d}{K}

. Additionally, we assume all latent actions share the same space (i.e., with the same prior / codebook), which reduces the learning problem to modeling the

|\mathcal{A}|

actions per factor rather than their

|\mathcal{A}|^{K}

joint combinations.

Figure 2: Latent action model.

Figure 3: Two training stages of FLAM. (a) A VQ-VAE is pretrained to extract features for latent action model learning. (b) FLAM infers latent actions and makes predictions for each factor independently, with all modules trained jointly to minimize the prediction error.

Dataset	Metric	FLAM (ours)	AdaWorld	Genie	World Model	PlaySlot	SlotFormer	Recon
Average	PSNR ( $\uparrow$ )	34.9	24.7	28.1	N/A	29.1	19.9	37.1
	SSIM ( $\uparrow$ )	0.890	0.862	0.859	N/A	0.847	0.814	0.909
	LPIPS ( $\downarrow$ )	0.051	0.096	0.091	N/A	0.120	0.223	0.035
	FVD ( $\downarrow$ )	419.6	749.5	717.5	N/A	858.0	1597.7	215.1
MultiGrid	PSNR ( $\uparrow$ )	56.5	24.0	33.9	56.1	56.5	24.9	56.5
	SSIM ( $\uparrow$ )	0.944	0.909	0.933	0.943	0.944	0.912	0.944

Table 1: Prediction performance of all methods on simulation and real-world datasets. Recon reports the reconstruction quality of the pretrained VQ-VAE, serving as a reference upper bound for prediction performance.

	FLAM(ours)	PlaySlot	SlotFormer	Random
D ( $\uparrow$ )	0.91	0.81	0.85	0.32
C ( $\uparrow$ )	0.91	0.81	0.85	0.30
I ( $\uparrow$ )	0.93	0.86	0.90	0.10

Table 2: Factor-agent correspondence in the MultiGrid environments, evaluated by DCI (disentanglement, completeness, informativeness) metrics.

$K$	2	4	8	16
PSNR ( $\uparrow$ )	53.4	56.5	56.5	56.5
SSIM ( $\uparrow$ )	0.943	0.944	0.944	0.944
LPIPS ( $\downarrow$ )	4e-4	3e-4	3e-4	3e-4
FVD ( $\downarrow$ )	3.62	2.87	2.80	2.73

Table 3: Prediction performance of FLAM on the MultiGrid dataset with 4 agents and different number of factors

K

Dataset	FLAM	vanilla BC	random
Bigfish (1k)	1.8	1.0	0.9
Bigfish (10k)	8.8	3.6	0.9
Starpilot (1k)	1.8	1.9	1.1
Starpilot (10k)	9.1	6.4	1.1

Table 4: Behavior cloning policy performance, evaluated by mean episodic return.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

	Simulation	Real-world
Encoder architecture	IMPALA	MAGVIT-v2
Decoder architecture	IMPALA	MAGVIT-v2
Image size	128 (Multigrid), 224 (Procgen)	224
$d_{z}$	128	128
Tokenizer quantizer	FSQ	FSQ
Tokenizer codebook size	1024	16384

Appendix

Table 5: Experiment hyperparameters. Simulation datasets are Multigrid and Procgen, while the real-world dataset is nuPlan.

Panel of Figure 4 — Figure 4: Prediction performance variation along with increasing number of entities in the scene.

Figure 5: UMAP projection of the learned latent actions on the MultiGrid dataset. Each point corresponds to a latent action inferred by the IDM for one factor on a single transition from the observation-only dataset. Points are colored by the ground-truth actions taken by the corresponding agent at that transition (action labels are used only for visualization and are not used for training).

Panel of Figure 6 — Figure 6: FLAM allows for changing / editing the motion of one entity without affecting the others. Here we show controllable video generation on MultiGrid through specifying latent actions for one of the movable agents, while other agents follow their original inferred latent actions.

Panel of Figure 7 — Figure 7: Two example rollouts where two entities (green and red ones) share the same actions. FLAM consistently maps the two correlated agents to the same factor.

Panel of Figure 8 — Figure 8: Ablation of FLAM on 4-agent MultiGrid dataset, with $K\in\{2,4,8,16\}$ . When $K$ is less than the number agents, FLAM assigns two agents to each factor to maximize latent action utilization. When $K$ is equal to or greater than the number of agents, FLAM consistently assigns each entity to a distinct factor, even with excessive $K$ , demonstrating FLAM’s robustness to over-specifying $K$ .

Metric	FLAM (ours)	FLAM w/o Factorizer Temporal Attention	Global-Coupled FLAM
PSNR ( $\uparrow$ )	56.5	51.6	56.3
SSIM ( $\uparrow$ )	0.944	0.941	0.944
LPIPS ( $\downarrow$ )	3e-4	3e-3	4e-4
FVD ( $\downarrow$ )	2.9	36.6	3.72

Appendix

Table 6: Prediction performance of FLAM and its ablative variants on the MultiGrid dataset.

	FLAM(ours)	FLAM w/o Factorizer Temporal Attention	Global-Coupled FLAM	Random
D ( $\uparrow$ )	0.91	0.33	0.79	0.32
C ( $\uparrow$ )	0.91	0.23	1.00	0.30
I ( $\uparrow$ )	0.93	0.20	0.58	0.10

Appendix

Table 7: Factor-agent correspondence of FLAM and its ablative variants on the MultiGrid dataset, evaluated by DCI (disentanglement, completeness, informativeness) metrics.

Panel of Figure 9 — Figure 9: Architecture ablations of FLAM on the 4-agent MultiGrid dataset. Removing temporal attention from the factorizer leads to temporally inconsistent factor assignments and degraded long-horizon predictions (notably in the last few timestamps). In contrast, the Global-Coupled variant can produce accurate predictions but tends to assign multiple entities into the same factor, resulting in poorer disentanglement than the original FLAM implementation.

Metric	$\beta=1e-5$	$\beta=1e-4$	$\beta=1e-3$	$\beta=1e-2$
PSNR ( $\uparrow$ )	56.5	56.5	56.3	24.9
SSIM ( $\uparrow$ )	0.944	0.944	0.944	0.911
LPIPS ( $\downarrow$ )	3e-4	3e-4	4e-4	0.057
FVD ( $\downarrow$ )	2.84	3.40	3.78	479.9

Appendix

Table 8: Prediction performance of FLAM on the MultiGrid dataset with different KL regularization coefficient

\beta

values.

Figure 10: Prediction rollouts of all methods on the MultiGrid dataset, with inaccuracies highlighted by red boxes (except for SlotFormer, where the errors are visually apparent). When applicable, we also show the slot-patch binding masks from slot attention, displayed on top of the ground-truth observations. FLAM maintains consistent factor-entity bindings over time and generates accurate predictions. In contrast, AdaWorld and SlotFormer fail to generate accurate predictions from the start, and Genie makes mistakes at a few steps. Meanwhile, PlaySlot and SlotFormer rely on object-centric representations learned from reconstruction which incorrectly bind the red and purple entities to the same slot.

Figure 11: Prediction rollouts of all methods on the Bigfish dataset, with inaccuracies highlighted by red circles (except for SlotFormer, where the errors are visually apparent). When applicable, we also show the slot-patch binding masks from slot attention, displayed on top of the ground-truth observations. FLAM maintains consistent factor-entity bindings over time, e.g., the brown slot consistently attends to the green fish controlled by the player. In terms of prediction quality, FLAM and Genie only make errors on the occluded fish. AdaWorld to track the speeds of the circled yellow and red fish, which move more slowly than in the ground truth. The World Model baseline also struggles to capture the motion of non-controllable fish, since only the player’s action is available. PlaySlot and SlotFormer rely on reconstruction-based object-centric representations that do not yield consistent slot assignments. As a result, PlaySlot fails to preserve consistent shapes and colors for multiple fish.

Figure 12: Prediction rollouts of all methods on the Leaper dataset, with inaccuracies highlighted by red boxes (except for PlaySlot and SlotFormer, where the errors are visually apparent). When applicable, we also show the slot-patch binding masks from slot attention, displayed on top of the ground-truth observations. FLAM generates accurate predictions, in particular capturing how the frog jumps between logs. Genie and World Model achieve the second best prediction quality, but fails to model the frog’s jumps in some timestamps. In contrast, AdaWorld and PlaySlot generate inaccurate predictions from the start, either removing the frog and cars from the scene or hallucinating non-existent.

Figure 13: Prediction rollouts of all methods on the Starpilot dataset, with inaccuracies highlighted by red circles (except for PlaySlot and SlotFormer, where the errors are visually apparent). When applicable, we also show the slot-patch binding masks from slot attention, displayed on top of the ground-truth observations. FLAM maintains a consistent binding for the player in the bottom-left corner (shown in cyan) and generates accurate predictions, except for the enemy bullet heading toward the player. Other entities (e.g., bullets, enemies, and meteors) move at constant speed, so even with less consistent bindings, FLAM can still achieve reasonable predictions. In contrast, AdaWorld, Genie, and the World Model baseline generate inaccurate predictions from the start, frequently shifting enemy appearances and missing the enemy bullet toward the player.

Figure 14: Prediction rollouts of all methods on the nuPlan dataset, with inaccuracies highlighted by red boxes (except for SlotFormer, where the errors are visually apparent). Models are conditioned on a fixed history window of size $w=5$ . When applicable, we also show the slot-patch binding masks from slot attention, displayed on top of the ground-truth observations. Compared to the baselines, FLAM generates relatively accurate predictions for the grey car moving from the center to the right. In contrast, AdaWorld and Genie fail to maintain a consistent appearance for this moving car. Meanwhile, PlaySlot and SlotFormer fail to model the dynamics and largely ignore the cars in the scene.

Explore similar work

DiLA: Disentangled Latent Action World Models
May 15, 2026Tianqiu Zhang, Muyang Lyu, Yufan Zhang +2Action-Conditioned Video Generation Latent Action Learning
Peking-Tsinghua Center for Life Sciences, Academy for Advanced Interdisciplinary Studies, IDG/McGovern Institute for Brain Research, Peking University. · Center of Quantitative Biology, Peking University. · School of Psychological and Cognitive Sciences, Key Laboratory of Machine Perception (Ministry of Education), Peking University.
Reconstructing Is Not Acting: Action-Centric Latent Dynamics Modeling
Sep 14, 2026Dingjie Fu, Dianxing Shi, Yangyang Xu +1Latent Action Learning Latent Dynamics Modeling
Harbin Institute of Technology (Shenzhen) · Beihang University
ActionParty: Multi-Subject Action Binding in Generative Video Games
Apr 2, 2026Alexander Pondaven, Ziyi Wu, Igor Gilitschenski +4Video World Models Video Diffusion Models
Snap Research · University of Oxford · University of Toronto +1

May 15, 2026J/K move · Enter open · S save

DiLA: Disentangled Latent Action World Models

Tianqiu Zhang, Muyang Lyu, Yufan Zhang, Fang Fang+1

Peking-Tsinghua Center for Life Sciences, Academy for Advanced Interdisciplinary Studies, IDG/McGovern Institute for Brain Research, Peking University. · Center of Quantitative Biology, Peking University. · School of Psychological and Cognitive Sciences, Key Laboratory of Machine Perception (Ministry of Education), Peking University.

Action-Conditioned Video Generation Latent Action Learning Disentangled Representation Learning World Model Learning Action-Conditioned World Models

Open paper AskPDF arXiv

Latent Action Models (LAMs) enable the learning of world models from unlabeled video by inferring abstract actions between consecutive frames. However, LAMs face a fundamental trade-off between action abstraction and generation fidelity. Existing methods typically circumvent this issue by using two-stage training with pre-trained world models or by limiting predictions to optical flow. In this paper, we introduce DiLA, a novel Disentangled Latent Action world model that aims to resolve this trade-off via content-structure disentanglement. Our key insight is that disentanglement and latent action learning are co-evolving: the predictive bottleneck inherent in latent action learning serves as a driving force for disentanglement, compelling the model to distill spatial layouts into the structure pathway while offloading visual details to a separate content pathway for generation. This synergy yields a continuous, semantically structured latent action space without compromising generative quality. DiLA achieves superior results in video generation quality, action transfer, visual planning, and manifold interpretability. These findings establish DiLA as a unified framework that simultaneously achieves high-level action abstraction and high-fidelity generation, advancing the frontier of self-supervised world model learning.

Abstract

Explore similar work

DiLA: Disentangled Latent Action World Models

Reconstructing Is Not Acting: Action-Centric Latent Dynamics Modeling

ActionParty: Multi-Subject Action Binding in Generative Video Games