LLMs have suddenly surpassed dedicated robotic models in performing spatial tasks. How might you train LLMs to do this?

To find out what data actually helps spatial learning in LLMs, I experimented to find the best textual representation of a robotic scene, also known as the structured state. I then propose an annotation scheme to generate this structured state at scale from video and build out a VLM pipeline for doing that.

statebench: Paper Coderobolabel: Paper Code

Sometime after it slopcoded your app or analyzed your Hinge conversation, the chatbot developed a stunning ability to operate surveillance drones or use robot arms dextrously. I don’t think any of these fancy eval companies have created GlockBench yet, but I bet it could do that too. What happened?

A humanoid robot drawing The Creation of Adam in simulation on the left, and a real robot arm drawing it on a whiteboard on the right.

Claude Opus 5.5 painting The Creation of Adam in sim vs real. Credit to @thermalpastor

So-called LANGUAGE models are increasingly nailing all sorts of visual and spatial tasks. Robotics, 3D design, image generation and more- all going the way of the transformer architecture. This is because, to the computer, text is undoubtedly the universal interface. Everything that the computer understands is ultimately relayed to it via some kind of text. As such, tasks that we perceive as visual can be distilled into textual formats called structured state representations and trained upon to improve visual decision making.

I could train a spatially-capable LLM if I found out what the best makeup for this structured state is and how to create a ton of them from any ol' video on the web.

The structured state

So what even is this structured state? It should describe the entire scene in compact text. In reading the state at some timestep, a model should have all the information it needs to take an action that produces a desired state in the next timestep.

Some tasks are expressed in file formats compact and legible enough that the raw data itself can be the bulk of the structured state. This becomes less true as the task widens in context; raw ingestion of every pixel value of Mario Kart to prompt a controller input is way less efficient than ingesting a simplified structural state.

A glossy purple blob titled structured state holding seven cream thought bubbles: where is everything?, what can I do next?, what's the goal?, is the arm holding something?, am I touching it?, how many tries so far?, and what did I just do?

Claude generated an SVG for me to visualize the structured state we are about to probe. I wonder how it got so good at generating SVGs?

The experiment

Let's take an easy case; I can probe this ideal structured state by running robotic tasks in a simulation engine.

In MuJoCo, I can get a report of every object's exact position at every step. This tracking data is the bulk of my structured state. I add a few mm of noise to mimic a real tracker, and store positions relative to each other, like where the cube is relative to the gripper.

The state also needs a menu of all possible actions the robot can take at some timestep. These actions can really range in complexity, but let’s first imagine them as fundamental movements like ‘move right’, ‘move down’, ‘move forward’, etc.

Lastly, the goal label: what the task is and what counts as done. Describing intent is a very qualitative observation that can be tough to label at scale. Outside of the state, you also need to find a clever verifier to score progress towards the goal. This is the metric you hillclimb on in an RL environment to teach a model the whole thing.

We can ablate across all these fields to see which hurt performance most in their absence, and at what degree of imprecision they hurt the model when present.

The planning architecture

Casting multimodal tasks as text prediction has always been possible. GPT2 could've taken a list of motor values and predicted something, albeit something bad.

But there have been serious obstacles for a language model predicting motor values at every timestep. One is the huge size of predicted values. A concise structured state helps this, but it's mostly addressed by better models with longer context and higher attention entropy. I'll leave that to the NLP PhDs.

Another problem is that these spatial tasks are way too slow and expensive for an LLM(though Ultrafast is intriguing). Robotic tasks demand real-time reactivity, but the LLM has to generate a decision for every step, or else follow a preset plan between refreshes, all while burning tokens.

The System 1/System 2 paradigm of models breaks this limitation. Jev, the heralded System 1 model, can act as a scout. It quickly reads the state, then picks among the actions on the menu. The slower, wiser System 2 model makes the initial plan and then, like a watchful guardian, ONLY steps in to rewrite the plan when Jev assesses that a step has failed or the task has changed considerably.

Sol writes a plan into the structured state, Jev reads the state and sends the next action to the robot in the scene, the scene feeds back into the structured state, and a dashed arrow returns to Sol only when a step along the plan fails.

An example of System 1/System 2 dynamics. Jev picks each move and checks if it worked, Sol plays the System 2 role of coach.

The System 2 LLM encodes its decree for Jev within the structured state, and Jev reads it at every decision. The LLM can write these directives live at inference-time, or offline to make tons of training labels at scale.

Trimming the fat

Proposed experimental labels often ended up being overshadowed by tracking data, even when accounting for tracking error and noise.

We could distinguish these experimental labels as:

  • Meaning labels: simple yes/no facts computed from the positions: "gripper lined up with the cube", "cube grasped", "cube in the tray"
  • Procedure labels: bookkeeping checks for the task in progress ("is my position estimate out of date?", "how many grab attempts since the last lift?")

These generally went nowhere, but some meaning labels that did help measured "fine contact" (i.e. "Is the pen touching the paper?") where a few millimeters can majorly impact the success of the task. In a real task, these labels would require a precision that outstrips tracking error, so you'd need additional haptic sensors.

The goal

The goal label mattered way more than any other field beyond tracking. You can’t just derive it from some sensor. Someone has to qualitatively evaluate what the task is and when it’s over.

Goals that poorly described ending conditions performed worse than no goal at all. One goal said "Pick up the red cube from the table and place it in the blue receptacle." The verifier, however, wanted the arm to retract when done. A good ending condition needed to be in harmony with the intent of the task AND the verifier. That ending condition also aided in learning from demos with failed actions like mishandled grabs or premature releases.

Left: a goal with no ending, and the arm hangs onto the cube over the tray. Right: the goal ends with “then move the gripper clear,” so it lets go and backs away.

Interestingly, the goal probing revealed the control that the action menu could exert. Even simple tasks without a goal label could be completed with narrow action menus. With an action menu like [approach, grasp, transport, release and retract] a model will look at a scene with one square block and a bin with a blind assumption that the block should be put in the bin. This is a very narrow example(what if the actual goal was something way different?) but it exemplifies how the menu subtly encodes the goal on its own. Such a narrow menu only works for that single task, but perhaps in the future that menu could be dynamically tailored to each task by the System 2 model.

Testing intent

Having the model control a robot to paint is a great test. It’s dextrous, can be graded by looking at it, and it’s fun to watch! With the System 2 model tasked with turning the goal into a discrete plan, that plan manifests in the painting task as a list of strokes to draw.

The Starry Night, The Great Wave and the Mona Lisa beside robot line drawings of each, one column from the prompt 'Draw this image.' and one from a goal with the complete intent, marked with red crosses and green checks.

Some initial drawing attempts

I wanted to play around with how the goal was presented to the model. Robotics has rewarded action-chunking in the past, so it might be good to try breaking down the goal into smaller sub-goals like the individual parts of a painting. Conversely, I also noticed in a lot of computer-use drawing demos that LLMs like to paint in a diffusion-like manner, first filling in the vague blobs of the scene and then iteratively refining at greater detail. I figured it would be good to try breaking the goal down into levels of detail, like first painting larger elements and then adding detail later.

Results of two paintings with different goal structures. ‘One shot’ gives the full plan at once, ‘rounds’ replan at set intervals, ‘parts’ divides the goal by major elements of the image, and coarse-to-fine moves from lowest to highest detail

Take a look! To me, the one-shot illustration performs the best. I found that the plans downstream of broken-down subgoals were badly miscalibrated. The model did not retain an accurate map of where previous strokes actually were, which led to a lot of clutter. Finding a good way to calibrate between subgoals would be vital.

Parts overlapped for a similar reason. With no grounding of the parts within a larger layout, each would drift into an overlapping mess in the center. Having a persistent layout that maps the parts helped a lot.

The System 2 model can create the initial plan from an image of a painting or just a text prompt, like ‘statue of liberty.’ Here you can see the parts overlapping badly without a layout.

Training on real data

It's easy to solve a task in a simulator, but we want this model to work in the real world! The same structured state can teach it real life actions if we can reconstruct the labels from regular videos.

Robust video-to-3D software can form the backbone of our tracking, and since tracking is our most vital signal, making it as accurate as possible would be a huge focus.

Apart from that, we’d need to create an action menu and a well-defined goal from that same video. To do this, I made an extensive annotation tool called robolabel.

Robolabel is a harness for annotating videos cheaply at scale using any LLM with image support. With it, we can reconstruct both the goal and the action menu at once. The tool structures multiple VLM passes to balance accuracy with cost.

  1. The first pass gets fed a small subset of static frames from the video, and key objects are picked out.
  2. The next pass watches the video at 2FPS and splits the clip into broad action tags like (approach, grasp, transport, release, retract).
  3. Then, a crawl pass explores the neighbors of those previous frames to refine the time boundaries for that action stage.
  4. We run these new boundary frames in order through the VLM to describe the progression of the scene: what happened in this stage? What is touching, visible, or held?
  5. These scene descriptions are read as text to generate highly accurate annotations of the goal

Despite several passes, cost is kept low since the model never actually ingests every frame of the video. The accuracy of annotations, however, turns out much higher in this scaffolded workflow than just ingesting the video raw. It’s the best of both worlds!

We get a clearly defined goal, and the descriptions of each action stage make great subgoals. The action menu can just be populated by the broad action tags we generate in the second pass.

We can also denote failed action stages and pair in contact/motor values when a video has them. The experiment includes a cost/speed/quality comparison across several models. Altogether, we can theoretically reconstruct our entire structured state from any video!

Let it rip

We now have our final structured state. Tracking data maps every object in the scene. An action menu encompasses what the model can actually tell the robot to do. A goal with a clear end defines the intent of the task. Additionally, optional contact labels describe haptic data that requires precision beyond the capabilities of the tracking data. Outside of this state lives a verifier that scores how far the model has progressed towards its intended goal.

A glossy purple blob labeled structured state holds four gold pills: tracking data, action menu, goal with a clear end, and contact labels, if needed. An arrow points to a round verifier dial with an orange progress ring, labeled scores progress toward the goal.

Our beautiful, platonically-ideal structured state

This covers our full training setup! We create an RL environment and let an LLM loose in it. We define each scene as a structured state, let the model pick from the action menu, and have the verifier score act as the reward function to maximize during training. With a data flywheel that can churn through whatever set of videos you want, we can train the model at a massive scale. I don’t have a fraction of the money to do actually run this, but I think this is something close to what the major labs are doing right now. This is how LLMs are able to learn incredible spatial skills through text!

A bar chart titled Mean Score Across Ten Tasks, where a pi 0.5 plus GPT 6 Astra system scores far above dedicated robot models, flanked by two astonished cartoon faces.