Index

A

act() (world_models.controller.iris_policy.IRISPolicy method) (world_models.models.iris_agent.IRISAgent method) 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_channels (world_models.configs.DiamondConfig attribute) actor_loss() (world_models.models.dreamer.Dreamer method) actor_lstm_dim (world_models.configs.DiamondConfig attribute) actor_res_blocks (world_models.configs.DiamondConfig attribute)

adam_epsilon (world_models.configs.DiamondConfig attribute) add() (world_models.memory.dreamer_memory.ReplayBuffer method) (world_models.memory.iris_memory.IRISOnPolicyBuffer method) (world_models.memory.iris_memory.IRISReplayBuffer 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) batch_size (world_models.configs.DiamondConfig attribute) benchmarks.atari_100k module

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) buffer_capacity (world_models.memory.iris_memory.IRISReplayBuffer property) burn_in_length (world_models.configs.DiamondConfig attribute)

C

CHANNELS (world_models.configs.DiTConfig attribute) classes (world_models.datasets.imagenet1k.ImageNetSubset property) clear() (world_models.memory.iris_memory.IRISOnPolicyBuffer method) close() (world_models.envs.GymImageEnv method) (world_models.envs.UnityMLAgentsEnv method) (world_models.utils.utils.StreamingVideoWriter method) (world_models.utils.utils.TorchImageEnvWrapper method) CNNFeatureExtractor (class in world_models.controller.iris_policy) collect_experience() (world_models.training.train_iris.IRISTrainer method) collect_random_episodes() (world_models.models.dreamer.Dreamer method)

combine_videos() (in module world_models.utils.utils) compute_human_normalized_score() (in module benchmarks.atari_100k) compute_lambda_return() (in module world_models.models.iris_agent) 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) decode() (world_models.vision.iris_decoder.DiscreteAutoencoder method) decode_from_embeddings() (world_models.vision.iris_decoder.IRISDecoder method) decode_from_indices() (world_models.vision.iris_encoder.IRISEncoder method) decode_indices() (world_models.vision.vq_layer.VectorQuantizer method) (world_models.vision.vq_layer.VectorQuantizerEMA method) 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) device (world_models.configs.DiamondConfig attribute) DiamondConfig (class in world_models.configs)

diffusion_channels (world_models.configs.DiamondConfig attribute) diffusion_cond_dim (world_models.configs.DiamondConfig attribute) diffusion_res_blocks (world_models.configs.DiamondConfig attribute) discount_factor (world_models.configs.DiamondConfig attribute) DiscreteAutoencoder (class in world_models.vision.iris_decoder) 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) encode() (world_models.vision.iris_decoder.DiscreteAutoencoder method) encode_to_indices() (world_models.vision.iris_encoder.IRISEncoder method) ensure_results_dir_exists() (in module world_models.utils.utils) entropy() (world_models.vision.dreamer_decoder.SampleDist method) entropy_weight (world_models.configs.DiamondConfig attribute)

environment_steps_per_epoch (world_models.configs.DiamondConfig attribute) Episode (class in world_models.memory.planet_memory) EPOCHS (world_models.configs.DiTConfig attribute) epsilon_greedy (world_models.configs.DiamondConfig attribute) eval_interval (world_models.configs.DiamondConfig attribute) evaluate() (world_models.models.dreamer.Dreamer method) (world_models.models.dreamer.DreamerAgent method) (world_models.training.train_iris.IRISTrainer 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.controller.iris_policy.CNNFeatureExtractor method) (world_models.controller.iris_policy.IRISActor method) (world_models.controller.iris_policy.IRISCritic method) (world_models.controller.iris_policy.IRISPolicy method) (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.iris_transformer.IRISTransformer method) (world_models.models.iris_transformer.IRISWorldModel 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) (world_models.vision.iris_decoder.DiscreteAutoencoder method) (world_models.vision.iris_decoder.IRISDecoder method) (world_models.vision.iris_decoder.ResidualBlock method) (world_models.vision.iris_decoder.UpsampleBlock method) (world_models.vision.iris_encoder.IRISEncoder method) (world_models.vision.iris_encoder.ResidualBlock method) (world_models.vision.iris_encoder.SelfAttentionBlock method) (world_models.vision.vq_layer.VectorQuantizer method) (world_models.vision.vq_layer.VectorQuantizerEMA method)

forward_actor_critic() (world_models.models.iris_agent.IRISAgent method) frameskip (world_models.configs.DiamondConfig attribute) FreezeParameters (class in world_models.utils.dreamer_utils)

G

game (world_models.configs.DiamondConfig attribute) 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_action() (world_models.controller.iris_policy.IRISActor method) get_arrays() (world_models.memory.iris_memory.IRISOnPolicyBuffer method) get_autoencoder_config() (world_models.configs.iris_config.IRISConfig method) 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_epsilon() (world_models.training.train_iris.IRISTrainer method) get_frame_shape() (world_models.configs.iris_config.IRISConfig method) get_mask() (in module world_models.utils.utils) get_parameters() (in module world_models.utils.dreamer_utils) get_rl_config() (world_models.configs.iris_config.IRISConfig method) get_transformer_config() (world_models.configs.iris_config.IRISConfig method) 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) imagination_horizon (world_models.configs.DiamondConfig attribute) imagine() (world_models.models.iris_transformer.IRISWorldModel method) imagine_rollout() (world_models.models.dreamer_rssm.RSSM method) (world_models.models.iris_agent.IRISAgent 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_hidden() (world_models.controller.iris_policy.IRISPolicy method) init_hidden_state() (world_models.controller.iris_policy.IRISActor method) (world_models.controller.iris_policy.IRISCritic method) init_state() (world_models.models.dreamer_rssm.RSSM method)

interpolate_pos_encoding() (world_models.models.vit.VisionTransformer method) IRISActor (class in world_models.controller.iris_policy) IRISAgent (class in world_models.models.iris_agent) IRISConfig (class in world_models.configs.iris_config) IRISCritic (class in world_models.controller.iris_policy) IRISDecoder (class in world_models.vision.iris_decoder) IRISEncoder (class in world_models.vision.iris_encoder) IRISOnPolicyBuffer (class in world_models.memory.iris_memory) IRISPolicy (class in world_models.controller.iris_policy) IRISReplayBuffer (class in world_models.memory.iris_memory) IRISTrainer (class in world_models.training.train_iris) IRISTransformer (class in world_models.models.iris_transformer) IRISWorldModel (class in world_models.models.iris_transformer)

J

JEPAConfig (class in world_models.configs)

L

lambda_returns (world_models.configs.DiamondConfig attribute) learning_rate (world_models.configs.DiamondConfig attribute) lineplot() (in module world_models.utils.utils) list_available_atari_envs() (in module world_models.envs) load() (world_models.models.iris_agent.IRISAgent method) 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_interval (world_models.configs.DiamondConfig attribute) 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 benchmarks.atari_100k) (in module world_models.training.train_iris) 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) max_noop (world_models.configs.DiamondConfig attribute) 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 benchmarks.atari_100k world_models.configs world_models.configs.iris_config world_models.controller.iris_policy 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.iris_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.iris_agent world_models.models.iris_transformer world_models.models.vit world_models.training.train_iris 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 world_models.vision.iris_decoder world_models.vision.iris_encoder world_models.vision.vq_layer

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)

num_conditioning_frames (world_models.configs.DiamondConfig attribute) num_epochs (world_models.configs.DiamondConfig attribute) num_sampling_steps (world_models.configs.DiamondConfig attribute) num_seeds (world_models.configs.DiamondConfig attribute)

O

obs_size (world_models.configs.DiamondConfig attribute) 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_mean (world_models.configs.DiamondConfig attribute) p_sample() (world_models.models.diffusion.DDPM.DDPM method) p_std (world_models.configs.DiamondConfig attribute) 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) predict_next_tokens() (world_models.models.iris_transformer.IRISTransformer method) preprocess_frame() (world_models.training.train_iris.IRISTrainer method) preprocess_img() (in module world_models.utils.utils) preprocess_obs() (in module world_models.models.dreamer) preset (world_models.configs.DiamondConfig attribute) print_results_table() (in module benchmarks.atari_100k)

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)

ResidualBlock (class in world_models.vision.iris_decoder) (class in world_models.vision.iris_encoder) ResizeImage (class in world_models.envs) restore_checkpoint() (world_models.models.dreamer.Dreamer method) reward_channels (world_models.configs.DiamondConfig attribute) reward_clip (world_models.configs.DiamondConfig attribute) reward_cond_dim (world_models.configs.DiamondConfig attribute) reward_lstm_dim (world_models.configs.DiamondConfig attribute) reward_res_blocks (world_models.configs.DiamondConfig attribute) RewardObs (class in world_models.envs) rho (world_models.configs.DiamondConfig attribute) ROOT_PATH (world_models.configs.DiTConfig attribute) RSSM (class in world_models.models.dreamer_rssm) RSSMPolicy (class in world_models.controller.rssm_policy) run_atari_100k() (in module benchmarks.atari_100k) run_single_game() (in module benchmarks.atari_100k)

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_next_tokens() (world_models.models.iris_transformer.IRISTransformer method) sample_random_action() (world_models.utils.utils.TorchImageEnvWrapper method) sample_sequence() (world_models.memory.iris_memory.IRISReplayBuffer method) sample_single() (world_models.memory.iris_memory.IRISReplayBuffer method) SampleDist (class in world_models.vision.dreamer_decoder) sampling_method (world_models.configs.DiamondConfig attribute) save() (world_models.models.dreamer.Dreamer method) (world_models.models.iris_agent.IRISAgent method) save_frames() (in module world_models.utils.utils) save_interval (world_models.configs.DiamondConfig attribute) save_video() (in module world_models.utils.utils) seed (world_models.configs.DiamondConfig attribute) SelectAction (class in world_models.envs) SelfAttentionBlock (class in world_models.vision.iris_encoder) seq_to_batch() (world_models.models.dreamer_rssm.RSSM method) sigma_data (world_models.configs.DiamondConfig attribute) sigma_max (world_models.configs.DiamondConfig attribute) sigma_min (world_models.configs.DiamondConfig attribute)

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) StreamingVideoWriter (class in world_models.utils.utils)

T

TanhBijector (class in world_models.vision.dreamer_decoder) TensorBoardMetrics (class in world_models.utils.utils) terminate() (world_models.memory.planet_memory.Episode method) terminate_on_life_loss (world_models.configs.DiamondConfig attribute) 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.training.train_iris.IRISTrainer method) train_epoch() (world_models.training.train_iris.IRISTrainer method) train_one_batch() (world_models.models.dreamer.Dreamer method) training_steps_per_epoch (world_models.configs.DiamondConfig attribute) 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)

update_actor_critic() (world_models.models.iris_agent.IRISAgent method) update_autoencoder() (world_models.models.iris_agent.IRISAgent method) update_transformer() (world_models.models.iris_agent.IRISAgent method) UpsampleBlock (class in world_models.vision.iris_decoder)

V

value_loss() (world_models.models.dreamer.Dreamer method) VectorQuantizer (class in world_models.vision.vq_layer) VectorQuantizerEMA (class in world_models.vision.vq_layer) VisionTransformer (class in world_models.models.vit) VisionTransformerPredictor (class in world_models.models.vit) visualize_latent_tsne() (in module world_models.utils.utils) visualize_latent_umap() (in module world_models.utils.utils)

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) weight_decay_actor (world_models.configs.DiamondConfig attribute) weight_decay_diffusion (world_models.configs.DiamondConfig attribute) weight_decay_reward (world_models.configs.DiamondConfig attribute) 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.configs.iris_config module world_models.controller.iris_policy 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.iris_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.iris_agent module world_models.models.iris_transformer module world_models.models.vit module world_models.training.train_iris 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 world_models.vision.iris_decoder module world_models.vision.iris_encoder module world_models.vision.vq_layer module write_frame() (world_models.utils.utils.StreamingVideoWriter method)