Building a GAIA-1 Style World Model for Indoor Robot Navigation

What is a World Model?

A world model is a neural network that learns to predict what the world will look like next given a current observation and an action. Instead of directly learning a policy (what action to take), you first learn the dynamics of the environment and then you can plan, imagine counterfactuals, and reason about consequences.

This is the core idea behind GAIA-1 from Wayve, which does this for autonomous driving at scale. I wanted to build something similar for indoor robot navigation. Wayve as a company I think we should be really proud of in the UK, originally starting in Cambridge they are now leading the way with autonomous vehicles and their cars are on our roads in the UK today. Home robotics is an area ripe with research and investment and is one area we are seeing huge advancements.

The best way to achieve my aim here and learn, as always, is to get hands on and try and solve a problem after all “What I cannot create, I do not understand” - Feynman.

The Setup

First things first I set about buying and building a raspberry PI controlled car. This would work as my data collection engine and through tele-operation, I would collect information about my home. It collects data in RGB frames at ~10fps paired with motor commands (left/right wheel speeds). Each trajectory is a few hundred frames of the robot navigating through rooms.

Did some tinkering and got the rig together, I then was able to SSH into the car and upload some .py files for it to take commands and send visuals from its camera back using pyzmq.

The pipeline has two stages:

Stage 1 — DINO VQ-VAE: Compress each 128×128 image frame into 64 discrete tokens using a Vector Quantised VAE with a DINO distillation loss. This gives semantically meaningful tokens — similar scene regions get similar token IDs. These tokens become useful as my Transformer that I want to train can use these tokens as a representation of say a couch or a rug as a means to view how it changes with the respective robot actions. We know through GPT models how powerful these models can use language tokens to help users and perform coding tasks etc. Instead of just minimizing pixel-wise error (which leads to blurry averages), the DINO distillation forces the VQ-VAE to group pixels into semantic clusters. It learns that a ‘rug’ is a single entity, making the Transformer’s job of moving the rug with actions much easier.

Stage 2 — GAIA Transformer: A causal GPT-style transformer trained on sequences of [frame_tokens, action, frame_tokens, action, ...]. Given N context frames and an action, it predicts the next frame’s tokens autoregressively. The transformer here is trained on the trajectories of these compressed latent representations of my home and learns how to recreate these latents frame by frame so when I feed the latent back into the decoder portion of my above mentioned VQ-VAE it should be able to build the images back up!

You could say I am creating a mini ‘world’ model of my own home here….we have Wayve at home!

Results

The model learned real spatial structure from trajectory data alone there is no 3D supervision, no depth sensors. Applying TURN_LEFT causes the scene to pan right; applying TURN_RIGHT enough steps eventually reveals a completely different part of the room.

Most impressively, the model maintains “object permanence.” If it pans away from the couch and then turns back, the couch reappears in the correct relative position. It has internalized the room’s geometry.

Below is model generated.

What’s Next

• Better action conditioning (continuous motor values instead of 5 discrete actions)
• Longer rollouts with less drift
• Using the world model for planning and policy learning