Package Overview

TorchWM is organized into focused modules so you can use only the pieces you need.

Core APIs

  • world_models.models: High-level models and agents (Dreamer, DreamerAgent, Planet, JEPAAgent)

  • world_models.configs: Configuration containers for Dreamer, JEPA, and diffusion runs

  • world_models.training: Script-style training entrypoints

Environment Integration

  • world_models.envs: Environment adapters for DMC, Gym/Gymnasium, Atari, MuJoCo, Unity ML-Agents

  • world_models.envs.wrappers: Common wrappers for action repeat, action normalization, time limits, and observation shaping

World Model Building Blocks

  • world_models.models.dreamer_rssm: Recurrent state-space model used by Dreamer

  • world_models.vision: Encoders/decoders and action heads for latent dynamics models

  • world_models.reward: Reward and value prediction heads

  • world_models.observations: Symbolic and visual observation reconstruction modules

Representation Learning and Diffusion

  • world_models.models.vit: Vision Transformer and JEPA predictor components

  • world_models.models.diffusion: DDPM scheduler and DiT model implementation

  • world_models.masks: Mask collators for JEPA-style context/target masking

Data and Memory

  • world_models.datasets: CIFAR-10, ImageNet-1K, and generic ImageFolder dataset loaders

  • world_models.memory: Replay buffers for Dreamer and episode-based memory for PlaNet/RSSM

Utilities

  • world_models.utils.dreamer_utils: Logging, parameter freezing, and TD(lambda) return computation

  • world_models.utils.jepa_utils: Optimizer schedules, distributed helpers, and training meters

  • world_models.transforms: Data augmentation pipelines used by JEPA/vision training

Which API Should I Use?

  • End-to-end Dreamer training: DreamerAgent

  • End-to-end JEPA training: JEPAAgent

  • Low-level model experimentation: Dreamer, RSSM, decoder/encoder modules

  • Custom data pipelines: make_cifar10, make_imagenet1k, make_imagefolder