LICENSE
README.md
pyproject.toml
setup.py
tests/test_agents.py
tests/test_checkpoint_hidden_state.py
tests/test_cli.py
tests/test_controller.py
tests/test_decoder.py
tests/test_diamond.py
tests/test_dreamer.py
tests/test_dreamer_utils.py
tests/test_encoder.py
tests/test_env_adapters.py
tests/test_envs.py
tests/test_genie.py
tests/test_iris_agent.py
tests/test_iris_decoder.py
tests/test_jepa_utils.py
tests/test_memory.py
tests/test_modular_rssm.py
tests/test_observations.py
tests/test_operators.py
tests/test_planet.py
tests/test_policy_hidden_state_imagination.py
tests/test_preconditioner_consistency.py
tests/test_public_api.py
tests/test_reproducible_imagination.py
tests/test_rewards.py
tests/test_rssm.py
tests/test_tinyworlds.py
tests/test_trainers.py
tests/test_transforms.py
tests/test_ui_selenium.py
tests/test_utils.py
tests/test_vector_env.py
tests/test_video_datasets.py
tests/test_vit.py
tools/__init__.py
tools/check_docs_render.py
tools/cli.py
torchwm/__init__.py
torchwm.egg-info/PKG-INFO
torchwm.egg-info/SOURCES.txt
torchwm.egg-info/dependency_links.txt
torchwm.egg-info/entry_points.txt
torchwm.egg-info/not-zip-safe
torchwm.egg-info/requires.txt
torchwm.egg-info/top_level.txt
world_models/__init__.py
world_models/api.py
world_models/catalog.py
world_models/benchmarks/__init__.py
world_models/benchmarks/adapters.py
world_models/benchmarks/atari_100k.py
world_models/benchmarks/cli.py
world_models/benchmarks/metrics.py
world_models/benchmarks/reporting.py
world_models/benchmarks/runner.py
world_models/blocks/__init__.py
world_models/blocks/attention.py
world_models/blocks/mhsa.py
world_models/blocks/st_transformer.py
world_models/configs/__init__.py
world_models/configs/diamond_config.py
world_models/configs/dit_config.py
world_models/configs/dreamer_config.py
world_models/configs/genie_config.py
world_models/configs/iris_config.py
world_models/configs/jepa_config.py
world_models/controller/__init__.py
world_models/controller/iris_policy.py
world_models/controller/rollout_generator.py
world_models/controller/rssm_policy.py
world_models/datasets/__init__.py
world_models/datasets/cifar10.py
world_models/datasets/diamond_dataset.py
world_models/datasets/imagenet1k.py
world_models/datasets/nuplan.py
world_models/datasets/tinyworlds.py
world_models/datasets/video_datasets.py
world_models/envs/__init__.py
world_models/envs/ale_atari_env.py
world_models/envs/ale_atari_vector_env.py
world_models/envs/brax_env.py
world_models/envs/diamond_atari.py
world_models/envs/dmc.py
world_models/envs/gym_env.py
world_models/envs/mujoco_env.py
world_models/envs/robotics_env.py
world_models/envs/unity_env.py
world_models/envs/vector_env.py
world_models/envs/wrappers.py
world_models/inference/__init__.py
world_models/inference/operators/__init__.py
world_models/inference/operators/base.py
world_models/inference/operators/dreamer_operator.py
world_models/inference/operators/iris_operator.py
world_models/inference/operators/jepa_operator.py
world_models/inference/operators/planet_operator.py
world_models/inference/operators/utils.py
world_models/masks/__init__.py
world_models/masks/default.py
world_models/masks/multiblock.py
world_models/masks/random.py
world_models/memory/__init__.py
world_models/memory/dreamer_memory.py
world_models/memory/iris_memory.py
world_models/memory/planet_memory.py
world_models/models/__init__.py
world_models/models/dreamer.py
world_models/models/dreamer_rssm.py
world_models/models/dynamics_model.py
world_models/models/genie.py
world_models/models/iris_agent.py
world_models/models/iris_transformer.py
world_models/models/jepa_agent.py
world_models/models/latent_action_model.py
world_models/models/modular_rssm.py
world_models/models/planet.py
world_models/models/rssm.py
world_models/models/vit.py
world_models/models/diffusion/DDPM.py
world_models/models/diffusion/DiT.py
world_models/models/diffusion/__init__.py
world_models/models/diffusion/actor_critic.py
world_models/models/diffusion/diamond_diffusion.py
world_models/models/diffusion/reward_termination.py
world_models/reward/__init__.py
world_models/reward/dreamer_v1_reward.py
world_models/reward/dreamer_v1_value.py
world_models/training/__init__.py
world_models/training/rl_harness.py
world_models/training/run_atari_100k.py
world_models/training/train_diamond.py
world_models/training/train_genie.py
world_models/training/train_iris.py
world_models/training/train_jepa.py
world_models/training/train_planet.py
world_models/training/train_rssm.py
world_models/utils/__init__.py
world_models/utils/data_utils.py
world_models/utils/dreamer_utils.py
world_models/utils/jepa_utils.py
world_models/utils/jit_utils.py
world_models/utils/logging_utils.py
world_models/utils/memory_utils.py
world_models/utils/utils.py
world_models/vision/__init__.py
world_models/vision/dreamer_decoder.py
world_models/vision/dreamer_encoder.py
world_models/vision/iris_decoder.py
world_models/vision/iris_encoder.py
world_models/vision/planet_decoder.py
world_models/vision/planet_encoder.py
world_models/vision/video_tokenizer.py
world_models/vision/vq_layer.py