Index

A

act_and_collect_data() (world_models.models.dreamer.Dreamer method) act_with_world_model() (world_models.models.dreamer.Dreamer method) action_size (world_models.utils.utils.TorchImageEnvWrapper property) action_space (world_models.envs.DeepMindControlEnv property) (world_models.envs.GymImageEnv property) (world_models.envs.NormalizeActions property) (world_models.envs.ObsDict property) (world_models.envs.OneHotAction property) (world_models.envs.UnityMLAgentsEnv property) ActionDecoder (class in world_models.vision.dreamer_decoder) ActionRepeat (class in world_models.envs) actor_loss() (world_models.models.dreamer.Dreamer method) add() (world_models.memory.dreamer_memory.ReplayBuffer method)

add_exploration() (world_models.vision.dreamer_decoder.ActionDecoder method) AllGather (class in world_models.utils.jepa_utils) AllReduce (class in world_models.utils.jepa_utils) AllReduceSum (class in world_models.utils.jepa_utils) append() (world_models.memory.planet_memory.Episode method) (world_models.memory.planet_memory.Memory method) apply_masks() (in module world_models.utils.utils) apply_model() (in module world_models.utils.utils) assign_type() (world_models.utils.utils.TensorBoardMetrics method) atanh() (world_models.vision.dreamer_decoder.TanhBijector method) Attention (class in world_models.models.vit) AttrDict (class in world_models.utils.utils) AverageMeter (class in world_models.utils.jepa_utils)

B

backward() (world_models.utils.jepa_utils.AllGather static method) (world_models.utils.jepa_utils.AllReduce static method) (world_models.utils.jepa_utils.AllReduceSum static method) BATCH (world_models.configs.DiTConfig attribute)

BETA_END (world_models.configs.DiTConfig attribute) BETA_START (world_models.configs.DiTConfig attribute) Block (class in world_models.models.vit) bottle() (in module world_models.utils.utils)

C

CHANNELS (world_models.configs.DiTConfig attribute) classes (world_models.datasets.imagenet1k.ImageNetSubset property) close() (world_models.envs.GymImageEnv method) (world_models.envs.UnityMLAgentsEnv method) (world_models.utils.utils.TorchImageEnvWrapper method) collect_random_episodes() (world_models.models.dreamer.Dreamer method) combine_videos() (in module world_models.utils.utils)

compute_return() (in module world_models.utils.dreamer_utils) ConvDecoder (class in world_models.vision.dreamer_decoder) ConvEmbed (class in world_models.models.vit) ConvEncoder (class in world_models.vision.dreamer_encoder) copy_imgnt_locally() (in module world_models.datasets.imagenet1k) CosineWDSchedule (class in world_models.utils.jepa_utils) CSVLogger (class in world_models.utils.jepa_utils)

D

DATASET (world_models.configs.DiTConfig attribute) DDPM (class in world_models.models.diffusion.DDPM) DeepMindControlEnv (class in world_models.envs) DenseDecoder (class in world_models.vision.dreamer_decoder) DEPTH (world_models.configs.DiTConfig attribute) detach_state() (world_models.models.dreamer_rssm.RSSM method) DiT (class in world_models.models.diffusion.DiT)

DiTConfig (class in world_models.configs) Dreamer (class in world_models.models.dreamer) DreamerAgent (class in world_models.models.dreamer) DreamerConfig (class in world_models.configs) DROP (world_models.configs.DiTConfig attribute) drop_path() (in module world_models.models.vit) DropPath (class in world_models.models.vit) dump_scalars_to_pickle() (world_models.utils.dreamer_utils.Logger method)

E

EMA (world_models.configs.DiTConfig attribute) EMA_DECAY (world_models.configs.DiTConfig attribute) ensure_results_dir_exists() (in module world_models.utils.utils) entropy() (world_models.vision.dreamer_decoder.SampleDist method)

Episode (class in world_models.memory.planet_memory) EPOCHS (world_models.configs.DiTConfig attribute) evaluate() (world_models.models.dreamer.Dreamer method) (world_models.models.dreamer.DreamerAgent method)

F

filter_dataset_() (world_models.datasets.imagenet1k.ImageNetSubset method) fix_init_weight() (world_models.models.vit.VisionTransformer method) (world_models.models.vit.VisionTransformerPredictor method) flatten_dict() (in module world_models.utils.utils) flush() (world_models.utils.dreamer_utils.Logger method) forward() (world_models.models.diffusion.DiT.DiT method) (world_models.models.diffusion.DiT.PatchEmbed method) (world_models.models.diffusion.DiT.PatchUnEmbed method) (world_models.models.diffusion.DiT.TransformerBlock method) (world_models.models.vit.Attention method) (world_models.models.vit.Block method) (world_models.models.vit.ConvEmbed method) (world_models.models.vit.DropPath method) (world_models.models.vit.MLP method) (world_models.models.vit.PatchEmbed method) (world_models.models.vit.VisionTransformer method) (world_models.models.vit.VisionTransformerPredictor method) (world_models.utils.jepa_utils.AllGather static method) (world_models.utils.jepa_utils.AllReduce static method) (world_models.utils.jepa_utils.AllReduceSum static method) (world_models.vision.dreamer_decoder.ActionDecoder method) (world_models.vision.dreamer_decoder.ConvDecoder method) (world_models.vision.dreamer_decoder.DenseDecoder method) (world_models.vision.dreamer_encoder.ConvEncoder method)

FreezeParameters (class in world_models.utils.dreamer_utils)

G

GaussianBlur (class in world_models.transforms.transforms) get_1d_sincos_pos_embed() (in module world_models.models.vit) get_1d_sincos_pos_embed_from_grid() (in module world_models.models.vit) get_2d_sincos_pos_embed() (in module world_models.models.vit) get_2d_sincos_pos_embed_from_grid() (in module world_models.models.vit) get_combined_params() (in module world_models.utils.utils)

get_dist() (world_models.models.dreamer_rssm.RSSM method) get_dit_config() (in module world_models.configs) get_mask() (in module world_models.utils.utils) get_parameters() (in module world_models.utils.dreamer_utils) gpu_timer() (in module world_models.utils.jepa_utils) grad_logger() (in module world_models.utils.jepa_utils) GymImageEnv (class in world_models.envs)

H

HEADS (world_models.configs.DiTConfig attribute)

I

ImageNet (class in world_models.datasets.imagenet1k) ImageNetSubset (class in world_models.datasets.imagenet1k) imagine_rollout() (world_models.models.dreamer_rssm.RSSM method) imagine_step() (world_models.models.dreamer_rssm.RSSM method)

IMG_SIZE (world_models.configs.DiTConfig attribute) init_distributed() (in module world_models.utils.jepa_utils) init_state() (world_models.models.dreamer_rssm.RSSM method) interpolate_pos_encoding() (world_models.models.vit.VisionTransformer method)

J

JEPAAgent (class in world_models.models.jepa_agent)

JEPAConfig (class in world_models.configs)

L

lineplot() (in module world_models.utils.utils) list_available_atari_envs() (in module world_models.envs) load_memory() (in module world_models.utils.utils) load_yml_config() (in module world_models.utils.utils) log() (world_models.utils.jepa_utils.CSVLogger method)

log_abs_det_jacobian() (world_models.vision.dreamer_decoder.TanhBijector method) log_scalar() (world_models.utils.dreamer_utils.Logger method) log_scalars() (world_models.utils.dreamer_utils.Logger method) log_videos() (world_models.utils.dreamer_utils.Logger method) Logger (class in world_models.utils.dreamer_utils) LR (world_models.configs.DiTConfig attribute)

M

main() (in module world_models.training.train_jepa) make_atari_env() (in module world_models.envs) make_atari_vector_env() (in module world_models.envs) make_cifar10() (in module world_models.datasets.cifar10) make_env() (in module world_models.models.dreamer) make_gym_env() (in module world_models.envs) make_half_cheetah_env() (in module world_models.envs) make_humanoid_env() (in module world_models.envs) make_imagefolder() (in module world_models.datasets.imagenet1k) make_imagenet1k() (in module world_models.datasets.imagenet1k) make_transforms() (in module world_models.transforms.transforms) make_unity_mlagents_env() (in module world_models.envs) MaskCollator (class in world_models.masks.multiblock) (class in world_models.masks.random) max_episode_steps (world_models.envs.GymImageEnv property) (world_models.envs.UnityMLAgentsEnv property) (world_models.utils.utils.TorchImageEnvWrapper property) mean() (world_models.vision.dreamer_decoder.SampleDist method) Memory (class in world_models.memory.planet_memory) MLP (class in world_models.models.vit) mode() (world_models.vision.dreamer_decoder.SampleDist method) module world_models.configs world_models.controller.rssm_policy world_models.datasets.cifar10 world_models.datasets.imagenet1k world_models.envs world_models.masks.multiblock world_models.masks.random world_models.memory.dreamer_memory world_models.memory.planet_memory world_models.models world_models.models.diffusion.DDPM world_models.models.diffusion.DiT world_models.models.dreamer world_models.models.dreamer_rssm world_models.models.jepa_agent world_models.models.vit world_models.training.train_jepa world_models.transforms.transforms world_models.utils.dreamer_utils world_models.utils.jepa_utils world_models.utils.utils world_models.vision.dreamer_decoder world_models.vision.dreamer_encoder

N

name (world_models.vision.dreamer_decoder.SampleDist property)

normalize_frames_for_saving() (in module world_models.utils.utils) NormalizeActions (class in world_models.envs)

O

obs_space (world_models.envs.RenderImage property) (world_models.envs.ResizeImage property) ObsDict (class in world_models.envs) observation_size (world_models.utils.utils.TorchImageEnvWrapper property) observation_space (world_models.envs.DeepMindControlEnv property) (world_models.envs.GymImageEnv property) (world_models.envs.ObsDict property) (world_models.envs.RewardObs property) (world_models.envs.UnityMLAgentsEnv property)

observe_rollout() (world_models.models.dreamer_rssm.RSSM method) observe_step() (world_models.models.dreamer_rssm.RSSM method) OneHotAction (class in world_models.envs)

P

p_sample() (world_models.models.diffusion.DDPM.DDPM method) PATCH (world_models.configs.DiTConfig attribute) PatchEmbed (class in world_models.models.diffusion.DiT) (class in world_models.models.vit) PatchUnEmbed (class in world_models.models.diffusion.DiT)

plot_metrics() (in module world_models.utils.utils) poll() (world_models.controller.rssm_policy.RSSMPolicy method) postprocess_img() (in module world_models.utils.utils) preprocess_img() (in module world_models.utils.utils) preprocess_obs() (in module world_models.models.dreamer)

Q

q_sample() (world_models.models.diffusion.DDPM.DDPM method)

R

render() (world_models.envs.DeepMindControlEnv method) (world_models.envs.GymImageEnv method) (world_models.envs.UnityMLAgentsEnv method) (world_models.utils.utils.TorchImageEnvWrapper method) RenderImage (class in world_models.envs) repeat_interleave_batch() (in module world_models.utils.jepa_utils) ReplayBuffer (class in world_models.memory.dreamer_memory) reset() (world_models.controller.rssm_policy.RSSMPolicy method) (world_models.envs.DeepMindControlEnv method) (world_models.envs.GymImageEnv method) (world_models.envs.ObsDict method) (world_models.envs.OneHotAction method) (world_models.envs.RenderImage method) (world_models.envs.ResizeImage method) (world_models.envs.RewardObs method) (world_models.envs.TimeLimit method) (world_models.envs.UnityMLAgentsEnv method) (world_models.utils.jepa_utils.AverageMeter method) (world_models.utils.utils.TorchImageEnvWrapper method)

ResizeImage (class in world_models.envs) restore_checkpoint() (world_models.models.dreamer.Dreamer method) RewardObs (class in world_models.envs) ROOT_PATH (world_models.configs.DiTConfig attribute) RSSM (class in world_models.models.dreamer_rssm) RSSMPolicy (class in world_models.controller.rssm_policy)

S

sample() (world_models.memory.dreamer_memory.ReplayBuffer method) (world_models.memory.planet_memory.Memory method) (world_models.models.diffusion.DDPM.DDPM method) (world_models.vision.dreamer_decoder.SampleDist method) sample_random_action() (world_models.utils.utils.TorchImageEnvWrapper method) SampleDist (class in world_models.vision.dreamer_decoder) save() (world_models.models.dreamer.Dreamer method) save_frames() (in module world_models.utils.utils) save_video() (in module world_models.utils.utils) SelectAction (class in world_models.envs) seq_to_batch() (world_models.models.dreamer_rssm.RSSM method) sign (world_models.vision.dreamer_decoder.TanhBijector property) sinusoidal_time_embedding() (in module world_models.models.diffusion.DiT) size (world_models.memory.planet_memory.Episode property) (world_models.memory.planet_memory.Memory property) stack_states() (world_models.models.dreamer_rssm.RSSM method)

step() (world_models.envs.ActionRepeat method) (world_models.envs.DeepMindControlEnv method) (world_models.envs.GymImageEnv method) (world_models.envs.NormalizeActions method) (world_models.envs.ObsDict method) (world_models.envs.OneHotAction method) (world_models.envs.RenderImage method) (world_models.envs.ResizeImage method) (world_models.envs.RewardObs method) (world_models.envs.SelectAction method) (world_models.envs.TimeLimit method) (world_models.envs.UnityMLAgentsEnv method) (world_models.masks.multiblock.MaskCollator method) (world_models.masks.random.MaskCollator method) (world_models.utils.jepa_utils.CosineWDSchedule method) (world_models.utils.jepa_utils.WarmupCosineSchedule method) (world_models.utils.utils.TorchImageEnvWrapper method)

T

TanhBijector (class in world_models.vision.dreamer_decoder) TensorBoardMetrics (class in world_models.utils.utils) terminate() (world_models.memory.planet_memory.Episode method) TimeLimit (class in world_models.envs) TIMESTEPS (world_models.configs.DiTConfig attribute) to_dict() (world_models.configs.JEPAConfig method) to_tensor_obs() (in module world_models.utils.utils)

TorchImageEnvWrapper (class in world_models.utils.utils) train() (world_models.models.diffusion.DiT.DiT class method) (world_models.models.dreamer.DreamerAgent method) (world_models.models.jepa_agent.JEPAAgent method) train_one_batch() (world_models.models.dreamer.Dreamer method) TransformerBlock (class in world_models.models.diffusion.DiT) trunc_normal_() (in module world_models.utils.jepa_utils)

U

UnityMLAgentsEnv (class in world_models.envs)

update() (world_models.utils.jepa_utils.AverageMeter method) (world_models.utils.utils.TensorBoardMetrics method)

V

value_loss() (world_models.models.dreamer.Dreamer method) VisionTransformer (class in world_models.models.vit) VisionTransformerPredictor (class in world_models.models.vit) vit_base() (in module world_models.models.vit) vit_giant() (in module world_models.models.vit)

vit_huge() (in module world_models.models.vit) vit_large() (in module world_models.models.vit) vit_predictor() (in module world_models.models.vit) vit_small() (in module world_models.models.vit) vit_tiny() (in module world_models.models.vit)

W

WarmupCosineSchedule (class in world_models.utils.jepa_utils) WIDTH (world_models.configs.DiTConfig attribute) WORKDIR (world_models.configs.DiTConfig attribute) world_model_loss() (world_models.models.dreamer.Dreamer method) world_models.configs module world_models.controller.rssm_policy module world_models.datasets.cifar10 module world_models.datasets.imagenet1k module world_models.envs module world_models.masks.multiblock module world_models.masks.random module world_models.memory.dreamer_memory module world_models.memory.planet_memory module world_models.models module world_models.models.diffusion.DDPM module

world_models.models.diffusion.DiT module world_models.models.dreamer module world_models.models.dreamer_rssm module world_models.models.jepa_agent module world_models.models.vit module world_models.training.train_jepa module world_models.transforms.transforms module world_models.utils.dreamer_utils module world_models.utils.jepa_utils module world_models.utils.utils module world_models.vision.dreamer_decoder module world_models.vision.dreamer_encoder module