Environments Guide#

TorchWM supports several environment backends for training, evaluation, and dataset collection. The detailed backend documentation now lives on separate pages so each environment family can document its own installation requirements, factory functions, observation/action contracts, and troubleshooting notes.

Quick start#

Use DreamerConfig.env_backend for Dreamer-compatible DMC, Gym/Gymnasium, MuJoCo, Gymnasium Robotics, Brax, and Unity environments. Choose the backend that matches your installed optional dependencies and task source.

from torchwm import DreamerConfig

cfg = DreamerConfig()
# env_backend may be one of: "dmc", "gym", "mujoco", "robotics", "brax", or "unity_mlagents"
cfg.env_backend = "dmc"
cfg.env = "walker-walk"
cfg.image_size = 64
cfg.action_repeat = 2
cfg.time_limit = 1000

For direct environment construction, use the top-level torchwm factories:

from torchwm import (
    DeepMindControlEnv,
    make_atari_env,
    make_brax_env,
    make_gym_env,
    make_robotics_env,
)

dmc_env = DeepMindControlEnv("cheetah-run", seed=0, size=(64, 64))
gym_env = make_gym_env("Pendulum-v1", seed=0, size=(64, 64))
brax_env = make_brax_env("ant", seed=0, image_size=(64, 64))
atari_env = make_atari_env("ALE/Pong-v5", obs_type="rgb", frameskip=4)
robotics_env = make_robotics_env("HalfCheetah-v2", seed=0, size=(64, 64))

Backend summary#

Backend

Page

Use when

DeepMind Control Suite

DMC

You want Dreamer-style continuous-control tasks with rendered images and native DMC state observations.

Gym/Gymnasium

Gym

You want classic control, Box2D, custom Gym environments, or generic rendered tasks converted to TorchWM image observations.

Brax

Brax

You want JAX/Brax continuous-control tasks wrapped in a Gym-like image adapter for TorchWM training loops.

Atari

Atari

You want Atari environments through Gymnasium/ALE, native ALE vectorization, or Atari-specific DIAMOND-style preprocessing.

MuJoCo

MuJoCo

You want Gymnasium MuJoCo task ids or native MJCF/MJB models.

Gymnasium Robotics

Gymnasium Robotics

You need any id registered by Gymnasium Robotics, including legacy MuJoCo v2/v3 ids.

Unity ML-Agents

Unity

You want to train against external Unity executables with continuous-control behaviors.

Vectorized environments

Vectorized

You need batched rollout collection across native ALE vector envs or multiprocessing Gym-like env factories.

Wrappers

Wrappers

You need action repeat, time limits, action normalization, one-hot actions, reward observations, or image transforms.

Shared TorchWM environment conventions#

  • Image-based training code generally expects an observation dictionary with an image entry.

  • DMC, Gym/Gymnasium, MuJoCo, Gymnasium Robotics, and Unity adapters return channel-first images shaped (3, H, W) with dtype uint8.

  • Atari can expose raw ALE observations, native vectorized observations, or DIAMOND-style preprocessed frames; check the Atari page before feeding observations directly to a model.

  • Dreamer environment creation applies ActionRepeat, NormalizeActions, and TimeLimit after constructing the selected backend. The lightweight catalog exposes Gymnasium Robotics ids to online world-model families except I-JEPA/JEPA, which uses image datasets rather than online Gymnasium environments.

  • Use torchwm envs list to inspect the lightweight backend catalog available to the CLI.

Common issues#

  • Missing optional dependency: install the backend-specific package listed on the backend page.

  • Observation shape mismatch: verify whether your selected factory returns CHW image dictionaries, HWC images, RAM, or vector observations.

  • Action mismatch: distinguish raw discrete action indices from one-hot vectors and normalized continuous actions.

  • Brax backend mismatch: set cfg.brax_backend only to backends supported by your installed Brax release (for example "generalized" or "mjx").

  • Slow environments: disable rendering where possible, use action repeat, and consider vectorized rollout collection.