Explore compact PyTorch reference implementations for the core mechanics of a world model: latent state, learned dynamics, action conditioning, imagined rollouts and planning.
Start small. First verify one-step state prediction. Then test repeated rollouts. Add planning only after the learned dynamics are useful for the target environment.