Training Object Permanence in World Models

Haotian Zhang, Fengyuan Yu, Dezhi Luo, Haoran Sun, Zehong Zhao, Qingying Gao, Yihan Li, Siyuan An, Huayi Qin, Yilan Zhang, Zhengze Jiang, Pinyuan Feng, Renrui Zhang, Ziyu Guo, Letian Wang, Mengyue Yang, Kangfu Mei, Maijunxian Wang, Ran Ji, Vikash Kumar, Freda Shi, Chandra Sripada, Vincent C. Muller, Philip Torr, Alan Yuille, Nikolaus Kriegeskorte, Felix Juefei-Xu, Lvmin Zhang, Jieneng Chen, Yilun Du, Hokin Deng

WROP provides a cognitively grounded benchmark and training corpus to teach video models object permanence and solidity.

Can we improve video generation models' understanding of object permanence and solidity by training them on a specialized, cognitively-grounded dataset?

Video generation models often fail to maintain object identity during occlusion or respect physical boundaries, causing objects to vanish or pass through solid barriers. These failures suggest a lack of foundational physical reasoning, or "core knowledge," that infants develop within their first half-year of life. The authors introduce WROP (World Reasoning with Object Permanence), a dataset of 150 hand-authored Blender generators that systematically vary physical scenarios while randomizing surface-level visual cues. By fine-tuning a 16B world model on this corpus, they demonstrate that physical reasoning is a trainable capability rather than an emergent property of scale alone. PWM-WROP, the resulting model, ranks first among true-continuation systems in blind human preference studies, performing competitively with frontier commercial models.

Paper Primer

The WROP dataset uses a core-cognition taxonomy to isolate physical reasoning from visual memorization. Each generator creates 10,000 samples by varying structural parameters—like trajectory and occlusion configuration—while independently randomizing surface parameters like lighting and color, forcing the model to learn the underlying physical rule rather than specific visual patterns.

The core mechanism is a targeted fine-tuning regime: the model is trained on a video-to-video contract where it must predict the consequence of a physical event given only the pre-event context. PWM-WROP is a 16B parameter model fine-tuned from a base world model using this corpus, effectively teaching the model to maintain object representations through occlusion and respect solidity constraints during contact.

Fine-tuning on WROP data significantly improves physical reasoning in true-continuation models.

In a blind pairwise Elo study, PWM-WROP achieved an Elo of 1679.5, outperforming the next-best true-continuation model (Grok Imagine) by 224 points. PWM-WROP ranks third overall among 14 models, trailing only two commercial reference-to-video systems.

Why is a new benchmark necessary if existing video models are already considered "world models"?

Existing benchmarks lack coverage of 3D physical inference and rely on multimodal language models as judges, which themselves suffer from core knowledge deficits. WROP provides a principled, human-evaluated framework to test if models actually understand object permanence and solidity.

Does the model's performance depend on its architecture or the training data?

The authors note that because models differ in both architecture and training, the performance gap cannot be attributed to training alone. However, the substantial improvement of PWM-WROP over its base model suggests that targeted training on concept-specific synthetic data is a critical lever for improving physical reasoning.

Introduction and Motivation

We expose why video models miss basic physical priors and outline a new benchmark to fix it.

Current video generation models produce realistic footage but repeatedly violate two core physical priors: objects disappear behind occluders and reappear at impossible locations, and solid objects pass through barriers without deflection.

OP is the expectation that an object continues to exist even when it is out of view.

OS is the expectation that solid objects cannot occupy the same space; they must collide or be deflected.

**Figure 1** **WROP: World Reasoning with Object Permanence.** We introduce a benchmark and training resource for object permanence and solidity in video generation models, built on 150 hand-designed Blender generators across six task families. Each generator randomises lighting, camera angle, speed, and other nuisance parameters while preserving the task’s core cognitive-scientific structure, yielding 10,000 samples per task. We release this 1.5M-sample training corpus, a 300-question evaluation exam with human Elo ratings across 14 video generation models, and PWM-WROP, a fine-tuned continuation model trained on WROP data achieving state-of-the-art performance in its category.

The gap between current video generation and physical reasoning remains large, motivating dedicated benchmarks like WROP.

Related Work

We survey video generation models, reasoning work, and benchmarks, highlighting gaps that VR‑OP&S addresses.

Modern video generation models, built on diffusion and transformer architectures, have produced high‑fidelity, temporally coherent videos. Proprietary systems such as Sora, MovieGen, Veo 3.1, and Kling 2.6, alongside open‑source models like Wan2.2, HunyuanVideo, CogVideoX‑1.5, and LTX‑2, demonstrate comparable capabilities. However, prior work largely targets two‑dimensional settings, focuses on image‑to‑video generation, and lacks a benchmark for structured physical inference about object identity, constraints, and causality.

Recent studies have probed reasoning abilities in these models, showing promise on maze solving, temporal induction, and abstract rule following. Existing benchmarks suffer from limited scale, minimal training splits, and heavy reliance on VLM‑based scoring, which is unreliable for intuitive‑physics evaluation due to known deficits in multimodal language models. Consequently, they fail to assess core physical priors such as object permanence and solidity.

The concepts of object permanence and solidity trace back to Kant’s formal precondition of experience, Wittgenstein’s view of causal intuition, and Piaget’s sensorimotor stage theory. Empirical work, including violation‑of‑expectation paradigms and findings by Falck et al., shows that infants and adults possess early, automatic expectations about object continuity and impenetrability. These priors constitute the most primitive layer of core physical knowledge that any world model should embody.

The WROP Dataset

We detail the WROP benchmark, its task taxonomy, and the pipeline that builds the dataset.

The existing video datasets lack systematic physical priors, making it hard to assess whether models truly understand object permanence or solidity. WROP fills this gap by providing a curated set of tasks that isolate those priors.

The benchmark groups 150 generators into six families that each isolate a single physical principle—either object permanence (OP) or object solidity (OS).

Input clip: frames 1–2 show the ball approaching the wall.

Target clip: frames 3–4 show the ball re‑emerging on the far side of the wall.

The model must generate the hidden motion between frames 2 and 3 despite no visual evidence.

This tiny example shows why the temporal split isolates the reasoning problem: the model cannot rely on continuous visual flow and must maintain a latent object representation.

Why not just use an existing video dataset instead of building WROP?

Existing datasets contain uncontrolled correlations between appearance and motion, allowing models to cheat by memorizing visual patterns. WROP deliberately decouples surface appearance from the underlying physical event, forcing genuine reasoning about permanence and solidity.

The pipeline turns abstract task specifications into concrete video‑prompt pairs while guaranteeing consistency and traceability.

What would happen if the validation step omitted the 60‑frame split check?

Without enforcing the exact split, some samples could contain overlapping or missing frames, breaking the input‑target alignment that the evaluation protocol relies on. Models would then be trained on inconsistent pairs, degrading performance on the intended reasoning tasks.

**Figure 2.** Taxonomy of the six WROP task families. The top row evaluates object permanence through dynamic occlusion, static-scene occlusion, and container-based concealment and movement; the bottom row evaluates object solidity through obstruction, support removal, and collision. Each panel shows representative frames from a sample drawn from one generator in the corresponding task family: an early input frame, the final input frame at the split boundary, and a target frame depicting the expected physical outcome.

**Figure 3** Generator-level composition of WROP’s 150 generators. The inner ring shows 90 OP and 60 OS generators; the outer ring shows the six task families.

**Figure 4** Generator-design example from `Marked_Boxes_Swap` using the current 60-frame input and 60-frame target clips. The sequence shows the objects in marked boxes, the closed boxes after crossing at the split, completion of the position exchange, and the identity-preserving reveal. The task requires the model to maintain the association between each hidden object and its marked container across the input–target boundary rather than infer identity from final screen position.

Evaluation Methodology

We fine‑tune PWM on the WROP corpus and evaluate its physical reasoning against fourteen video models.

Evaluation asks two questions: do current video generators respect object permanence and solidity, and can those principles be learned from a focused dataset.

Fine‑tuning PWM is like a seasoned athlete adding a new drill: the backbone stays the same, but the weights shift to master the specific motion of physical continuity.

How does PWM‑WROP fine‑tuning differ from training a brand‑new video model from scratch?

Fine‑tuning starts from a fully‑trained backbone that already captures generic visual dynamics; the only new learning signal is the WROP contract, so the model only adjusts weights needed for object permanence. Training from scratch would require learning both generic motion and the physical constraints simultaneously, demanding far more data and compute.

Step 1: Load the pretrained Cosmos3‑Nano weights (unchanged).

Step 2: For each sample, feed the first 5 frames as input and compute the loss against the last 5 frames.

Step 3: Perform one SGD update per sample (learning rate = 1e‑4), accumulating gradients over the 4 samples.

Step 4: After processing all 4 samples, the model has completed one epoch; the total wall‑clock time is ≈ 2 seconds on a single GPU.

Step 5: Evaluate on a held‑out clip; the model now predicts the correct continuation of the occluded object in 3 of the 4 cases.

This tiny run shows that the only thing the model needed to change was its mapping from the conditioning frames to the target frames; the rest of the network already knew how to generate plausible motion.

Fourteen video‑to‑video systems are benchmarked: PWM‑WROP plus thirteen external models grouped into three interface classes—true continuation, reference‑to‑video, and edit/transfer.

**Table 1.** The fourteen evaluated models, grouped by provenance, with access route, interface class and native output geometry. $^\dagger$Our model emits a 117-frame clip that replays its 57 conditioning frames before the 60 predicted frames; the replay is trimmed before evaluation. $^\ddagger$Stitched or padded outputs whose source prefix is trimmed before evaluation (Section 4.3).

All models are run through a unified inference harness that supplies only the conditioning video and prompt, never the target, and trims any source replay before evaluation.

Human preference is the primary metric: raters view paired continuations and choose the better one (or a tie) based on alignment, motion naturalness, and physical plausibility, then scores are fit with a Bradley–Terry model.

As a secondary check we compute full‑reference metrics (MSE, PSNR, SSIM, LPIPS, CLIP similarity, FID, etc.) after resampling predictions to the target’s frame count and resolution.

Quantitative Results

PWM‑WROP ranks third overall, leading true‑continuation models and excelling on object‑permanence tasks.

Video generation models often ignore basic physical priors such as object permanence and solidity; this section reports how the fine‑tuned PWM‑WROP model performs on the WROP benchmark.

PWM‑WROP achieves the highest Elo among true‑continuation models, ranking third overall.

Table 2 and Figure 5

Reference‑to‑video models dominate the leaderboard, with Wan 3.0 Prime and MiniMax H3 tying at the top. Their advantage stems from regenerating the entire scene rather than extending the input clip, which gives them more freedom to maintain physical coherence.

Bradley–Terry Elo converts pairwise win‑loss data into a single rating per model, analogous to chess Elo but derived from a probabilistic win model.

How does Bradley–Terry Elo differ from the traditional Elo system used in chess?

Traditional Elo assumes a fixed logistic win probability based on rating differences, whereas Bradley–Terry explicitly models each pairwise comparison and incorporates ties, then rescales the fitted scores to the Elo range for interpretability.

**Figure 5** Human-preference Elo of the 14 models. Points are Bradley–Terry strengths on an Elo scale (mean 1500, dashed line), bars are 95% rater-clustered bootstrap intervals, colour is interface class, and PWM-WROP (filled) is the top true-continuation model.

**Figure 6** Within-family human-preference ranks across the six WROP task families. Each cell shows the Bradley–Terry rank within that family (1 = best, 14 = worst), fitted separately on the 20-rater pairwise judgments for each family (36–96 games per family; 361 total). Color encodes rank from green (top) to red (bottom). The outlined row marks PWM-WROP; overall Elo (right) is from the joint fit over all judgments. Row order follows overall Elo rank; column groups correspond to the task families in Figure 2. Interface class is indicated by the label color, matching Figure 5.

**Figure 9.** Per-family human-preference Elo leaderboards for all six WROP task families. Points are Bradley–Terry MLE strengths on an Elo scale (mean 1500, dashed line); bars are 95% rater-clustered bootstrap intervals. Colour encodes interface class as in Figure 5: filled blue circle is PWM-WROP (ours); open circles are other true-continuation models (blue), reference-to-video models (orange), and edit/transfer models (green). Each panel is sorted independently by within-family rank. Intervals are especially wide in OS families (36–44 total games each); adjacent ranks are rarely distinguishable. Numerical values and within-family ranks are reported in Table 6.

Table 2 lists the full human‑preference leaderboard; the top two positions are occupied by reference‑to‑video models, while PWM‑WROP secures third place with an Elo of 1679.5.

PWM‑WROP’s fine‑tuned performance rivals state‑of‑the‑art models, establishing it as the strongest true‑continuation system on the WROP benchmark.

Qualitative Analysis

Qualitative examples reveal where models succeed or fail on object permanence and solidity.

PWM‑WROP succeeds on all six representative object‑permanence and solidity tasks.

Qualitative analysis of the six cases (G43, G19, G66, G10, G27, G137) shows correct behavior in each.

**Figure 7** Qualitative comparison on three object-permanence task families (Figure 2, top row). Each row shows the same ground-truth input (two key frames, blue), the PWM-WROP target output (two frames, green), and the full baseline output (three frames spanning the generation, red), with $\checkmark/\times$ badges indicating physical correctness. **G43 (OP-1):** Gemini Omni Flash collapses all three balls into a tight cluster at the tunnel exit, violating their lane identities and count, while PWM-WROP keeps each ball in its correct lane with continuous motion. **G19 (OP-2):** Seendance 2.5 generates an oversized occluder over an empty region of the scene while the original objects remain visible, then makes all three objects suddenly appear when the occluder moves away, with no clear continuity from their previous positions; PWM-WROP correctly preserves the scene and reveals the same unchanged configuration. **G66 (OP-3):** Seendance 2.5 correctly animates the 180° turntable rotation but then lifts the cup at the original front position—a different physical cup—rather than tracking the cup containing the ball to its new location; PWM-WROP correctly tracks the ball with the rotating cup and lifts the correct cup.

**Figure 8** Qualitative comparison on three object-solidity task families (Figure 2, bottom row). Layout follows Figure 7: shared input (blue), PWM-WROP target output (green), and full baseline output (red). G10 (OS-1): Seedance 2.5 allows the ball to pass through the first barrier despite its aperture being half the ball diameter; PWM-WROP correctly brings the ball to rest at the barrier. G27 (OS-2): Seedance 2.5 removes the entire support structure rather than only the plate, then leaves the ball suspended in mid-air with no gravitational response; PWM-WROP generates an immediate, correctly-timed fall to the floor. G137 (OS-3): Gemini Omni Flash hallucinates additional balls before the collision, then allows the purple ball to vanish on impact, violating count conservation; PWM-WROP preserves all seven balls with physically plausible diverging trajectories.

The remaining errors fall into two categories: representation dropout, where objects disappear or reappear incorrectly, and causal decoupling, where physical cause‑effect relationships are broken.

Automatic Metrics and Conclusion

PWM‑WROP tops visual similarity metrics and human ranking, highlighting the impact of WROP training.

PWM‑WROP attains the top human Elo rank among 14 video generation models in a blind pairwise study.

Rank 1 out of 14 models, evaluated by 20 raters on the WROP exam.

**Table 3.** Reports full-reference metrics against the target video, computed at 320x192 on each model's trimmed predicted span (Section 4.3). At this resolution, PWM-WROP performs best on perceptual distance (LPIPS 0.081; next best 0.105), structural similarity (MS-SSIM 0.921; next best 0.877), pixel error (MSE, MAE), and final-frame MSE, and ranks second on CLIP similarity. It ranks lower on PSNR and SSIM, and lower still on the resolution-sensitive measures computed at 720p—final-frame LPIPS (0.166) and FID (20.444)—where its 320x192 output is upsampled four times before comparison.

These full‑reference metrics quantify visual similarity to the target video but do not directly assess object permanence or solidity; an edit model can repaint the static scene and achieve high $SSIM$ or $LPIPS$ without correctly handling occluded objects.

Fine‑tuning on the cognitively grounded WROP benchmark yields a model that outperforms commercial systems on perceptual quality despite lower native resolution.

Task Inventory Details

All 150 WROP generators are listed with their physical scenario descriptions.

The appendix enumerates every generator used in the WROP benchmark, grouping them into six physical‑reasoning families.

**Table 4.** Task inventory of all 150 WROP generators across the six task families. Each entry names the generator and its core physical scenario.

OP‑1 (Baillargeonian Occlusion) contains 26 generators such as G02, G04, and G122 that test whether objects remain visible when they pass behind static occluders.

OP‑2 (Object Static Occlusion) comprises 29 generators (e.g., G15, G24, G105) that evaluate persistence of object identity behind stationary screens or doors.

OP‑3 (Container Permanence) lists 35 generators—including G08, G11, and G30—that probe whether objects stay inside containers when the container moves or is opened.

OS‑1 (Baillargeonian Obstruction) defines 20 generators (e.g., G03, G07, G57) that present solid barriers forcing a moving object to stop, deflect, or be blocked.

OS‑2 (Object Drop) enumerates 21 generators such as G01, G27, and G150 that cause objects to fall when support is removed or a trapdoor opens.

Task Inventory Continued

Appendix A enumerates the second half of the task inventory used to probe object permanence and solidity.

The first group of tasks tests whether an object can traverse a ramp and encounter either a solid or transparent barrier without passing through it. Variations include vehicles meeting walls, holes, gates, or doors, and a descending guillotine gate that can stop a moving object. Additional scenarios probe size‑based constraints such as tall vehicles blocked by low beams and small objects passing through apertures while larger ones are blocked.

Table 4 (continued) lists concrete interaction primitives used in the benchmark, pairing IDs with short descriptive names. Examples are “two balls collide and bounce” (newtons cradle) and “glancing oblique collision” (head‑on velocity exchange). These entries form the basis for the Object‑Permanence world‑model training tasks.

The OS‑3 “Object Collision” suite comprises 19 distinct tasks (IDs G44–G143) that cover a range of physical interactions: heavy‑light collisions, off‑center impacts, pendulum strikes, block toppling, and billiard‑style bank shots. Each task describes a specific configuration—e.g., a heavy ball striking a light ball, a ball rolling into a triangular rack, or a ball rolling onto a see‑saw lever. Together they provide a comprehensive probe of object permanence and solidity in simulated physics.

Implementation Details

Implementation details, performance breakdown, and human preference results for PWM on Trainium2.

The released PWM package trains and samples Cosmos3‑Nano on AWS Trainium2 using native PyTorch, avoiding XLA graph tracing. It loads the public diffusers‑layout weights, shards them across a two‑dimensional device mesh of tensor parallelism × FSDP2 via DTensor, and compiles each Mixture‑of‑Transformers block once with static shapes. Training runs with an $fp32$ master copy under FSDP2 while compute uses $bf16$, and checkpoints are written per rank with bit‑exact resumption.

Experiments run on a single trn2.48xlarge instance, exposing sixteen Trainium2 chips as 64 logical NeuronCores linked by NeuronLink. The production layout for the 36‑layer model uses tensor parallelism degree $4$ and FSDP degree $16$ over the 64 cores. Inference runs on the tensor‑parallel group alone ($bf16$, 29 GB weights, 7.3 GB per core), fitting a four‑core trn2.3xlarge.

**Table.** List of tasks categorized by ID, Name, and Description, including sections for OS-1: `Baillargeonian_Obstruction` and OS-2: `Object_Drop`.

The configuration column records the progressive optimisations: the initial run, a rewrite of the reduce‑scatter copy‑in using row concatenation, and a multi‑tensor AdamW with a single synchronisation per step. Correspondingly, tokens processed per second rose from $6.74$ k to $4{,}705$, and finally to $10{,}564$.

At the $5.71$ s step, the block forward and backward passes consume $3.4$ s; the remaining $1.8$ s stems from $75$ FSDP2 all‑gather calls, $1.4$ s from reduce‑scatters, $1.25$ s from the root unit’s all‑gather of embedding, head, and loss, another $1.25$ s from the optimiser, and $0.37$ s from gradient clipping. Reducing the FSDP degree does not improve throughput because the collectives are latency‑bound (200 MB in $24$ ms), so eliminating collectives is the effective lever.

Loading the full $14.6$ GB $fp32$ tensor‑parallel shard before FSDP2 partitioning exhausts device memory, causing the first backward pass to fail. The solution is to interleave loading and sharding, fragmenting the tensor as it is placed on the device.

The default chunk‑concatenation copy‑in consumed $7.9$ s of a $15.9$ s step, making it the dominant bottleneck. Replacing it with a row‑concatenation variant, verified to produce identical outputs, cut the copy‑in time by $2.25\times$.

Parameters replicated across the tensor‑parallel group diverged by $5.9\times10^{-5}$ after $100$ steps. The fix is to explicitly synchronise replicated gradients at every step.

The default AdamW implementation launches roughly ten kernels and performs $22$ host synchronisations per parameter, inflating optimiser time. A multi‑tensor AdamW with a single synchronisation per step reduced optimiser time from $1{,}324$ ms to $126$ ms per step on a rank handling $662$ tensors, while preserving bitwise equality over five steps.

The root unit’s $2.49$ GB $fp32$ reduce buffer cannot accommodate the $151{,}936\times4{,}096$ embedding when it is replicated across the tensor‑parallel group. Two remedies are to shard the vocabulary across the group or to freeze the text embedding during fine‑tuning; this remains an open issue.

A Neuron Kernel Interface flash‑attention path is implemented, but its backward pass is limited to sequences of at most $8{,}192$ tokens, excluding $720p$ geometries. It has not yet been A/B‑tested against the compiled scaled‑dot‑product path, so the default remains the latter.

Each block is compiled once with static shapes; any operator that falls back to the host inflates a sub‑second block to several seconds. To avoid this, input prompts are padded to the training text length, masking vision tokens so every input compiles to a single shape.

Before measuring throughput, the port passed several correctness gates: bitwise forward parity with the upstream CPU reference, two‑layer real‑weight parity of $7\times10^{-7}$, a $64$‑rank bit‑exact resume, and overfitting checks on six upstream clips (sample‑to‑source cosine $0.979$) and on $64$ physics renders (continuation‑half cosine $0.993$–$0.999$ versus $0.85$–$0.995$ for the base model). Sampling agreement with an independent earlier port yielded latent cosine $0.88$–$0.97$ on five prompts, with differences attributable to kernel‑level numerical variation over $35$ sampler steps. A full $64$‑rank checkpoint of the $36$‑layer model occupies $162$ GiB.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers