A world action model in a sandbox

In September Black Forest Labs released FLUX 3 Action, a 7-billion-parameter world action model. It takes camera frames, a robot's joint state and a text instruction, and produces the robot's next few dozen commands while also predicting the next video frames. I read through the training code and wanted to understand it properly. You can't do that with a 7B model on a 16 GB MacBook, so me and my trusted assistant Claude built a toy version from scratch, with every part shrunk until it trains on a laptop.

The code is on GitHub: mgrue/world-action-model.

The sandbox

The robot is a car seen from above. toycar is a small deterministic 2D simulator: an oval track, a 64 × 64 overhead camera, and a car with a yellow nose so you can tell which way it's pointing. A seed reproduces an episode exactly.

The one interesting feature of the track is the Fork. The top straight splits into an outer and an inner branch with a grass island between them, and the two rejoin further on. A scripted pure-pursuit driver, the expert, follows one of three Instructions: outer, inner, or any. Under any it picks a branch at random.

That fork is why the project exists. Under any, both "steer left" and "steer right" are correct. A plain regression network, trained to minimise squared error, learns the average of all correct answers, and the average of left and right is straight ahead into the island. Generative models trained with diffusion or flow matching don't average. They sample one answer, so they should pick a branch and commit to it. FLUX 3 Action generates its actions this way for exactly this reason, and the toy is built to reproduce the effect.

The model

The pipeline follows the FLUX design, scaled down by roughly a factor of a thousand:

• A frame compressor. A 3.2 M-parameter convolutional VAE (vae-v1) squeezes each 64 × 64 frame into an 8 × 8 × 16 latent. The car covers about 0.7 % of the picture. On the first try the VAE learned to paint an empty track and score well anyway. Weighting car pixels 31 times higher fixed that. The VAE only passes if position and heading can still be read from the reconstruction: 95 % of frames within 0.3 m and 10°.
• A transformer. 8.2 M parameters, 8 blocks with parallel attention and SwiGLU, adaLN conditioning on the noise level, and RoPE over time and grid position. Each example is one row of 162 tokens: 16 for the current frame, 1 for the car's state, 1 for the instruction, 128 for eight imagined future frames, and 16 for the actions. The first 18 are given. The last 144 start out as noise.
• Flow matching. In training, the true future and the true actions are mixed with noise at a random level, and the network learns which way leads back to clean data. To sample, it starts from pure noise and takes a few Euler steps. Actions and future frames are generated together in one pass, and action k shares a time position with frame k, so the plan and the imagined video stay consistent.

Training on about a million windows of expert driving took a little over three hours on an M1 Pro. On held-out data wam-v1 came close to the expert: steering error 0.06 on a −1…1 scale. At the fork, 99 % of its sampled plans clearly commit to one branch. Under any, 47 % go outer, and almost every situation gets samples going both ways. As a side effect you can decode the imagined future frames and watch what the model expects to happen.

Letting it drive

None of that means the model can drive. Grading on recorded data always places the car where the expert put it. The real test is a closed loop: the model's actions move the car, the simulator draws the next frame, and the model has to cope with wherever it ended up.

Each tick, the model gets the frame, the state and the instruction, and draws four candidate plans of 16 actions each. Only the first action is executed before it replans. Chunks are never averaged, because averaging two plans is the regression failure again. Instead, the loop keeps the candidate that best continues the plan it was already following. Without that, the car can change its mind about the branch from one tick to the next.

Each driver plays the same 120 unseen episodes of 60 seconds, about four laps, and has to pass four gates: at most 5 % crashes, at least 3.5 laps on average, the instructed branch taken in at least 95 % of episodes, and a 30–70 % outer/inner split under any.

Before the model got anywhere near the gates, I checked the gates themselves. Each saboteur is the expert with one deliberate defect: one never steers, one is slow, one takes the wrong branch, one always goes outer under any, and one averages the outer and inner steering at the fork. Each one has to fail exactly the gate it breaks. The averaging saboteur is the regression failure in miniature: every one of its 40 any episodes ends in the island on the first pass.

The fork works, the corner doesn't

The result for wam-v1:

Everything the model was built to show works. It takes the instructed branch (2 wrong out of 380 passages), splits any evenly, and never hits the island. Only 0.5 % of its plans near the fork point between the branches. Episodes without a crash run at full expert pace.

It still crashes in almost a third of the episodes, and nearly all the crashes are in the same place: 35 of 37 happen driving clockwise, at the exit of the left-hand bend just before the fork. The car ends up a little further toward the outside of that bend on each lap, until one lap it leaves the road.

This is compounding error, the classic weakness of imitation learning, and here it can be measured. When the car is within 0.3 m of the expert's line, the model's steering differs from what the expert would do in the same spot by 0.19. At 0.6–1 m off the line, the gap grows to 0.50. The expert hardly ever drives that far off its own line, so the model has seen very few examples of getting back. On average the model drives 0.29 m off the line, against 0.01 m for the expert.

So the generative part works, and the imitation part fails the way imitation is known to fail. The fix isn't a bigger network. The model needs to see recoveries: expert demonstrations with injected noise (DART) or expert corrections of the model's own mistakes (DAgger). That's the next model, and it will be graded by the same loop, unchanged.

How it was built

The project is planned in eleven steps, each with something that can be checked before moving on. Every step got its own specification, built up from small decision tickets. Each ticket had one question ("what happens when the car leaves the road?", "how do chunked policies replan in practice?") and was resolved by research or a throwaway prototype that measured something. I did the ticket work together with Claude, and the full trail is in the repo under .scratch/. Thresholds were fixed before any model was trained, and every frozen checkpoint records the exact code, data and recipe it came from. This step-by-step routine is why a 30 % crash rate counts as a clear result here and not a mystery.

Steps 1 to 9 took a few days. Still to come are the experiments that motivated the project: the actual regression baseline (does it really drive into the island?), whether predicting future frames helps the actions at all, classifier-free guidance, chunk length, and swapping the VAE for a JEPA encoder that predicts the meaning of future frames instead of their pixels. And, more practically, a car that makes it round the left-hand bend.