Introduction
Last week I stumbled upon SmolVLA (Shukor et al. 2025), a small vision-language-action model for robots. The pitch is lovely: most VLAs are billions of parameters and need a datacenter to breathe, and SmolVLA says no, you can train something useful on community-collected data from cheap robot arms and run it on a consumer GPU. I immediately thought, cool I wonder how was LeWorldModel (Maes et al. 2026) would do in this setting. If you don’t remember LeWorldModel you can read my previous post about world models and JEPA. Basically it trains a tiny world model (about 15M parameters, single GPU, a few hours) that stays stable end-to-end using only two loss terms.
Now, a VLA is a policy. You show it a camera image and an instruction, and it tells you what action to take. A world model (Ha and Schmidhuber 2018; Schmidhuber 1990) is a different animal. Give it a state and an action, and it tells you the state you’ll land in. So what if you took SmolVLA’s cheap, affordable, community-data setup, and instead of training a policy, you trained a small world model in the same spirit as LeWorldModel? A “SmolWM”, if you will. Can’t say I dig the name but you know, there are two things that are hard in computer science. Cache invalidation and naming things.
Long story short, I built it. This post is the writeup of the parts that worked well, a cool demo along with the part where I ran into a wall.
From “what do I do” to “what happens next”
A policy learns a mapping from observation to action. Good for control, but it’s a bit of a black box about the world itself. A world model learns the dynamics. Given where you are now and what you do, where do you end up? Write the current state as a latent \(z_t\), the action as \(a_t\), and the model predicts the next latent \(\hat{z}_{t+1}\). That’s it.
The word “latent” is doing a lot of work here and is the cornerstone in a JEPA (Joint-Embedding Predictive Architecture), which means the model predicts in a compact representation space, never in pixels. Why would you throw away the pixels? Because most of what’s in a camera frame is irrelevant to the physics. The exact texture of the table, the lighting, the JPEG noise. If you force a model to predict all of that, it spends its capacity being a video codec. Predict in latent space and it can spend its capacity on the stuff that actually matters, like where the cube is and which way it’s about to slide. There’s a nice writeup by Yann LeCun (LeCun 2022) that talks at length about this. I also built a simple world model from scratch previously. Check it out if you’re curious.
So SmolWM is three small pieces. A little vision transformer that turns a camera image into a latent. A tiny MLP that folds in the robot’s joint positions and a frozen language embedding of the instruction. And a small transformer predictor that takes \((z_t, a_t)\) and coughs up \(\hat{z}_{t+1}\). On the full config it’s about 25M trainable parameters, which is firmly in the “trains on one GPU over lunch” regime.
The whole thing fits on one diagram (below). The top row encodes the current observation into \(z_t\), the predictor pushes it forward under the action to get \(\hat{z}_{t+1}\), and the bottom row encodes the next observation with the same encoder to give the target \(z_{t+1}\). The two losses live on the right.
The two-loss trick, and why JEPAs usually fall over
Here’s the problem with training a thing like this. If your only goal is “make the predicted next latent match the actual next latent”, there’s a stupid solution that wins every time: map every image to the same constant vector. Prediction error zero, model useless. This is called representation collapse, and it’s the ghost that haunts every JEPA. The usual exorcism involves a second “target” network updated with an exponential moving average, stop-gradients, and a small pile of hyperparameters you have to tune by feel.
LeWorldModel’s contribution, which is really quite clever, is a regularizer called SIGReg (Balestriero and LeCun 2025) that replaces all of that machinery with one idea: force the distribution of latents to look like an isotropic Gaussian. If your latents are spread out like a nice round Gaussian ball, they can’t all collapse to a point. And crucially, it does this without an EMA target and without stop-gradients. Two loss terms, one knob.
\[ \mathcal{L} = \underbrace{\lVert \hat{z}_{t+1} - z_{t+1} \rVert^2}_{\text{predict the next latent}} + \lambda \cdot \underbrace{\text{SIGReg}(z)}_{\text{keep latents Gaussian}} \]
The way SIGReg checks “does this look Gaussian” is itself neat. It’s a goodness-of-fit test on the empirical characteristic function along random 1-D projections (an Epps-Pulley statistic, if you want to look it up). I ported it faithfully from their reference code, because when someone hands you a regularizer that collapses six hyperparameters into one, you don’t get creative. You copy it exactly and say thank you.
Ok, but does it actually learn any physics?
This is the question I’m interested in. The prediction loss can look great while the model has learned nothing useful, so I needed a test that a lazy model would fail.
Here’s one way of doing that. Take a real transition, e.g., the robot was here, it did this, it ended up there. Now corrupt it. Keep the starting state, keep the observed outcome, but swap in a different action. If the model genuinely understands that actions cause consequences, its prediction error should shoot up, because the outcome no longer matches the command. If the model is just doing pattern-matching on “the next frame usually looks like this frame”, it won’t notice a thing.
Score every transition by prediction error, label the corrupted ones as “implausible”, and compute the AUROC. Chance is 0.5. Here’s what happens over training on the PushT task:
The wrong_action curve climbing from 0.5 to about 0.94 is the model teaching itself that actions have consequences, with nobody ever labelling a single frame. Nobody told it what “impossible” means. It figured out physics well enough to be surprised when physics gets violated. So yeah, that’s cool.
Catching a floating cube
Right, so if the model can smell an impossible action, can it smell an impossible video? Let’s see shall we?
I took a real clip of an SO100 arm (Knight et al., n.d.) doing a stacking task, and I made a doctored copy where, for a stretch in the middle, I splice in frames from the end of the episode. The arm and cubes suddenly jump to a future state and then snap back. To your eye it’s a glitch, a little teleport. To a world model it should be a physically impossible transition.
Now, the naive thing is to color each frame by the model’s raw prediction error. That doesn’t work, and the reason is instructive. Raw error is high whenever a transition is hard to predict, which includes perfectly real but fast motion. So both the real clip and the doctored clip light up red and you’ve learned nothing. The fix is to look at the excess error: how much more surprised is the model by the doctored frame than by the real one at the same moment? Outside the tampered window the two clips are identical, so the excess is zero and everything stays calm. Inside the window, the doctored physics push the excess up.
And it separates cleanly. Here are the error distributions for real versus doctored transitions:
Why do I care about this beyond it being a fun GIF? Because a world model that reliably flags “that’s physically impossible” is exactly the thing you want as a safety monitor sitting next to a robot policy. Before you execute an action, you can ask the world model whether the state it’s about to walk into is even possible. A policy can’t tell you that. It only knows what to do, and whether the state it lands in even makes sense is outside its job. This is the kind of check I’d want between a learned controller and anything with actual torque.
Can it plan? Here’s where I hit the wall
A world model is supposed to let you plan. You imagine rolling different action sequences forward in latent space, score them, and pick the best. No decoder needed, because you never render the imagined future, you just reason about it in latent space. LeWorldModel reports solving PushT this way and even runs the planning orders of magnitude faster than pixel-space world models. So I wired up a Cross-Entropy Method planner where you sample action sequences, roll the latent forward with the predictor, keep the good ones, repeat, execute the first chunk in the real simulator, and go again.
And, in no way whatsoever, did it solve the task. Coverage barely moved. I spent an embarrassing number of iterations blaming the planner, tuning the search, trying different goals. I finally did the one measurement I probably should have started with.
Planning rolls the latent forward many steps. Every metric I had was one step. So the real question is this: does a free-running multi-step rollout stay faithful to reality, or does it drift off into fantasy after two or three steps? I rolled the predictor eight steps forward with real actions and compared each predicted latent to the encoder’s actual latent.
That flat ratio around 0.23 was kind of a “huh” moment. The model rolls out beautifully. Eight steps, twenty-four frames of imagined future, and the error doesn’t compound. My “the rollouts must be diverging” theory was just wrong. The model was never the problem.
So what was? The goal. I was scoring plans by how close the rollout got to a goal latent, and it turns out the latent of a whole scene is dominated by the agent’s position, because the agent is the biggest thing that moves. So my planner was cheerfully solving “put the agent where it is in the goal picture” while completely ignoring the little T-shaped block it was supposed to be pushing. I tried a coverage reward head instead. It could only explain about 40% of the variance, because “is the block on the target” is a fussy geometric readout the latent doesn’t hand you for free, and the search happily gamed the weak signal.
The thing is, LeWorldModel solves this, and they don’t use a reward head or anything fancy. They plan toward real expert states set onto the actual trajectory, where the goal is always a reachable near-future and the agent’s motion in that goal is the pushing motion. My synthetic goals didn’t have that coupling, so matching the agent didn’t move the block. Reproducing their setup faithfully needs their expert dataset (a 13GB download I wasn’t going to inflict on a metered connection at 11pm in the highlands of Sweden) or their evaluation stack. So I stopped, wrote it down as work-in-progress, and shipped the parts that you’re reading about now.
Conclusion
Here’s where I landed. The world model itself is a small success. It learns robot dynamics from cheap data with no labels, its multi-step rollouts stay accurate out to two dozen frames, and it can look at a video and tell you, unprompted, when the physics have been faked. That last one has a real use as a plausibility monitor sitting between a policy and a robot which is a useful safety layer. Also, it’s cheap to build.
The planning is the asterisk. Everything I measured says the model is capable of supporting it, and I still couldn’t make it solve the task, because getting good reachable goals into a latent planner is its own hard problem that I underestimated. I spent most of my time there blaming the wrong component, which is a very normal way to spend an evening in this line of work, and a good reminder to measure the thing before you theorize about it.
If you’ve done latent-space planning and know the trick for goals that actually couple to the object you care about, I would genuinely love to hear it, because I suspect it’s simpler than I’m making it. And if I’ve butchered anything about JEPAs or SIGReg above, tell me. You can find me in the usual places.