PaperlayerView exact revision on arXiv

Source-linked AI summary

Latent Particle World Models: Self-supervised Object-centric Stochastic Dynamics Modeling

Tal Daniel, Carl Qi, Dan Haramati, Amir Zadeh, Chuan Li, Aviv Tamar, Deepak Pathak, David Held

arXiv:2603.04553v1cs.LG

TL;DR

Existing video generators are costly for practical decision-making, while conventional image representations often lack semantic object structure. LPWM learns an object-centric latent world model end-to-end from videos, with stochastic per-particle dynamics and conditioning on actions, language, images, and multi-view inputs. It achieves state-of-the-art object-centric video prediction across diverse datasets and demonstrates applicability to imitation learning, although its current scope is limited to recurring scenarios with small camera motion.

  • Problem

    High-fidelity video generators are computationally expensive, and patch-based image representations lack explicit semantic object structure for modeling interactions.

  • Method

    LPWM combines self-supervised latent particles with a learned per-particle latent action module that supports stochastic rollouts and conditioning on actions, language, images, and multi-view inputs.

  • Results

    LPWM achieves state-of-the-art object-centric video prediction on diverse real-world and simulated multi-object datasets and demonstrates imitation-learning applicability.

  • Takeaways & Limitations

    LPWM provides an object-centric world-modeling framework that preserves object permanence, models complex interactions, and produces multimodal rollouts for decision-making tasks.

  • Takeaways & Limitations

    LPWM currently depends on datasets with small camera motion and recurring scenarios and is not yet applicable to general-purpose large-scale video data.

Abstract

from arXiv · show

We introduce Latent Particle World Model (LPWM), a self-supervised object-centric world model scaled to real-world multi-object datasets and applicable in decision-making. LPWM autonomously discovers keypoints, bounding boxes, and object masks directly from video data, enabling it to learn rich scene decompositions without supervision. Our architecture is trained end-to-end purely from videos and supports flexible conditioning on actions, language, and image goals. LPWM models stochastic particle dynamics via a novel latent action module and achieves state-of-the-art results on diverse real-world and synthetic datasets. Beyond stochastic video modeling, LPWM is readily applicable to decision-making, including goal-conditioned imitation learning, as we demonstrate in the paper. Code, data, pre-trained models and video rollouts are available: https://taldatech.github.io/lpwm-web

1 INTRODUCTION

Modern video generators achieve high visual fidelity but remain computationally expensive, motivating world models for decision-making. LPWM addresses this gap with self-supervised object-centric representations, stochastic latent dynamics, and flexible conditioning.

  • High-fidelity video generation requires thousands of GPU hours for training, while diffusion-based inference remains slow and resource-intensive.
  • Image patchification provides scalable representations but does not explicitly encode semantic content or meaningful object interactions.
  • Object-centric representations have shown empirical benefits across video prediction, reinforcement learning, imitation learning, and microscopy.
  • LPWM is a self-supervised object-centric world model trained end-to-end on complex real-world video without explicit particle tracking.
  • LPWM supports conditioning on actions, language, images, and multi-view inputs for video modeling and decision-making applications.

2 RELATED WORK

Related work spans holistic latent world models and unsupervised object-centric approaches using patches, slots, or particles. LPWM is positioned as addressing scalability, temporal consistency, and conditioning limitations in prior methods.

  • General latent world models improve long-horizon prediction through discrete latents, hierarchy, attention, and language conditioning but usually lack explicit object decomposition.
  • Patch-based object-centric methods model local object attributes but rely on post-hoc matching across frames, limiting scalability to complex or real-world videos.
  • Slot-based methods typically separate decomposition and dynamics training, while suffering from inconsistent decompositions and blurry predictions.
  • LPWM advances particle-based modeling toward world models by handling occlusion, appearance, and random movements with action, language, and goal-image conditioning.
  • Prior work has explored latent actions as global variables for controllable or playable environments, often using quantization and two-stage training.

3 BACKGROUND

The background introduces variational and temporal latent-variable models, then describes Deep Latent Particles as a self-supervised object-centric image representation. DLP decomposes scenes into foreground particles and a background particle with disentangled attributes and masks.

  • VAEs maximize the evidence lower bound using an encoder, decoder, and prior, with reconstruction and KL regularization terms.
  • Temporal VAEs condition each timestep’s latent prior on previous latents through a dynamics model to learn coherent video dynamics.
  • DLP models each image as M foreground latent particles plus one background particle in a VAE-based self-supervised representation.
  • Foreground particle attributes encode position, scale, depth, transparency, and local visual features.
  • DLP learns an alpha-channel mask for each particle, enabling pixel-space foreground-background decomposition.
  • Particle notation uses temporal subscripts, particle superscripts, and attribute subscripts such as position p and background bg.

4 LATENT PARTICLE WORLD MODELS (LPWM)

LPWM is an end-to-end latent world model that represents scenes as particles, reconstructs images, and predicts stochastic next states using per-particle latent actions. Its context and dynamics modules support conditioning and multimodal rollouts without explicit particle tracking.

  • Architecture: LPWM encodes frames into particle sets, reconstructs images, and autoregressively predicts future particle states with a jointly trained VAE.The four components are the ENCODER, DECODER, CONTEXT, and DYNAMICS modules.
  • Latent actions: The CONTEXT module assigns each particle a latent action, allowing simultaneous local interactions and multimodal stochastic transitions from identical states.Unlike global latent-action models, this formulation can represent alternatives such as moving left or right from the same state.
  • Latent actions: During training, inverse dynamics infers latent actions from observed transitions, while a learned latent policy regularizes them and supplies samples during inference.The latent policy models action distributions conditioned on the current state and can incorporate global actions, language, or image goals.
  • Dynamics: The DYNAMICS module combines current particles with their corresponding latent actions to predict next-step particle distributions autoregressively.It uses a causal spatio-temporal transformer and conditions particles on latent actions through AdaLN.
  • Architecture: The encoder retains all M particles with their patch-origin identities across timesteps, eliminating explicit particle tracking and enabling parallel frame encoding.Particles can move around their origins, providing an intermediate regime between fixed patches and explicitly tracked object-centric particles.

5 EXPERIMENTS

LPWM is evaluated for unconditional and conditioned video prediction across diverse synthetic and real-world datasets, with ablations and an imitation-learning application. It outperforms baselines on stochastic video metrics and supports compact, object-centric decision-making models.

  • Experimental design: The experimental suite benchmarks unconditional, action-conditioned, and language-conditioned prediction, studies design choices, and evaluates imitation learning on multi-object long-horizon tasks.Image conditioning and multi-view training are additionally demonstrated for goal-conditioned imitation learning.
  • Experimental design: LPWM is tested on diverse real-world and simulated datasets spanning deterministic and stochastic dynamics, varied interaction densities, and multiple domains.The benchmark includes OBJ3D, PHYRE, and Mario among its simulated datasets.
  • Baselines and metrics: The main DVAE baseline matches LPWM’s architecture and parameter count but uses fixed grid patch embeddings without explicit object-centric attributes.Additional comparisons include PlaySlot, G-SWM, SlotFormer/OCVP, and DDLP when applicable.
  • Results: LPWM outperforms all baselines on LPIPS and FVD across stochastic dynamic datasets under varied conditioning settings.It also preserves object permanence, models complex interactions, and produces diverse plausible rollouts from identical initial conditions.
  • Ablations: Ablations show that per-particle latent actions are essential for strong performance, while performance remains robust near the effective particle-dimension action size.The studies also compare global and per-particle actions, latent-action dimensionality, and positional embeddings.
  • Imitation learning: LPWM’s latent actions can support policy learning by mapping predicted per-particle actions to executable global actions.A compact two-layer attention-pooling transformer performs this mapping during goal-conditioned rollout.

Task GCIVL HIQL LPWM (Ours)

LPWM is evaluated for goal-conditioned imitation learning on PandaPush and OGBench-Scene, using a single model per environment and comparing imagined trajectories with executions. It achieves competitive results on several multi-object tasks but not the hardest OGBench tasks.

  • Evaluation setup: Table 3 reports imitation-learning success rates for PandaPush and OGBench-Scene, covering cube manipulation and long-horizon planning with diverse objects.LPWM uses one model per environment across tasks, while PandaPush baselines train separate task-specific policies.
  • PandaPush: On PandaPush, LPWM outperforms all baselines except EC Diffuser and matches it on the 1-cube task.The reported LPWM variant uses multiple camera views simultaneously.
  • OGBench-Scene: On OGBench, LPWM outperforms all baselines on tasks involving up to four atomic behaviors, despite highly suboptimal unstructured play data.All methods fail on task4 and task5 except HIQL, which reaches 20% success on task4.
  • Trajectory comparison: Figure 4 compares LPWM’s imagined goal-conditioned trajectories with actual OGBench-Scene executions, which closely match according to the caption.The trajectories use a learned mapping from latent predictions to actions.

6 CONCLUSION

LPWM advances self-supervised object-centric world modeling for real-world data by decomposing scenes into latent particles and modeling their temporal evolution. It supports diverse conditioning and shows potential for decision-making, while remaining limited to datasets with small camera motion and recurring scenarios.

  • 6 CONCLUSION: LPWM discovers keypoints, bounding boxes, and masks self-supervised, then models scene evolution through latent particles and novel latent action and dynamics modules.The model supports stochastic object-centric video generation from real-world data.
  • 6 CONCLUSION: LPWM supports action, language, image, and multi-view conditioning for object-centric video generation.
  • 6 CONCLUSION: LPWM shows potential for decision-making tasks, including imitation learning.
  • 6 CONCLUSION: LPWM presently depends on datasets with small camera motion and recurring scenarios, such as robotics or video games, and is not applicable to general-purpose large-scale video.

7 ETHICS STATEMENT

The work evaluates its video generation model on simulated and real-world robotics datasets in controlled, non-sensitive environments. It is not demonstrated on general-purpose video or sensitive content, and broader future deployments require attention to potential misuse.

  • 7 ETHICS STATEMENT: The model is evaluated on simulated and real-world robotics datasets, using publicly available or controlled, non-sensitive data.
  • 7 ETHICS STATEMENT: The work is not demonstrated on general-purpose video or applied to sensitive content.
  • 7 ETHICS STATEMENT: Future extensions to broader or less-controlled domains should consider potential misuse and responsible deployment.

8 REPRODUCIBILITY STATEMENT

The authors support reproducibility by providing implementation details, hyperparameters, source code, and pretrained checkpoints. These resources are intended to help others build on the framework.

  • 8 REPRODUCIBILITY STATEMENT: The appendix provides code excerpts, extended implementation details, and the full list of hyperparameters.
  • 8 REPRODUCIBILITY STATEMENT: Source code and pretrained model checkpoints are available to facilitate reuse of the framework.

A.1 LARGE LANGUAGE MODELS (LLMS) ASSISTANCE DISCLOSURE

Large language models assisted with sentence-level writing polish and grammar improvement. The authors reviewed and approved all suggestions.

  • A.1 LARGE LANGUAGE MODELS (LLMS) ASSISTANCE DISCLOSURE: LLMs were used to polish writing and improve grammar at the sentence level.
  • A.1 LARGE LANGUAGE MODELS (LLMS) ASSISTANCE DISCLOSURE: All LLM-generated suggestions were reviewed and approved by the authors.

A.2 PRELIMINARIES: SPATIAL SOFTMAX (SSM) AND SPATIAL TRANSFORMER NETWORK (STN)

Spatial Softmax converts feature-map heatmaps into probability distributions whose moments provide differentiable keypoint coordinates and uncertainty, while STNs extract and recompose localized image regions through differentiable warping.

  • Spatial Softmax: Spatial Softmax applies softmax over heatmap spatial dimensions to obtain a probability distribution for keypoint localization.Each entry represents the probability of a keypoint at a spatial position.
  • Spatial Softmax: The distribution’s expected coordinates and covariance summarize keypoint position and spatial uncertainty.Normalized coordinate grids typically span [-1, 1] across image width and height.
  • Spatial Transformer Network: Spatial Transformer Networks perform differentiable spatial transformations, including translation, scaling, rotation, and general warping.Grid sampling maps output coordinates to input coordinates and uses bilinear interpolation.
  • Spatial Transformer Network: In DLP, differentiable warping extracts particle-specific glimpses during encoding and stitches decoded glimpses back into reconstructed images.Both grid generation and sampling are differentiable and implemented in modern frameworks such as PyTorch.

A.3 DEEP LATENT PARTICLES (DLP)

Deep Latent Particles represent images as foreground particles and a background particle, with disentangled stochastic attributes supporting object-centric scene reconstruction.

  • Particle representation: A foreground latent particle contains disentangled stochastic attributes for position, scale, depth, transparency, and visual features.The particle representation combines explicit attributes with appearance information.
  • Particle attributes: Position encodes normalized (x, y) location, while scale defines particle height and width as bounding-box size.Position uses a Gaussian latent with a prior derived from Spatial Softmax; scale is constrained to [0, 1] with a Sigmoid.
  • Particle attributes: Depth determines relative compositing order but need not represent physical 3D depth in monocular RGB inputs.Its role is ordering decoded objects during reconstruction.
  • Particle attributes: Transparency controls particle visibility continuously through a Beta-distributed latent rather than a Bernoulli presence variable.Values of 0 and 1 represent fully transparent and fully visible particles, respectively.
  • Particle attributes: Visual features encode foreground appearance, while a single background particle represents global background appearance.Foreground and background features are modeled as Gaussian latents.

A.4 LATENT PARTICLE WORLD MODELS - EXTENDED METHOD DETAILS

LPWM is an end-to-end VAE world model that encodes video frames into object-centric particles, models conditioned stochastic transitions with latent actions, and reconstructs future observations.

  • World-model formulation: The model learns a compact self-supervised object-centric latent representation and predicts future observations autoregressively from video sequences.It can operate without conditioning signals and does not require pretrained image tokenization.
  • Architecture: LPWM jointly trains an encoder, decoder, context module, and dynamics module as a variational autoencoder.The encoder produces particles, the decoder reconstructs frames, context samples latent actions, and dynamics predicts next-step particle states.
  • Encoding and decoding: The encoder produces M foreground particles from learned per-patch keypoints plus one background particle, while decoder filtering can reduce them to L ≤ M.Filtering is deferred to decoding to preserve particle identities and avoid explicit tracking downstream.
  • Latent actions: The context module predicts per-particle latent actions from particle sequences and optional conditioning signals such as controls or language.Its inverse-dynamics and policy heads support inference during training and stochastic action sampling at inference.
  • Dynamics and training: The dynamics module predicts next-step particle distributions conditioned on current particles and latent actions, with distribution parameters used in KL training.The temporal ELBO combines static and dynamic terms, while transparency masks inactive-particle contributions.
  • Training objective: Reconstruction uses pixel-wise MSE in simulated environments and perceptual loss on real-world data.The perceptual term uses VGG features, with γ = 0.1 controlling its contribution.

A.4.1 ENCODER Eϕ

LPWM’s encoder builds a hierarchical particle representation from images while preserving particle identities and supporting stochastic, conditioned, multi-view dynamics. Its design combines keypoint proposals, contextual attributes, per-particle latent actions, and decoder-side filtering.

  • Hierarchical encoding: The encoder hierarchically proposes particle locations, infers attributes, and extracts foreground and background appearance features.Keypoint proposals come from a patch-based network with a spatial softmax; attributes include offsets, scales, depths, and transparencies.
  • Hierarchical encoding: Particle positions combine spatial-softmax keypoint proposals with learned offsets that correct proposals toward object centers.The offset correction is important for modeling object dynamics because proposals need not align with object centers.
  • Particle identity: LPWM defers particle filtering to the decoder, retaining all encoded particles and their identities throughout the encoder.This creates a particle-grid regime in which particles move locally around their originating patch centers while preserving positional identity.
  • Attribute encoding: DLPv3-style encoding adds particle attention for depth and residual appearance refinement, while stable transparency parameterization targets more reliable training.Residual appearance modeling lets features incorporate contextual information from all particles.
  • Latent dynamics: LPWM replaces global latent actions with per-particle latent action distributions that support simultaneous interactions and stochastic sampling.The latent policy conditions these distributions on the current state and can incorporate actions, language, or image-based goals.
  • Multi-view conditioning: The model synchronizes particle dynamics across camera views, enabling multi-view training for occlusion-prone decision-making tasks.Pre-trained LPWMs can later support goal-conditioned policy learning from image or language goals.
Loading 2603.04553v1…