While studying world models, I noticed that each new paper introduces a mix of new components and architectures—but the core structure stays surprisingly consistent. Yet there wasn’t an easy way to swap these components in and out, experiment with them independently, or understand how they interact to form a complete world model.
PyTorch-World aims to solve that: it provides a clean, modular framework where you can plug in different components, compare approaches, and learn how world models work from the inside out.
New updates rolling out soon, this is v0.0.3!
You can also install the library from pip https://pypi.org/project/pytorch-world/
Currently this library supports PlaNet world model by Google, Here's the code to train the model in a CartPole-v1 environment:
from world_models.models.planet import Planet
p = Planet(env="CartPole-v1", bit_depth=5, headless=True, max_episode_steps=100, action_repeats=1, results_dir="my_experiment")
p.train(epochs=1)