chandar-lab/semantic-wm

This is research code from a paper titled "Reconstruction or Semantics? What Makes a Latent Space Useful for Robotic World Models." A "world model" in this context is an AI system that can predict what will happen next in a video, imagine a robot arm watching itself work, then simulating future frames to plan its next move. The central question the project investigates: when building such a world model, should you represent video frames using a space optimized for visual reconstruction (making the image look right) or semantic meaning (capturing what is actually happening)? The paper concludes that semantic encoders, which understand meaning rather than just pixels, better preserve information about robot actions, task progress, and downstream behavior, even if the reconstructed images are less sharp on standard visual quality measures. The code trains what is called a "latent diffusion world model", a type of generative AI that learns to predict future robot video frames step-by-step from noise, conditioned on what actions the robot takes. It supports multiple different encoder types for compressing video frames and both single-camera and multi-camera robot setups. This codebase is intended for AI researchers working on robot learning and video generation. Setup requires Python with specific deep learning dependencies. Users can download a robot demonstration dataset, train a compression adapter, train the world model itself, and then run evaluation using several metrics for visual quality, spatial accuracy, and how well the model responds to action inputs.