A weekend trying to teleport objects in Yann LeCun's world model
I tried editing a JEPA world model's latents at inference. The arm steers, the block doesn't, and the training data explains why.
I recently read LeWorldModel, a 15M-parameter JEPA-based world model trained end-to-end and small enough to run on my laptop. It’s trained on simple simulation tasks like Push-T, Two-Rooms, Reacher, OGBench-Cube - small environments where the model learns to predict its own next-frame embedding given an action. To prevent embedding collapse - the failure mode where the encoder maps everything to the same point, which JEPA models normally fight off with complex mechanisms - LeWM uses a single shiny new regularizer called SIGReg that pushes the latent space toward an isotropic Gaussian: embeddings forced to spread evenly across the whole space. This replaced complicated training tricks with one elegant regularizer, one hyperparameter.
The fact that the LeWM model is so small and elegant let me iterate on it fast from the convenience of my laptop (MacBook Air M4) - and led me down a rabbit hole asking how much directions in its latent are semantically meaningful, and whether they can be controlled at inference. I wanted to see if I could reach into the activations and steer the world model’s output in a controlled way - something like Golden Gate Claude for robotics, where you pick a concept, slide it, and watch the world rearrange itself around it.
I started with Push-T, a simple 2D task where a blue circular arm pushes a gray T-shaped block to a goal pose. Top-down view, minimal physics, no perception noise - the cleanest possible environment for asking whether a world model’s latent space is editable.
It partially worked, less than I expected. Which methods worked, and why, is the rest of the post.
The setup
JEPA is the line of work, championed by Yann LeCun, betting that prediction in embedding space is more useful than next-pixel or next-token prediction - and it’s becoming a leading paradigm for self-supervised learning over continuous signals. The model never reconstructs the world; it predicts what its own internal representation will look like a moment later.
The line has grown fast: it started with I-JEPA on static images, extended to video with V-JEPA and V-JEPA 2, and crossed into world modeling for robot planning with DINO-WM, which uses a frozen DINOv2 encoder and predicts future patch embeddings instead of pixels. It even spread into vertical foundation models such as Brain-JEPA for fMRI and US-JEPA for medical ultrasound. LeJEPA gave the framework a theoretical cleanup, and LeWM is the tiny end-to-end world-model version using this cleanup.
The thing that makes LeWM specifically interesting for editing is SIGReg: instead of stop-gradients, EMA teachers, and contrastive heuristics to prevent collapse, you prove the optimal embedding distribution is an isotropic Gaussian and enforce it directly. One regularizer, one hyperparameter, latent is provably nice. My hunch was that clean, steerable linear directions would follow.
I tried three editing methods that followed the literature in the space, all editing encoder activations directly without retraining the model. I left the predictor alone to keep the scope tight and focus on static features (positions, angle) rather than dynamics:
Linear steering - collect frames where a concept like the arm is at the left vs right side of the screen, take a direction in latent space, add it to the encoder activations.
Activation patching - borrow specific token activations from a “donor” frame and splice them into a receiver’s forward pass at one of the encoder layers.
SAE feature clamping - train a sparse autoencoder (SAE) on top of a single encoder layer output to find interpretable features, then clamp the features tied to the target concept. SAEs were made famous by Anthropic’s mechanistic interpretability research and have recently been applied to vision-language-action models. I most wanted this to work because it requires no reference simulation environment, and discovers concepts in an unsupervised manner - the most elegant path for steerability.
To measure each edit’s effect, I trained five small MLP heads - encoder(frame) → state for arm position, block position, and block angle. These act as verification readouts: if an edit doesn’t move a head’s reading, the latent didn’t actually change on that concept. I held out a test set (all R² ≥ 0.984) before trusting them as a measurement instrument. I also trained a small CNN decoder back to pixels so I could visually check that edits produced coherent images, not just moved numbers.
To make the comparison rigorous, I locked the evaluation: 10 fixed frames where both arm and block sit near the center (where data distribution mostly is). For each frame I measured how much the target concept shifted versus how much everything else leaked - what I called selectivity. A method that moves arm_x and nothing else scores high; one that shifts everything equally scores near 1×. I then plotted a selectivity table for each of the methods attempted.
Two methods steer the arm
Ridge-based linear steering wins consistently. Instead of averaging extremes and subtracting (the textbook recipe), I fit concept = β · z and steered along β. Smooth dose-response, target deltas of 60–340 px, scene stays coherent. Diff-of-means plateaus at ~2× - ridge dominates because ridge uses coefficient information; DOM only uses bin extremes. This contradicts the LLM steering literature, where ITI and CAA both report diff-of-means matching or beating probe directions. My guess is that the 192-d bottleneck and SIGReg’s flattening change the geometry enough that ridge’s regularization helps where it usually hurts.
Activation patching works once you control for angle. When patching tokens for “move arm left,” the donor’s block angle has to match the receiver’s. Without that constraint, angle drift dominates the metric and the method looks broken. With it, arm_x_neg hits 5.49× selectivity. I almost called this method dead before this finding.
SAEs didn’t work for me, for a specific reason
Single-feature SAE clamping failed everywhere for me. I trained 4× expansion SAEs (768 hidden, top-k=64) on 15K encoder activations, then used ridge regression on the SAE’s feature activations to rank which features correlate with each concept. Of the top 5 ridge features for block_y, only 2 actually had block_y as their own dominant concept - the other 3 were really tracking block_x or agent-block distance. The SAE entangles concepts that move together. When I ablated the “block_y features,” other concepts moved more than block_y did. I tried a few extensions - clamping multiple top features at once, suppressing features at individual spatial patches (patch-token SAEs), training higher-expansion SAEs (up to 64×), and selecting features on held-out data per Joseph et al. to avoid selection-on-the-same-data bias - none produced coherent selective editing. Reconstruction quality isn’t the issue: the patch-token SAEs reconstructed cleanly (R² 0.97-0.99).
Huang & Chang reported a decoupling between SAE feature decodability and causal effect on a counting ViT - following their work, I also tried single-feature clamping at multiple sites to see if a different layer fixed it, but the bad results held across the board.
The block doesn’t steer at all
No method I tried produced clean block translation. Every approach that moved block_x also moved arm_x by nearly as much - the selectivity never crossed the green selectivity gate for any of the block directions.
The best explanation I have behind this is that the encoder ties arm movement and block movement together because the training data does. The block only moves when the arm moves. There are no frames where the block moves alone. Hence, no inference-time intervention can recover a separation that the encoder never had to learn. In this case, steerability is data-bound, not method-bound, and the ceiling is set by the arm-block correlation in the training distribution.
What the layers are doing
A surprising finding came from running linear probes across every layer of the model - all 12 ViT encoder blocks, the SIGReg projection, all 6 predictor blocks. A linear probe is just a ridge regression from a layer’s activations to some physical quantity (arm position, block position, etc.). High R² means the concept is linearly decodable at that layer - directly readable as a weighted sum of features, which is also what makes it editable by linear-steering methods. Low R² means it hasn’t been computed yet, or it’s stored in a non-linear form this probe can’t reach.
Joseph et al. reported a “Physics Emergence Zone” in vision world models trained on more complex worlds - physical concepts gradually consolidate, peaking around one third of the network depth.
LeWM tells a different story. The encoder is largely reading block position off the patches rather than building it up slowly - block_x sits at 0.73 R² in the very first ViT block. Arm position and block angle climb more gradually and saturate around L9, closer to the gradual-construction pattern. Even after accounting for randomness (a randomly-initialized ViT-Tiny gets ~0.20 R² on block position from pixel statistics alone), the trained-encoder lift is real and concepts get refined steadily up to the last layers.
Caveat: Joseph’s PEZ is about motion direction in dynamics, not static positions. LeWM’s predictor only has 6 layers and 3 frames of history (vs V-JEPA 2’s 16+) - probably too short for a middle-⅓ peak to show up here either way.
Caveat: Joseph’s PEZ is about motion direction in dynamics, not static positions. LeWM’s predictor only has 6 layers and 3 frames of history (vs V-JEPA 2’s 16+) - probably too short for a middle-⅓ peak to show up here either way. A side experiment running Joseph’s iterative-orthogonalized-probe method on LeWM (peel off one probe direction, retrain on the residual, repeat) finds arm position at enc_z spans ~21 independent directions before R² collapses to chance - same order of magnitude as their ~40-50 dims for motion direction in V-JEPA 2. Seems what lets single-vector steering work here is that the dominant direction happens to be causally clean enough to push on.
As expected, dynamics are different - velocity, speed, and distance-to-block hover near the random-init baseline throughout the encoder, since it only sees one frame at a time. That information emerges inside the predictor, which gets 3 frames of history.
The encoder-to-predictor story looks similar - dynamics R² jumps +0.23 between enc_z and pred_L0 and then plateaus across the rest of the predictor. It peaks almost immediately - though it’s only 6 layers.
One more question I asked at the layer level: does the best layer for reading a concept also happen to be the best layer for editing it? Huang & Chang found these can decouple - in their counting ViT, the layer that’s most causally active for a concept isn’t always the layer with the best decodability. I ran all three editing methods at every encoder layer and measured causal selectivity at each site.
In LeWM, the peaks coincide but the trajectories don’t. Probe R² is already high early (block_x sits at 0.73 in the very first ViT block), while selectivity stays near random-baseline until L9 and then jumps at L10. So decodability is available throughout the encoder, but causal handles for editing only appear at the last few layers. Different shape from Huang & Chang’s counting ViT (where peak decodability and peak causality decouple), but not a clean collocation either. My read: SIGReg pushes the causal handles toward the projector site, while earlier layers carry the decodable signal without the same edit leverage. I didn’t experiment further with this.
What I’d try next
Now that inference-only methods seem to reach their limits, three directions stand out:
Touch the data. Generate Push-T frames where agent and block are placed at independently sampled poses. Fine-tune the encoder on the mix. The arm-block correlation should drop and the encoder be forced to allocate separate latent directions, improving block steering.
Touch the training. SIGReg’s isotropic-Gaussian prior may be working against object-level disentanglement: it spreads variance evenly without giving the model a reason to allocate distinct subspaces to distinct objects. C-JEPA replaces patch-level masking with object-level masking, so the encoder has to infer one object from the others - a counterfactual-style training signal (”what is A given B, C?”) that forces relations between objects into the latent. Worth re-running the steering experiments on a C-JEPA-trained world model and seeing if block finally moves.
Squeeze more out of the frozen model with AI-in-the-loop. I ran most of this through a tight loop - propose a hypothesis from the literature, implement it, let the MLP decoders tell me whether the direction worked. That let me offload a lot of the iteration to coding agents while keeping the measurements honest. With more compute and a tighter loop, the bottleneck shifts from labor to taste.
A great learning experience. Thanks to the LeJEPA / LeWM team for making world models research so accessible!