Can a Robot Predict Future Visual States in Latent Space?¶
Completed
Motivation¶
A robot cannot test every action in the physical world. A world model provides a cheaper space in which to predict future observations before acting.
Problem¶
The model had to retain spatial information from images, learn temporal dynamics, and reconstruct visible future states without training a large visual backbone from scratch.
Architecture¶
Trident-WM follows a Vision–Memory–Decoder design. A frozen DINOv2 backbone and trainable linear neck encode observations. A causal Transformer predicts latent dynamics. A deconvolutional decoder reconstructs images from predicted latents.
Training combines latent-space MSE with pixel reconstruction MSE. Inputs are resized to 224×224 for DINOv2. YAML configurations define CPU, GPU Sprint, and GPU Heavy regimes.
Implementation¶
The model was trained on the LeRobot dataset for 100 epochs using PyTorch Lightning. Weights & Biases retained training and visualization output. Nix, uv, and a Justfile define the development environment and commands.
Results¶
Final validation losses were approximately 0.0027 latent MSE and 0.0023 pixel MSE. Generated sequences preserved coarse geometry of the manipulation target, but became blurry and drifted toward grayscale.
Lessons Learned¶
Low MSE does not imply perceptually sharp futures. The reconstruction artifacts are consistent with averaging under pixel MSE, and the qualitative output is a more useful account of that limitation than the scalar loss alone.
Future Work¶
The documented alternatives are categorical prediction, diffusion, or VQ-VAE-style representations, as used by systems such as DreamerV3, GAIA-1, and IRIS. A second version would need to evaluate one of those objectives against the existing MSE baseline.