Skip to content

pytorch-world

pytorch-world

World models for PyTorch. Learn how a world evolves, then plan by imagining what comes next. Shared parts keep every model short and easy to compare.

uv add torch_world

Get started Models API

A world model learns how an environment evolves: from an observation and an action, it predicts what happens next. With one, an agent can plan by imagining the outcomes of its actions before taking them.

torch_world gives world-model research shared parts. Models are organized into families that share a way of learning, and every model is built from the same layers, objectives, planners, environments and datasets. The first family is JEPA, which predicts the future in a learned representation space instead of in pixels. More families and models are on the way.

  • Model families

    Models grouped by how they learn. JEPA comes first, with seven models from image pretraining to multi-robot control. More families are coming.

    Models

  • One world-model interface

    encode, predict and rollout are the same for every world model, so they can be trained, compared and planned with interchangeably.

    Concepts

  • Planning

    Search for actions inside a model's imagination. The cross-entropy method works with any world model, with no extra code.

    Planning guide

  • Environments

    Batched pixel worlds that render on the GPU, and MuJoCo pick-and-place with UR10e and SO-101 arms.

    Environments

  • Datasets

    Windows over collected experience for training dynamics, plus medical images and questions for vision-language models.

    Datasets

  • Building blocks

    Vision transformers for images and video, conditioned predictors, masking strategies and objectives, shared by every family.

    API reference

Thirty seconds

from torch_world.envs import PointMass
from torch_world.models import LeWM
from torch_world.nn import CLSEncoder
from torch_world.planning import CEM

model = LeWM(CLSEncoder(img_size=32, patch_size=8, embed_dim=64), action_dim=2, embed_dim=64)
episodes = PointMass(batch_size=64).collect(4)                 # pixels (64, 4, 3, 32, 32), actions (64, 4, 2)
loss = model(episodes["pixels"], episodes["actions"])["loss"]  # prediction + SIGReg

model.eval()
z, goal = model.encode(episodes["pixels"][:, :1]), model.encode(episodes["pixels"][:, -1:])[:, 0]
actions = CEM(horizon=5, action_dim=2)(model, z, goal)         # (64, 5, 2)

This trains LeWM, a JEPA world model, on a point mass and plans toward a goal image. Read the quick start for what each line does, or open the tour notebook.