Trident-WM Visual World Model¶
Completed
Overview¶
Trident-WM is an experimental world model I developed to investigate visual prediction in latent space.
The objective was to learn temporal relationships between video observations and predict future frames without training a large visual encoder from scratch.
I used a Vision–Memory–Decoder architecture combining a pretrained visual backbone, a causal Transformer, and an image reconstruction decoder.
Implementation¶
The architecture consists of three components:
- Vision: A frozen DINOv2 backbone extracts visual features from 224×224 images. A trainable linear layer projects these features into the model's latent representation.
- Memory: A causal Transformer processes sequences of visual embeddings to predict future latent states.
- Decoder: A deconvolutional network reconstructs images from the predicted latent representations.
The Transformer operates on visual observations without conditioning on robot actions.
For future prediction, the model receives an initial sequence of observations and generates subsequent predictions autoregressively, using previous predictions to continue the sequence.
Training combines two mean squared error (MSE) objectives: one measuring the difference between predicted and target latent representations, and another measuring pixel reconstruction error.
I implemented the training pipeline using PyTorch Lightning and trained the model for 100 epochs on a LeRobot dataset. Weights & Biases was used to track training metrics and visualize the generated sequences.
Results and Limitations¶
The final recorded validation losses were:
- Latent MSE: approximately
0.0027. - Pixel reconstruction MSE: approximately
0.0023.
The generated sequences preserved some coarse spatial structure of the manipulation scenes, but the reconstructed images were blurry and gradually lost color, tending toward grayscale.
These artifacts are consistent with limitations of pixel-based MSE reconstruction, although I did not isolate the contribution of the loss function from the architecture itself.
The model could generate future visual observations, but reconstruction quality degraded over the predicted sequences.
Because the model was not conditioned on robot actions, it could not directly evaluate the consequences of different actions. It was therefore not suitable for action-conditioned planning in its implemented form.
Trident-WM remained an experimental research project and was not deployed on a physical robot.