Skip to content

Projects

Six hands-on projects build a complete world-model pipeline from scratch. Work through them in order: the encoder from P01 becomes the observation encoder in P02, the dynamics model from P02 becomes the backbone in P03 and the baseline in P04, the two trained systems from P03 and P04 are compared in P05, and P06 probes those same systems for causal fidelity. Each project is a notebook-first tutorial that runs on CPU, GPU, or TPU, uses only synthetic data, and passes a checkpoint to the next stage.

Hardware requirements

Every notebook in this section was developed and run on Google Colab with a single T4 GPU. If you have comparable or better compute, an Nvidia GPU, an AMD GPU, or a TPU from the same or a later generation, you can run all five projects without any changes. On Colab, we recommend subscribing to Colab Pro, because the free tier only occasionally has an idle T4 available, so a paid plan gives you the reliable access these projects assume.

Markdown pages only include narrative text and code. Any outputs, plots, tables, or other artifacts live in the corresponding .ipynb notebook files.

Open any notebook in Jupyter or Colab and run it top to bottom. If an upstream checkpoint is missing, the notebook falls back to random initialization so it still works as a smoke test, but the cross-project comparisons only become meaningful once the real checkpoints are present.

Project sequence

#ProjectPrerequisiteSavesDeliverable
P01Train a VAE EncoderL02 Part Avae_encoder.ptCNN VAE on 64×64 frames; ELBO loss curve; latent traversals showing disentangled dimensions
P02Build an RSSM Dynamics ModelP01, L02 Part Brssm.ptGRU, MDN-RNN, and RSSM compared; rollout plots; 1-step to 5-step prediction error curves
P03Train a Dreamer AgentP02, L03 Part Bdreamer.ptEncoder + RSSM + latent Actor-Critic training loop; reward curve; FID and reward-correlation self-evaluation
P04Swap the Dynamics BackboneP02, L03 Part Atransformer_wm.ptRSSM replaced by a STORM-style categorical VAE plus causal Transformer; architecture comparison report
P05World Model Evaluation DashboardP03, P04, L04--Both trained models loaded and scored side by side: PSNR, reward correlation, token loss, and latent drift
P06Counterfactual Action-Conditioned World ModelP03, P04causal_wm.ptPearl-ladder analysis: interventional and counterfactual rollouts, an inverse-dynamics-regularized world model, and an action-influence metric

How the checkpoints chain together

The projects share a single set of weight files passed forward through the pipeline. P01 trains the VAE and writes vae_encoder.pt. P02 loads that encoder, trains the dynamics models, and writes rssm.pt. From there the path forks: P03 combines the encoder and RSSM into a Dreamer agent saved as dreamer.pt, while P04 reuses the RSSM as a baseline and trains a Transformer backbone saved as transformer_wm.pt. P05 loads both dreamer.pt and transformer_wm.pt for the accuracy evaluation. P06 then loads the same two checkpoints to probe causal fidelity, training one action-regularized model of its own that it saves as causal_wm.pt.