Source-linked AI summary
LeCuration: A Tiny World Model as a Data Curation Multi-Tool
Mayank Sengupta, Nirmit Desai, Eric Song, Kunal Sawarkar
TL;DR
Physical-AI data curation needs ways to characterize datasets beyond hand filtering and rule-based filters. LeCuration trains a small predictive world model for curation and uses its embeddings to cluster gameplay and flag prediction surprises. The CS:GO case study is qualitative, with downstream training results and quantitative curation metrics still unreported.
Problem
Physical-AI data curation needs a small, quickly trainable model that characterizes typical, rare, and surprising data, since hand filtering does not scale and rule-based filters may require inaccessible physics-engine state.
Source: §1 · p. 1Method
LeCuration adapts a small predictive world model to CS:GO gameplay, using its embeddings and predictions for curation and a diffusion decoder for visual inspection.
Source: §1 · p. 1, §1.2 · p. 3Results
The embeddings organize gameplay into recognizable content clusters, and prediction divergence supplies a candidate signal for surprising frames.
Source: §3.1 · p. 5, §4 · p. 7Takeaways & Limitations
The study demonstrates qualitative curation primitives on CS:GO gameplay, including content-based clustering and anomaly signaling from embedding divergence.
Source: §3.1 · p. 5, §4 · p. 7The model detects visual surprise rather than physics-specific implausibility, and downstream training benefits and broader real-world data remain untested.
Source: §3.3 · p. 6, §5 · p. 8
Abstract
from arXiv · showhide
Many applications of physical AI run within finite or closed physical worlds with a limited set of physical laws governing object behavior. Examples include robots working in a warehouse and agents moving around in a video game. In order to better organize, filter, and curate data for physical AI applications, we propose a new approach centered on the unique settings and physical laws of individual datasets. We train LeCuration, a small world model intended to serve as a data curation tool for a separate, larger downstream model. To build this model, we choose LeWorldModel (LeWM)as our latent encoder and predictor, adding a diffusion transformer (DiT) decoder to add visuals to autoregressive gameplay rollout. We find that the embeddings of this model can be used as an anomaly detection signal and as a content-based clustering heuristic, and that auto-regressively predicting the game state with this model allows us to qualitatively check for action-state consistency. This paper presents a qualitative, proof-of-concept case study on CS:GO gameplay data; we do not yet report quantitative curation metrics or downstream training results, which we identify as the key next step.
1 Introduction
The paper motivates a small, quickly trainable model for characterizing datasets and adapting world-model representations to data curation rather than planning. LeCuration demonstrates qualitative clustering, anomaly detection, and visualized autoregressive prediction on CS:GO gameplay.
1 Introduction
- Data curation needs a small, quickly trainable model that can characterize what is typical, rare, or surprising without relying on hand filtering or hidden physics-engine state.The authors argue that the effect of curation cannot be understood without training a model on the curated data.Source: §1 · p. 1
- The proof of concept is limited to qualitative validation on one CS:GO map, with no quantitative curation metrics or downstream-model training results yet.The authors identify downstream training results and quantitative evaluation as open next steps.Source: §1 · p. 2
1.1 Background and Prior Work
- Unlike prior JEPA uses for planning or benchmark evaluation, this paper applies learned representations to analyze and curate a separate model’s training corpus.The paper describes prior work as using the JEPA model as the end consumer of its representations.Source: §1.1 · p. 2
1.2 Main Contribution
- LeCuration adapts LeWorldModel as a compact curation tool for a larger downstream model, rather than as a planner or control-benchmark model.The encoder has roughly 2M parameters and the predictor roughly 13M; the model is trained on CS:GO recordings.Source: §1 · p. 1, §1.2 · p. 2
- Its contributions include qualitatively coherent gameplay clusters, an embedding-divergence anomaly signal, and autoregressive gameplay-frame reconstruction from predicted world states.The diffusion transformer decoder conditions on predicted world states to generate TAESD latents for gameplay frames.Source: §1.2 · p. 3, Figure 1 · p. 3
2 Architecture and Training
LeCuration adapts a predictive world model to TAESD latents, using a latent encoder and action-conditioned predictor to support dataset analysis. A separately trained diffusion decoder visualizes embeddings, while self-supervised training and autoregressive post-training prepare the model for CS:GO data.
2 Architecture and Training
- LeCuration adapts the LeWorldModel architecture and training recipe for data curation, with modifications intended to improve training speed and fit that use.The model operates on TAESD latents rather than encoding raw pixels directly.Source: §2 · p. 3, §2.1 · p. 3
2.1 Stable Diffusion/TAESD Latent Space
- TAESD compresses each 224 × 224 × 3 frame to a 4×28×28 latent, paired with a 51-dimensional one-hot action vector representing keys pressed.The action vector includes key presses and mouse movements.Source: §2.1 · p. 3
2.2 Latent Encoder and Predictor
- A patch-ViT encoder maps each frame’s TAESD latent to a 196-dimensional CLS world state, and a causal predictor uses six prior embeddings plus action embeddings to predict the next state.The prediction-to-observation gap is the primary curation signal.Source: §2.2 · pp. 3–4
2.3 Diffusion Transformer Decoder
- A separately trained DiT decoder denoises TAESD latents conditioned on frozen CLS embeddings so learned representations can be inspected visually.The decoder is trained after JEPA pre-training and does not participate in the curation pipeline.Source: §2.3 · p. 4
2.4 Training Recipe
- Training pre-encodes frames once, compares predictions with stop-gradient EMA-encoder targets, and uses pre-training followed by 50-frame autoregressive post-training.Post-training emphasizes later frames with a discount factor of 1.03; the reported pre-training and post-training durations are 120 and 60 minutes, respectively.Source: §2.4 · p. 4
2.5 Dataset and Compute
- Experiments use unfiltered CS:GO recordings from varied skill levels on de_dust2, sampled at 16Hz for frames and key presses.The experiments ran on a 40GB Nvidia A100 GPU accessed through Google Cloud Platform.Source: §2.5 · p. 5
3 Evaluation: Embedding and Predictor Quality
The evaluation qualitatively finds that embeddings organize gameplay by semantic content and support retrieval and curation, while coverage- and rarity-based selections remain unvalidated downstream.
3.1 Sorting the Dataset by Clustering Embeddings
- One million clustered training-corpus embeddings form recognizable gameplay situations without semantic labels, with rows showing coherent clusters.The reported examples include open sightlines, close-quarter interiors, smoke-obscured views, and player silhouettes at doorframes.Source: §3.1 · p. 5, Figure 2 · p. 6
- The cluster structure supports rebalancing overrepresented situations and filtering loading screens, spectator views, scope views, or heavily compressed frames without per-frame labels.The figure’s first row illustrates anomalous clusters that can be sifted out without searching for them specifically.Source: §3.1 · p. 5, Figure 2 · p. 6
3.2 Auto-regressive Nearest-Neighbor Retrieval
- Nearest-neighbor-snapped rollouts match ground-truth scene context, player orientation, and lighting, supporting embedding coherence and a 1M-frame retrieval corpus.The authors also suggest selecting training frames near a target embedding without running the DiT.Source: §3.2 · p. 5
- Retrieval exposes weapon misrepresentation and flicker, plausibly tied to the weapon’s small screen footprint and rare firing or switching actions; proposed finetuning would oversample those actions.The authors suggest this as a way to improve accuracy in that aspect for a larger model.Source: §3.2 · p. 5
3.3 Using Clusters for Coverage-Based Curation
- Farthest-first episode-centroid selection maximizes spread in 196 dimensions, but downstream quality at reduced training budgets remains untested.The selected subsets have not yet been used to train a downstream world model.Source: §3.3 · p. 6
- Rare-trajectory upsampling targets episodes in sparse embedding regions, but its hypothesized disproportionate training signal remains untested.The authors plan comparisons with random-sampling and full-corpus baselines in future work.Source: §3.3 · p. 6
4 Predictor Quality and Anomaly Detection via Embedding Divergence
During autoregressive rollout, embedding divergence is treated as a candidate anomaly signal: spikes can accompany surprising events, but may have multiple causes and are not yet a validated curation rule.
4 Predictor Quality and Anomaly Detection via Embedding Divergence
- At each rollout step, the Euclidean distance between the predicted and ground-truth encoded world states, d_t = ||ê_t − e_t||_2, tracks embedding divergence.The rollout starts from six ground-truth frames; predicted states are decoded with the DiT for visualization.Source: §4 · pp. 6–7, Figure 3 · p. 7
- A spike can reflect surprise, accumulated rollout error, or out-of-distribution actions; no quantitative threshold or selection metric is established.The authors treat d_t as a candidate signal for further investigation, not a validated curation rule.Source: §4 · p. 7
- The predictor’s world state diverges during a rare or confusing action sequence, with the scene teleporting elsewhere on the map before returning after about 350 frames.The divergence period is visible in the Figure 3 graph.Source: §4 · p. 7, Figure 3 · p. 7
- High-d_t frames are candidate curation targets for exclusion when degenerate or oversampling when they represent rare, physically meaningful interactions.The signal is derived from the gap between predicted and observed embeddings.Source: §4 · p. 7
- The signal is not physics-specific: unusual paths, abrupt recording cuts, and wall collisions can all produce spikes.Separating physics-caused divergence from other causes is left for future work and would require pretraining on a wide real-world corpus.Source: §4 · p. 7
5 Discussion and Limitations
LeCuration is presented as a small, domain-trained model for curating data for a larger downstream model. Its embeddings and predictions offer qualitative clustering and anomaly signals, but the model captures visual appearance and motion continuity rather than physical plausibility.
5 Discussion and Limitations
- LeCuration’s embeddings organize gameplay into content-based clusters without labels, supporting the paper’s claim that a small world model can aid data curation.The model is described as a small tool for a separate, larger downstream model.Source: §5 · p. 7
- Autoregressive embedding divergence provides a learned, dataset-specific anomaly signal for identifying frames to exclude or oversample.Oversampling may help a downstream model capture rare, physically meaningful interactions, but that use remains qualitative.Source: §5 · pp. 7–8
- Embedding extraction and anomaly scoring require forward passes that can be cached and reused across curation experiments without retraining.The paper describes this cost as negligible compared with training the downstream world model.Source: §5 · p. 8
- The main limitation is that the embeddings and predictions encode appearance and motion continuity, not physical plausibility.The authors also want to test the approach on larger, more varied data, preferably real-world datasets.Source: §5 · p. 8
A Comparing Nearest-Neighbor Retrieval and DiT Output
Nearest-neighbor retrieval can make autoregressive rollouts more interpretable than DiT-rendered videos in this dataset. The authors recommend combining both visualization approaches for future work, especially with more varied data.
A Comparing Nearest-Neighbor Retrieval and DiT Output
- Nearest-neighbor retrieval produced more interpretable rollout videos, mainly because retrieved dataset frames were clearer than DiT-generated frames.The DiT was designed as a visualization tool separate from the encoder and predictor; it was not necessary once the model had seen enough data.Source: A Comparing Nearest-Neighbor Retrieval… · p. 8
- A passing player that obscured much of the screen could cause divergence, while the DiT struggled to depict such scenes because they were rare in its training data.Source: A Comparing Nearest-Neighbor Retrieval… · p. 8
- Nearest-neighbor retrieval worked well here because the footage came from the relatively small, closed de_dust2 map; the authors recommend pairing it with a DiT for more varied data.Source: A Comparing Nearest-Neighbor Retrieval… · p. 9
B Code and Data
The project code is available, but the exact dataset is not public; the authors point to a similar public Counter-Strike 2 dataset in the same format.
B Code and Data
- The project code is available, but the dataset supplied by an anonymous data partner is not public; a similar Counter-Strike 2 dataset in the same format is public.Source: B Code and Data · p. 9