torch-world 0.0.1

PyTorch World (torch_world) is a library of world models for PyTorch. It gives world-model research shared parts (layers, masking, objectives, planners, environments) so that a new model is a short file and existing models are easy to compare.

uv add torch_world                 # core: depends only on torch
uv add "torch_world[examples]"     # + matplotlib, for the notebooks' figures
uv add "torch_world[action]"       # + mujoco, for the RobotCube arms and Push-T

Quick start

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)

Structure

Models

JEPA

Joint-embedding predictive architectures (JEPAs) predict the representation of a target from the representation of a context, instead of reconstructing pixels. Action-conditioned JEPAs are world models: they predict the next latent state from the current one and an action, and planning searches for the actions that bring the imagined latent closest to a goal.

Adding a model

Subclass JEPA (representation learning) or WorldModel (action-conditioned dynamics) and build it from the shared modules. A WorldModel that implements encode, predict and forward gets rollout, cost and CEM planning with no extra code. See examples/custom_model.ipynb and CONTRIBUTING.md.

Tests

pytest

Contributors

Contributions are welcome: see CONTRIBUTING.md.

Supported by

License

Apache 2.0