Index

A

act() (torchwm.IRISAgent method) (torchwm.IRISPolicy method) (world_models.controller.iris_policy.IRISPolicy method) (world_models.IRISAgent method) (world_models.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_embedding_dim (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute) action_encoder_depth (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute) action_encoder_dim (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute) action_noise (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) action_num_heads (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) action_pooling (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (torchwm.LatentActionModelConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.genie_config.LatentActionModelConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.configs.LatentActionModelConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.LatentActionModelConfig attribute) (world_models.training.train_genie.GenieConfig attribute) action_repeat (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) action_size (torchwm.IRISActor attribute) (torchwm.IRISReplayBuffer attribute) (world_models.configs.wm_config.WMControllerConfig attribute) (world_models.configs.wm_config.WMMDNRNNConfig attribute) (world_models.controller.iris_policy.IRISActor attribute) (world_models.IRISActor attribute) (world_models.IRISReplayBuffer attribute) (world_models.memory.iris_memory.IRISReplayBuffer attribute) (world_models.models.controller.Controller attribute) (world_models.utils.utils.TorchImageEnvWrapper property) action_space (world_models.envs.BraxImageEnv property) (world_models.envs.BSuiteImageEnv property) (world_models.envs.DeepMindControlEnv attribute) (world_models.envs.DeepMindControlEnv property) (world_models.envs.dmc.DeepMindControlEnv attribute) (world_models.envs.dmc.DeepMindControlEnv property) (world_models.envs.dmlab.DMLabEnv property) (world_models.envs.DMLabEnv property) (world_models.envs.gym_env.GymImageEnv attribute) (world_models.envs.gym_env.GymImageEnv property) (world_models.envs.GymImageEnv attribute) (world_models.envs.GymImageEnv property) (world_models.envs.mujoco_env.MuJoCoImageEnv property) (world_models.envs.MuJoCoImageEnv property) (world_models.envs.NormalizeActions property) (world_models.envs.ObsDict property) (world_models.envs.OneHotAction property) (world_models.envs.procgen_env.ProcgenImageEnv property) (world_models.envs.ProcgenImageEnv property) (world_models.envs.unity_env.UnityMLAgentsEnv attribute) (world_models.envs.unity_env.UnityMLAgentsEnv property) (world_models.envs.UnityMLAgentsEnv attribute) (world_models.envs.UnityMLAgentsEnv property) (world_models.envs.wrappers.NormalizeActions property) (world_models.envs.wrappers.ObsDict property) (world_models.envs.wrappers.OneHotAction property) action_vocab_size (torchwm.DynamicsModelConfig attribute) (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.DynamicsModelConfig attribute) (world_models.configs.genie_config.DynamicsModelConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.DynamicsModelConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute)

ActionDecoder (class in torchwm) (class in world_models) (class in world_models.vision.dreamer_decoder) ActionRepeat (class in world_models.envs) (class in world_models.envs.wrappers) actions (torchwm.IRISOnPolicyBuffer attribute) (world_models.IRISOnPolicyBuffer attribute) (world_models.memory.iris_memory.IRISOnPolicyBuffer attribute) actor (torchwm.IRISPolicy attribute) (world_models.controller.iris_policy.IRISPolicy attribute) (world_models.IRISPolicy attribute) actor_channels (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.diamond_config.ModelPreset attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) actor_learning_rate (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) actor_loss() (world_models.models.dreamer.Dreamer method) actor_lstm_dim (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.diamond_config.ModelPreset attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) actor_res_blocks (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) ActorCritic (class in world_models.training.rl_harness) ActorCriticNetwork (class in torchwm) (class in world_models) (class in world_models.models.diffusion.actor_critic) AdaLNNormalization (class in torchwm) (class in world_models) adam_epsilon (torchwm.DiamondConfig attribute) (torchwm.DreamerConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DiamondConfig attribute) (world_models.DreamerConfig attribute) AdaptiveGroupNorm (class in world_models.models.diffusion.diamond_diffusion) add() (torchwm.IRISOnPolicyBuffer method) (torchwm.IRISReplayBuffer method) (torchwm.ReplayBuffer method) (world_models.datasets.diamond_dataset.ReplayBuffer method) (world_models.IRISOnPolicyBuffer method) (world_models.IRISReplayBuffer method) (world_models.memory.dreamer_memory.ReplayBuffer method) (world_models.memory.iris_memory.IRISOnPolicyBuffer method) (world_models.memory.iris_memory.IRISReplayBuffer method) (world_models.ReplayBuffer method) add_exploration() (torchwm.ActionDecoder method) (world_models.ActionDecoder method) (world_models.vision.dreamer_decoder.ActionDecoder method) agent_types (world_models.datasets.nuplan.NuPlanSample attribute) agents_future (world_models.datasets.nuplan.NuPlanSample attribute) agents_mask (world_models.datasets.nuplan.NuPlanSample attribute) agents_past (world_models.datasets.nuplan.NuPlanSample attribute) algo (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) aliases (torchwm.EnvBackendSpec attribute) (torchwm.ModelSpec attribute) (world_models.api.EnvBackendSpec attribute) (world_models.api.ModelSpec attribute) (world_models.EnvBackendSpec attribute) (world_models.ModelSpec attribute) 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() (torchwm.Episode method) (torchwm.Memory method) (world_models.Episode method) (world_models.Memory method) (world_models.memory.dreamer_memory.Episode method) (world_models.memory.dreamer_memory.Memory method) (world_models.memory.planet_memory.Episode method) (world_models.memory.planet_memory.Memory method) apply_gradient_checkpointing() (in module world_models.utils.memory_utils) apply_masks() (in module world_models.utils.utils) apply_model() (in module world_models.utils.utils) assert_finite() (in module world_models.utils.logging_utils) assert_finite_values() (in module world_models.utils.logging_utils) assign_type() (world_models.utils.utils.TensorBoardMetrics method) atanh() (torchwm.TanhBijector method) (world_models.TanhBijector method) (world_models.vision.dreamer_decoder.TanhBijector method) Attention (class in world_models.models.vit) AttentionBlock (class in world_models.models.diffusion.diamond_diffusion) attn_drop_rate (torchwm.DynamicsModelConfig attribute) (torchwm.STTransformerConfig attribute) (world_models.configs.DynamicsModelConfig attribute) (world_models.configs.genie_config.DynamicsModelConfig attribute) (world_models.configs.genie_config.STTransformerConfig attribute) (world_models.configs.STTransformerConfig attribute) (world_models.DynamicsModelConfig attribute) (world_models.STTransformerConfig attribute) AttrDict (class in world_models.utils.utils) autoregressive_sample() (torchwm.DynamicsModel method) (world_models.DynamicsModel method) (world_models.models.dynamics_model.DynamicsModel method) AverageMeter (class in world_models.utils.jepa_utils)

B

BackboneBase (class in world_models.models.modular_rssm) 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) BaseAdapter (class in world_models.benchmarks.adapters) BATCH (torchwm.DiTConfig attribute) (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) (world_models.DiTConfig attribute) batch() (torchwm.OperatorABC method) (world_models.inference.operators.base.OperatorABC method) (world_models.OperatorABC method) batch_size (torchwm.DiamondConfig attribute) (torchwm.DreamerConfig attribute) (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (torchwm.IRISReplayBuffer attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.configs.wm_config.WMMDNRNNConfig attribute) (world_models.datasets.tinyworlds.TinyWorldsConfig attribute) (world_models.datasets.video_datasets.DatasetConfig attribute) (world_models.DiamondConfig attribute) (world_models.DreamerConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.IRISReplayBuffer attribute) (world_models.memory.iris_memory.IRISReplayBuffer attribute) (world_models.training.train_genie.GenieConfig attribute) BenchmarkRunner (class in world_models.benchmarks.runner) BETA_END (torchwm.DiTConfig attribute) (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) (world_models.DiTConfig attribute)

BETA_START (torchwm.DiTConfig attribute) (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) (world_models.DiTConfig attribute) Block (class in world_models.models.vit) bootstrap_ci() (in module world_models.benchmarks.metrics) bootstrap_iqm_ci() (in module world_models.benchmarks.metrics) bottle() (in module world_models.utils.utils) brax_auto_reset (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) brax_backend (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) brax_jit (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) brax_suppress_warp_warnings (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) BraxImageEnv (class in world_models.envs) BSuiteImageEnv (class in world_models.envs) buffer_capacity (torchwm.IRISReplayBuffer property) (world_models.IRISReplayBuffer property) (world_models.memory.iris_memory.IRISReplayBuffer property) buffer_size (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.datasets.wm_dataset.RolloutDataset attribute) (world_models.DreamerConfig attribute) burn_in_length (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute)

C

cache_dir (world_models.datasets.tinyworlds.TinyWorldsConfig attribute) causal_mask (world_models.blocks.st_transformer.STTemporalAttention attribute) CHANNELS (torchwm.DiTConfig attribute) (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) (world_models.DiTConfig attribute) checkpoint_interval (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) checkpoint_path (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) classes (world_models.datasets.imagenet1k.ImageNetSubset property) clear() (torchwm.IRISOnPolicyBuffer method) (world_models.IRISOnPolicyBuffer method) (world_models.memory.iris_memory.IRISOnPolicyBuffer method) close() (world_models.envs.BraxImageEnv method) (world_models.envs.BSuiteImageEnv method) (world_models.envs.dmlab.DMLabEnv method) (world_models.envs.DMLabEnv method) (world_models.envs.gym_env.GymImageEnv method) (world_models.envs.GymImageEnv method) (world_models.envs.mujoco_env.MuJoCoImageEnv method) (world_models.envs.MuJoCoImageEnv method) (world_models.envs.procgen_env.ProcgenImageEnv method) (world_models.envs.ProcgenImageEnv method) (world_models.envs.unity_env.UnityMLAgentsEnv method) (world_models.envs.UnityMLAgentsEnv method) (world_models.envs.vector_env.VectorizedEnv method) (world_models.envs.WorldModelEnv method) (world_models.utils.logging_utils.MetricsLogger method) (world_models.utils.utils.StreamingVideoWriter method) (world_models.utils.utils.TorchImageEnvWrapper method) cnn_activation_function (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) CNNDecoder (class in torchwm) (class in world_models) (class in world_models.vision.planet_decoder) CNNEncoder (class in torchwm) (class in world_models) (class in world_models.vision.planet_encoder) CNNFeatureExtractor (class in torchwm) (class in world_models) (class in world_models.controller.iris_policy) collate_fn() (in module world_models.datasets.diamond_dataset) collect_experience() (world_models.training.train_iris.IRISTrainer method) collect_random_episodes() (world_models.models.dreamer.Dreamer method) collect_steps (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) collect_system_stats() (in module world_models.utils.logging_utils) collect_trajectories() (world_models.training.rl_harness.PPOTrainer method) combine_videos() (in module world_models.utils.utils) commitment_weight (torchwm.LatentActionModelConfig attribute) (torchwm.VideoTokenizerConfig attribute) (world_models.configs.genie_config.LatentActionModelConfig attribute) (world_models.configs.genie_config.VideoTokenizerConfig attribute) (world_models.configs.LatentActionModelConfig attribute) (world_models.configs.VideoTokenizerConfig attribute) (world_models.LatentActionModelConfig attribute) (world_models.VideoTokenizerConfig attribute)

compute_aggregate_metrics() (in module world_models.benchmarks.metrics) compute_gae() (world_models.training.rl_harness.PPOTrainer method) compute_lambda_return() (in module torchwm) (in module world_models) (in module world_models.models.iris_agent) compute_lambda_returns() (world_models.models.diffusion.actor_critic.RLLoss method) compute_return() (in module world_models.utils.dreamer_utils) config_path (torchwm.ModelSpec attribute) (world_models.api.ModelSpec attribute) (world_models.ModelSpec attribute) Controller (class in world_models.models.controller) conv_vae_loss_fn() (in module world_models.losses.convae_loss) ConvBlock (class in world_models.models.diffusion.reward_termination) ConvDecoder (class in torchwm) (class in world_models) (class in world_models.models.modular_rssm) (class in world_models.vision.dreamer_decoder) ConvEmbed (class in world_models.models.vit) ConvEncoder (class in torchwm) (class in world_models) (class in world_models.models.modular_rssm) (class in world_models.vision.dreamer_encoder) ConvVAE (class in world_models.vision.VAE.ConvVAE) ConvVAEDecoder (class in world_models.vision.VAE.ConvVAE) ConvVAEEncoder (class in world_models.vision.VAE.ConvVAE) copy_imgnt_locally() (in module world_models.datasets.imagenet1k) CosineWDSchedule (class in world_models.utils.jepa_utils) create_config() (in module torchwm) (in module world_models) (in module world_models.api) create_dataloader() (world_models.datasets.tinyworlds.TinyWorldsDataLoader static method) create_dit() (in module torchwm) (in module world_models) (in module world_models.models.diffusion.DiT) create_dynamics_model() (in module torchwm) (in module world_models) (in module world_models.models.dynamics_model) create_efficient_dataloader() (in module world_models.utils.data_utils) create_genie() (in module torchwm) (in module world_models) (in module world_models.models.genie) create_genie_large() (in module torchwm) (in module world_models) (in module world_models.models.genie) create_genie_small() (in module torchwm) (in module world_models) (in module world_models.models.genie) create_genie_trainer() (in module world_models.training.train_genie) create_latent_action_model() (in module torchwm) (in module world_models) (in module world_models.models.latent_action_model) create_model() (in module torchwm) (in module world_models) (in module world_models.api) create_modular_rssm() (in module torchwm) (in module world_models) (in module world_models.models.modular_rssm) create_rl_harness_example() (in module world_models.training.rl_harness) create_st_transformer() (in module world_models.blocks.st_transformer) create_tinyworlds_dataloader() (in module world_models.datasets.tinyworlds) create_video_dataloader() (in module world_models.datasets.video_datasets) create_video_dataset_from_config() (in module world_models.datasets.video_datasets) create_video_tokenizer() (in module torchwm) (in module world_models) (in module world_models.vision.video_tokenizer) CSVLogger (class in world_models.utils.jepa_utils)

D

d (torchwm.RSSMPolicy attribute) (world_models.controller.rssm_policy.RSSMPolicy attribute) (world_models.RSSMPolicy attribute) data_dir (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.configs.wm_config.WMMDNRNNConfig attribute) (world_models.configs.wm_config.WMVAEConfig attribute) (world_models.DreamerConfig attribute) data_loader_num_workers (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) data_pass() (in module world_models.training.train_mdn_rnn) data_source (world_models.datasets.video_datasets.VideoDatasetBase attribute) (world_models.datasets.video_datasets.VideoDatasetConfig attribute) DATASET (torchwm.DiTConfig attribute) (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) (world_models.DiTConfig attribute) DATASET_CONFIGS (world_models.datasets.tinyworlds.TinyWorldsDataset attribute) dataset_name (world_models.datasets.tinyworlds.TinyWorldsConfig attribute) DATASET_NAMES (world_models.datasets.tinyworlds.TinyWorldsDataLoader attribute) dataset_type (world_models.datasets.video_datasets.VideoDatasetConfig attribute) DatasetConfig (class in world_models.datasets.video_datasets) DDPM (class in torchwm) (class in world_models) (class in world_models.models.diffusion.DDPM) decode() (torchwm.LatentActionModel method) (torchwm.VideoTokenizer method) (world_models.LatentActionModel method) (world_models.models.latent_action_model.LatentActionModel method) (world_models.utils.dreamer_utils.TwoHotEncoder method) (world_models.VideoTokenizer method) (world_models.vision.iris_decoder.DiscreteAutoencoder method) (world_models.vision.video_tokenizer.VideoTokenizer method) decode_from_embeddings() (torchwm.IRISDecoder method) (world_models.IRISDecoder method) (world_models.vision.iris_decoder.IRISDecoder method) decode_from_indices() (torchwm.IRISDecoder method) (torchwm.IRISEncoder method) (world_models.IRISDecoder method) (world_models.IRISEncoder method) (world_models.vision.iris_decoder.IRISDecoder method) (world_models.vision.iris_encoder.IRISEncoder method) decode_indices() (torchwm.VectorQuantizer method) (torchwm.VectorQuantizerEMA method) (torchwm.VideoTokenizer method) (world_models.VectorQuantizer method) (world_models.VectorQuantizerEMA method) (world_models.VideoTokenizer method) (world_models.vision.video_tokenizer.VideoTokenizer method) (world_models.vision.vq_layer.VectorQuantizer method) (world_models.vision.vq_layer.VectorQuantizerEMA method) decode_observation() (torchwm.ModularRSSM method) (world_models.models.modular_rssm.ModularRSSM method) (world_models.ModularRSSM method) decode_reward() (torchwm.ModularRSSM method) (world_models.models.modular_rssm.ModularRSSM method) (world_models.ModularRSSM method) decoder (world_models.vision.VAE.ConvVAE.ConvVAE attribute) decoder_depth (torchwm.VideoTokenizerConfig attribute) (world_models.configs.genie_config.VideoTokenizerConfig attribute) (world_models.configs.VideoTokenizerConfig attribute) (world_models.VideoTokenizerConfig attribute) decoder_dim (torchwm.VideoTokenizerConfig attribute) (world_models.configs.genie_config.VideoTokenizerConfig attribute) (world_models.configs.VideoTokenizerConfig attribute) (world_models.VideoTokenizerConfig attribute) DecoderBase (class in world_models.models.modular_rssm) DeepMindControlEnv (class in world_models.envs) (class in world_models.envs.dmc) DefaultCollator (class in world_models.masks) (class in world_models.masks.default) denoise() (world_models.models.diffusion.diamond_diffusion.EDMPreconditioner method) dense_activation_function (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) DenseDecoder (class in torchwm) (class in world_models) (class in world_models.vision.dreamer_decoder) DEPTH (torchwm.DiTConfig attribute) depth (torchwm.DynamicsModelConfig attribute) (torchwm.STTransformerConfig attribute) DEPTH (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) depth (world_models.configs.DynamicsModelConfig attribute) (world_models.configs.genie_config.DynamicsModelConfig attribute) (world_models.configs.genie_config.STTransformerConfig attribute) (world_models.configs.STTransformerConfig attribute) DEPTH (world_models.DiTConfig attribute) depth (world_models.DynamicsModelConfig attribute) (world_models.STTransformerConfig attribute) description (torchwm.EnvBackendSpec attribute) (torchwm.ModelSpec attribute) (world_models.api.EnvBackendSpec attribute) (world_models.api.ModelSpec attribute) (world_models.EnvBackendSpec attribute) (world_models.ModelSpec attribute) detach_state() (torchwm.ModularRSSM method) (torchwm.RSSM method) (world_models.models.dreamer_rssm.RSSM method) (world_models.models.modular_rssm.ModularRSSM method) (world_models.ModularRSSM method) (world_models.RSSM method) detect_anomaly (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) deter_size (torchwm.DreamerConfig attribute) (torchwm.ModularRSSM property) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) (world_models.models.modular_rssm.BackboneBase attribute) (world_models.models.modular_rssm.ModularRSSM property) (world_models.ModularRSSM property) deterministic_state_fwd() (torchwm.RecurrentStateSpaceModel method) (world_models.models.rssm.RecurrentStateSpaceModel method) (world_models.RecurrentStateSpaceModel method) device (torchwm.DiamondConfig attribute) (torchwm.RSSMPolicy attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.configs.wm_config.WMMDNRNNConfig attribute) (world_models.configs.wm_config.WMVAEConfig attribute) (world_models.controller.rollout_generator.RolloutGenerator attribute) (world_models.controller.rssm_policy.RSSMPolicy attribute) (world_models.DiamondConfig attribute) (world_models.RSSMPolicy attribute) DiamondAdapter (class in world_models.benchmarks.adapters) DiamondAgent (class in world_models.training.train_diamond) DiamondAtariWrapper (class in world_models.envs.diamond_atari) DiamondConfig (class in torchwm) (class in world_models) (class in world_models.configs) (class in world_models.configs.diamond_config) diffusion_channels (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.diamond_config.ModelPreset attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute)

diffusion_cond_dim (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.diamond_config.ModelPreset attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) diffusion_res_blocks (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.diamond_config.ModelPreset attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) DiffusionUNet (class in world_models.models.diffusion.diamond_diffusion) dim (torchwm.DynamicsModelConfig attribute) (torchwm.STTransformerConfig attribute) (world_models.configs.DynamicsModelConfig attribute) (world_models.configs.genie_config.DynamicsModelConfig attribute) (world_models.configs.genie_config.STTransformerConfig attribute) (world_models.configs.STTransformerConfig attribute) (world_models.DynamicsModelConfig attribute) (world_models.STTransformerConfig attribute) disc_loss_coeff (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) discount (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) discount_factor (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) DiscreteAutoencoder (class in world_models.vision.iris_decoder) display (world_models.configs.wm_config.WMControllerConfig attribute) DiT (class in torchwm) (class in world_models) (class in world_models.models.diffusion.DiT) DiTConfig (class in torchwm) (class in world_models) (class in world_models.configs) (class in world_models.configs.dit_config) dmlab_action_repeat (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) dmlab_action_set (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) dmlab_config (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) dmlab_observations (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) dmlab_renderer (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) DMLabEnv (class in world_models.envs) (class in world_models.envs.dmlab) DownBlock (class in world_models.models.diffusion.diamond_diffusion) download_all_datasets() (in module world_models.datasets.tinyworlds) Dreamer (class in world_models.models.dreamer) DreamerAdapter (class in world_models.benchmarks.adapters) DreamerAgent (class in world_models.models.dreamer) DreamerConfig (class in torchwm) (class in world_models) (class in world_models.configs) (class in world_models.configs.dreamer_config) DreamerOperator (class in torchwm) (class in world_models) (class in world_models.inference.operators.dreamer_operator) DreamerPolicyExport (class in world_models.export) DreamerRewardModel (in module torchwm) (in module world_models) DreamerV1Adapter (class in world_models.benchmarks.adapters) DreamerV2Adapter (class in world_models.benchmarks.adapters) DreamerValueModel (in module torchwm) (in module world_models) DROP (torchwm.DiTConfig attribute) (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) (world_models.DiTConfig attribute) drop_path() (in module world_models.models.vit) drop_path_rate (torchwm.DynamicsModelConfig attribute) (torchwm.STTransformerConfig attribute) (world_models.configs.DynamicsModelConfig attribute) (world_models.configs.genie_config.DynamicsModelConfig attribute) (world_models.configs.genie_config.STTransformerConfig attribute) (world_models.configs.STTransformerConfig attribute) (world_models.DynamicsModelConfig attribute) (world_models.STTransformerConfig attribute) drop_rate (torchwm.DynamicsModelConfig attribute) (torchwm.STTransformerConfig attribute) (world_models.configs.DynamicsModelConfig attribute) (world_models.configs.genie_config.DynamicsModelConfig attribute) (world_models.configs.genie_config.STTransformerConfig attribute) (world_models.configs.STTransformerConfig attribute) (world_models.DynamicsModelConfig attribute) (world_models.STTransformerConfig attribute) DropPath (class in world_models.blocks.st_transformer) (class in world_models.models.vit) dtype (torchwm.TensorSpec attribute) (world_models.inference.operators.base.TensorSpec attribute) (world_models.TensorSpec attribute) dump_scalars_to_pickle() (world_models.utils.dreamer_utils.Logger method) dynamics_depth (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute) dynamics_dim (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute) dynamics_num_heads (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute) DynamicsModel (class in torchwm) (class in world_models) (class in world_models.models.dynamics_model) DynamicsModelConfig (class in torchwm) (class in world_models) (class in world_models.configs) (class in world_models.configs.genie_config)

E

early_stopping_patience (world_models.configs.wm_config.WMMDNRNNConfig attribute) (world_models.configs.wm_config.WMVAEConfig attribute) EarlyStopping (class in world_models.utils.train_utils) EDMPreconditioner (class in world_models.models.diffusion.diamond_diffusion) ego_future (world_models.datasets.nuplan.NuPlanSample attribute) ego_past (world_models.datasets.nuplan.NuPlanSample attribute) EMA (torchwm.DiTConfig attribute) (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) (world_models.DiTConfig attribute) EMA_DECAY (torchwm.DiTConfig attribute) ema_decay (torchwm.VideoTokenizerConfig attribute) EMA_DECAY (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) ema_decay (world_models.configs.genie_config.VideoTokenizerConfig attribute) (world_models.configs.VideoTokenizerConfig attribute) EMA_DECAY (world_models.DiTConfig attribute) ema_decay (world_models.VideoTokenizerConfig attribute) embed_size (torchwm.ModularRSSM property) (world_models.models.modular_rssm.EncoderBase attribute) (world_models.models.modular_rssm.ModularRSSM property) (world_models.ModularRSSM property) embedding_dim (torchwm.DynamicsModelConfig attribute) (torchwm.LatentActionModelConfig attribute) (torchwm.VideoTokenizerConfig attribute) (world_models.configs.DynamicsModelConfig attribute) (world_models.configs.genie_config.DynamicsModelConfig attribute) (world_models.configs.genie_config.LatentActionModelConfig attribute) (world_models.configs.genie_config.VideoTokenizerConfig attribute) (world_models.configs.LatentActionModelConfig attribute) (world_models.configs.VideoTokenizerConfig attribute) (world_models.DynamicsModelConfig attribute) (world_models.LatentActionModelConfig attribute) (world_models.VideoTokenizerConfig attribute) embedding_size (world_models.models.modular_rssm.GRUBackbone property) (world_models.models.modular_rssm.LSTMBackbone property) (world_models.models.modular_rssm.TransformerBackbone property) enable_console_metrics (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) enable_jsonl (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) enable_mixed_precision() (in module world_models.utils.memory_utils) enable_tensorboard (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) enable_wandb (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) encode() (torchwm.LatentActionModel method) (torchwm.VideoTokenizer method) (world_models.LatentActionModel method) (world_models.models.latent_action_model.LatentActionModel method) (world_models.utils.dreamer_utils.TwoHotEncoder method) (world_models.VideoTokenizer method) (world_models.vision.iris_decoder.DiscreteAutoencoder method) (world_models.vision.video_tokenizer.VideoTokenizer method) encode_to_indices() (torchwm.IRISEncoder method) (world_models.IRISEncoder method) (world_models.vision.iris_encoder.IRISEncoder method) encode_video() (torchwm.Genie method) (world_models.Genie method) (world_models.models.genie.Genie method) encoder (world_models.vision.VAE.ConvVAE.ConvVAE attribute) encoder_depth (torchwm.LatentActionModelConfig attribute) (torchwm.VideoTokenizerConfig attribute) (world_models.configs.genie_config.LatentActionModelConfig attribute) (world_models.configs.genie_config.VideoTokenizerConfig attribute) (world_models.configs.LatentActionModelConfig attribute) (world_models.configs.VideoTokenizerConfig attribute) (world_models.LatentActionModelConfig attribute) (world_models.VideoTokenizerConfig attribute) encoder_dim (torchwm.LatentActionModelConfig attribute) (torchwm.VideoTokenizerConfig attribute) (world_models.configs.genie_config.LatentActionModelConfig attribute) (world_models.configs.genie_config.VideoTokenizerConfig attribute) (world_models.configs.LatentActionModelConfig attribute) (world_models.configs.VideoTokenizerConfig attribute) (world_models.LatentActionModelConfig attribute) (world_models.VideoTokenizerConfig attribute) EncoderBase (class in world_models.models.modular_rssm) ensure_results_dir_exists() (in module world_models.utils.utils) entropy() (torchwm.SampleDist method) (world_models.SampleDist method) (world_models.vision.dreamer_decoder.SampleDist method)

entropy_weight (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) env (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.controller.rollout_generator.RolloutGenerator attribute) (world_models.DreamerConfig attribute) env_backend (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) env_instance (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) EnvBackendSpec (class in torchwm) (class in world_models) (class in world_models.api) environment_steps_per_epoch (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) Episode (class in torchwm) (class in world_models) (class in world_models.memory.dreamer_memory) (class in world_models.memory.planet_memory) episode_gen (world_models.controller.rollout_generator.RolloutGenerator attribute) episodes (torchwm.IRISReplayBuffer attribute) (torchwm.Memory attribute) (world_models.IRISReplayBuffer attribute) (world_models.Memory attribute) (world_models.memory.iris_memory.IRISReplayBuffer attribute) (world_models.memory.planet_memory.Memory attribute) EPOCHS (torchwm.DiTConfig attribute) (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) (world_models.DiTConfig attribute) eps_lengths (torchwm.Memory attribute) (world_models.Memory attribute) (world_models.memory.planet_memory.Memory attribute) epsilon_greedy (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) EulerSampler (class in world_models.models.diffusion.diamond_diffusion) eval_interval (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) evaluate (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) evaluate() (in module world_models.training.train_controller) (in module world_models.training.train_rssm) (world_models.benchmarks.adapters.BaseAdapter method) (world_models.benchmarks.adapters.DiamondAdapter method) (world_models.benchmarks.adapters.DreamerAdapter method) (world_models.benchmarks.adapters.IRISAdapter method) (world_models.models.dreamer.Dreamer method) (world_models.models.dreamer.DreamerAgent method) (world_models.training.train_diamond.DiamondAgent method) (world_models.training.train_iris.IRISTrainer method) exp_name (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) experience_replay (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) export() (torchwm.ExportableAgentMixin method) (world_models.export.ExportableAgentMixin method) (world_models.ExportableAgentMixin method) export_any() (in module torchwm) (in module world_models) (in module world_models.export) export_csv() (in module world_models.benchmarks.reporting) export_latex() (in module world_models.benchmarks.reporting) export_markdown() (in module world_models.benchmarks.reporting) export_model() (in module torchwm) (in module world_models) (in module world_models.export) ExportableAgentMixin (class in torchwm) (class in world_models) (class in world_models.export) extensions (world_models.datasets.video_datasets.VideoDatasetConfig attribute) extra (world_models.configs.wm_config.WMMDNRNNConfig attribute) (world_models.configs.wm_config.WMVAEConfig attribute)

F

factory_path (torchwm.EnvBackendSpec attribute) (world_models.api.EnvBackendSpec attribute) (world_models.EnvBackendSpec attribute) 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) flatten_parameters() (in module world_models.training.train_controller) flush() (world_models.utils.dreamer_utils.Logger method) (world_models.utils.logging_utils.MetricsLogger method) forward() (torchwm.ActionDecoder method) (torchwm.ActorCriticNetwork method) (torchwm.AdaLNNormalization method) (torchwm.CNNDecoder method) (torchwm.CNNEncoder method) (torchwm.CNNFeatureExtractor method) (torchwm.ConvDecoder method) (torchwm.ConvEncoder method) (torchwm.DenseDecoder method) (torchwm.DiT method) (torchwm.DynamicsModel method) (torchwm.Genie method) (torchwm.IRISActor method) (torchwm.IRISCritic method) (torchwm.IRISDecoder method) (torchwm.IRISEncoder method) (torchwm.IRISPolicy method) (torchwm.LatentActionModel method) (torchwm.MultiHeadSelfAttention method) (torchwm.OperatorABC method) (torchwm.PatchEmbed method) (torchwm.PatchUnEmbed method) (torchwm.RecurrentStateSpaceModel method) (torchwm.RewardModel method) (torchwm.RewardTerminationModel method) (torchwm.RMSNorm method) (torchwm.RSSM method) (torchwm.STTransformer method) (torchwm.ValueModel method) (torchwm.VectorQuantizer method) (torchwm.VectorQuantizerEMA method) (torchwm.VideoTokenizer method) (world_models.ActionDecoder method) (world_models.ActorCriticNetwork method) (world_models.AdaLNNormalization method) (world_models.blocks.mhsa.MultiHeadSelfAttention method) (world_models.blocks.st_transformer.DropPath method) (world_models.blocks.st_transformer.STMLP method) (world_models.blocks.st_transformer.STSpatialAttention method) (world_models.blocks.st_transformer.STTemporalAttention method) (world_models.blocks.st_transformer.STTransformer method) (world_models.blocks.st_transformer.STTransformerBlock method) (world_models.CNNDecoder method) (world_models.CNNEncoder method) (world_models.CNNFeatureExtractor method) (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.ConvDecoder method) (world_models.ConvEncoder method) (world_models.DenseDecoder method) (world_models.DiT method) (world_models.DynamicsModel method) (world_models.export.DreamerPolicyExport method) (world_models.export.IRISActorCriticExport method) (world_models.Genie method) (world_models.inference.operators.base.OperatorABC method) (world_models.IRISActor method) (world_models.IRISCritic method) (world_models.IRISDecoder method) (world_models.IRISEncoder method) (world_models.IRISPolicy method) (world_models.LatentActionModel method) (world_models.models.controller.Controller method) (world_models.models.diffusion.actor_critic.ActorCriticNetwork method) (world_models.models.diffusion.diamond_diffusion.AdaptiveGroupNorm method) (world_models.models.diffusion.diamond_diffusion.AttentionBlock method) (world_models.models.diffusion.diamond_diffusion.DiffusionUNet method) (world_models.models.diffusion.diamond_diffusion.DownBlock method) (world_models.models.diffusion.diamond_diffusion.ResBlock method) (world_models.models.diffusion.diamond_diffusion.TimestepEmbedding method) (world_models.models.diffusion.diamond_diffusion.UpBlock 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.diffusion.reward_termination.ConvBlock method) (world_models.models.diffusion.reward_termination.RewardTerminationLoss method) (world_models.models.diffusion.reward_termination.RewardTerminationModel method) (world_models.models.dreamer_rssm.RSSM method) (world_models.models.dynamics_model.DynamicsModel method) (world_models.models.genie.Genie method) (world_models.models.iris_transformer.IRISTransformer method) (world_models.models.iris_transformer.IRISWorldModel method) (world_models.models.latent_action_model.LatentActionModel method) (world_models.models.mdrnn.MDRNN method) (world_models.models.mdrnn.MDRNNCell method) (world_models.models.modular_rssm.BackboneBase method) (world_models.models.modular_rssm.ConvDecoder method) (world_models.models.modular_rssm.ConvEncoder method) (world_models.models.modular_rssm.DecoderBase method) (world_models.models.modular_rssm.EncoderBase method) (world_models.models.modular_rssm.GRUBackbone method) (world_models.models.modular_rssm.LSTMBackbone method) (world_models.models.modular_rssm.MLPDecoder method) (world_models.models.modular_rssm.MLPEncoder method) (world_models.models.modular_rssm.TransformerBackbone method) (world_models.models.modular_rssm.TransformerBlock method) (world_models.models.modular_rssm.ViTEncoder method) (world_models.models.rssm.RecurrentStateSpaceModel 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.MultiHeadSelfAttention method) (world_models.OperatorABC method) (world_models.PatchEmbed method) (world_models.PatchUnEmbed method) (world_models.RecurrentStateSpaceModel method) (world_models.RewardModel method) (world_models.RewardTerminationModel method) (world_models.RMSNorm method) (world_models.RSSM method) (world_models.STTransformer method) (world_models.training.rl_harness.ActorCritic 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.ValueModel method) (world_models.VectorQuantizer method) (world_models.VectorQuantizerEMA method) (world_models.VideoTokenizer 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.planet_decoder.CNNDecoder method) (world_models.vision.planet_encoder.CNNEncoder method) (world_models.vision.VAE.ConvVAE.ConvVAE method) (world_models.vision.VAE.ConvVAE.ConvVAEDecoder method) (world_models.vision.VAE.ConvVAE.ConvVAEEncoder method) (world_models.vision.video_tokenizer.VideoTokenizer method) (world_models.vision.vq_layer.VectorQuantizer method) (world_models.vision.vq_layer.VectorQuantizerEMA method)

forward_actor_critic() (torchwm.IRISAgent method) (world_models.IRISAgent method) (world_models.models.iris_agent.IRISAgent method) frame_shape (torchwm.CNNFeatureExtractor attribute) (torchwm.IRISActor attribute) (torchwm.IRISCritic attribute) (torchwm.IRISPolicy attribute) (world_models.CNNFeatureExtractor attribute) (world_models.controller.iris_policy.CNNFeatureExtractor attribute) (world_models.controller.iris_policy.IRISActor attribute) (world_models.controller.iris_policy.IRISCritic attribute) (world_models.controller.iris_policy.IRISPolicy attribute) (world_models.IRISActor attribute) (world_models.IRISCritic attribute) (world_models.IRISPolicy attribute) frameskip (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) free_nats (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) FreezeParameters (class in world_models.utils.dreamer_utils) from_config() (torchwm.DiT class method) (torchwm.Genie class method) (torchwm.IRISAgent class method) (world_models.DiT class method) (world_models.Genie class method) (world_models.IRISAgent class method) (world_models.models.diffusion.DiT.DiT class method) (world_models.models.dreamer.Dreamer class method) (world_models.models.dreamer.DreamerAgent class method) (world_models.models.genie.Genie class method) (world_models.models.iris_agent.IRISAgent class method) (world_models.models.jepa_agent.JEPAAgent class method) (world_models.training.train_diamond.DiamondAgent class method) from_dict() (torchwm.JEPAConfig class method) (world_models.configs.jepa_config.JEPAConfig class method) (world_models.configs.JEPAConfig class method) (world_models.JEPAConfig class method) from_pretrained() (torchwm.DiT class method) (torchwm.Genie class method) (torchwm.IRISAgent class method) (world_models.DiT class method) (world_models.Genie class method) (world_models.IRISAgent class method) (world_models.models.diffusion.DiT.DiT class method) (world_models.models.dreamer.Dreamer class method) (world_models.models.dreamer.DreamerAgent class method) (world_models.models.genie.Genie class method) (world_models.models.iris_agent.IRISAgent class method) (world_models.models.jepa_agent.JEPAAgent class method) (world_models.training.train_diamond.DiamondAgent class method)

G

game (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) GaussianBlur (class in world_models.transforms.image) generate() (torchwm.Genie method) (world_models.Genie method) (world_models.models.genie.Genie method) generate_rollouts() (in module world_models.training.train_world_model) Genie (class in torchwm) (class in world_models) (class in world_models.models.genie) GenieConfig (class in torchwm) (class in world_models) (class in world_models.configs) (class in world_models.configs.genie_config) (class in world_models.training.train_genie) GenieSmallConfig (class in torchwm) (class in world_models) (class in world_models.configs) (class in world_models.configs.genie_config) GenieTrainer (class in world_models.training.train_genie) 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() (torchwm.ActorCriticNetwork method) (torchwm.IRISActor method) (world_models.ActorCriticNetwork method) (world_models.controller.iris_policy.IRISActor method) (world_models.IRISActor method) (world_models.models.diffusion.actor_critic.ActorCriticNetwork method) (world_models.training.rl_harness.ActorCritic method) get_actions() (torchwm.ActorCriticNetwork method) (world_models.ActorCriticNetwork method) (world_models.models.diffusion.actor_critic.ActorCriticNetwork method) get_arrays() (torchwm.IRISOnPolicyBuffer method) (world_models.IRISOnPolicyBuffer method) (world_models.memory.iris_memory.IRISOnPolicyBuffer method) get_autoencoder_config() (torchwm.IRISConfig method) (world_models.configs.iris_config.IRISConfig method) (world_models.configs.IRISConfig method) (world_models.IRISConfig method) get_available_memory() (in module world_models.models.dreamer) get_combined_params() (in module world_models.utils.utils) get_dataset_info() (world_models.datasets.tinyworlds.TinyWorldsDataLoader static method) get_default_device() (in module world_models.configs.diamond_config) get_dist() (torchwm.ModularRSSM method) (torchwm.RSSM method) (world_models.models.dreamer_rssm.RSSM method) (world_models.models.modular_rssm.ModularRSSM method) (world_models.ModularRSSM method) (world_models.RSSM method) get_dit_config() (in module torchwm) (in module world_models) (in module world_models.configs) (in module world_models.configs.dit_config) get_embed_size() (world_models.models.modular_rssm.EncoderBase method) get_env_backend_spec() (in module torchwm) (in module world_models) (in module world_models.api)

get_epsilon() (world_models.training.train_iris.IRISTrainer method) get_frame_shape() (torchwm.IRISConfig method) (world_models.configs.iris_config.IRISConfig method) (world_models.configs.IRISConfig method) (world_models.IRISConfig method) get_hidden_size() (torchwm.ActorCriticNetwork method) (world_models.ActorCriticNetwork method) (world_models.models.diffusion.actor_critic.ActorCriticNetwork method) get_info() (world_models.datasets.tinyworlds.TinyWorldsDataset method) get_init_hidden() (world_models.models.mdrnn.MDRNN method) (world_models.models.mdrnn.MDRNNCell method) get_init_state() (torchwm.RecurrentStateSpaceModel method) (world_models.models.rssm.RecurrentStateSpaceModel method) (world_models.RecurrentStateSpaceModel method) get_loss() (in module world_models.training.train_mdn_rnn) get_mask() (in module world_models.utils.utils) get_mask_prob() (world_models.models.dynamics_model.MaskGITSampler method) get_model_spec() (in module torchwm) (in module world_models) (in module world_models.api) get_num_parameters() (torchwm.Genie method) (world_models.Genie method) (world_models.models.genie.Genie method) get_operator() (in module torchwm) (in module world_models) get_package_logger() (in module world_models.utils.logging_utils) get_parameters() (in module world_models.utils.dreamer_utils) get_posterior() (torchwm.RSSM method) (world_models.models.dreamer_rssm.RSSM method) (world_models.RSSM method) get_preconditioners() (world_models.models.diffusion.diamond_diffusion.EDMPreconditioner method) get_prior() (torchwm.RSSM method) (world_models.models.dreamer_rssm.RSSM method) (world_models.RSSM method) get_rl_config() (torchwm.IRISConfig method) (world_models.configs.iris_config.IRISConfig method) (world_models.configs.IRISConfig method) (world_models.IRISConfig method) get_transformer_config() (torchwm.IRISConfig method) (world_models.configs.iris_config.IRISConfig method) (world_models.configs.IRISConfig method) (world_models.IRISConfig method) get_value() (torchwm.ActorCriticNetwork method) (world_models.ActorCriticNetwork method) (world_models.models.diffusion.actor_critic.ActorCriticNetwork method) gmm_components (world_models.configs.wm_config.WMMDNRNNConfig attribute) gmm_loss() (in module world_models.losses.gmm_loss) gpu_timer() (in module world_models.utils.jepa_utils) grad_clip_norm (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) grad_logger() (in module world_models.utils.jepa_utils) GRUBackbone (class in world_models.models.modular_rssm) gym_render_mode (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) GymImageEnv (class in world_models.envs) (class in world_models.envs.gym_env)

H

H (torchwm.RSSMPolicy attribute) (world_models.controller.rssm_policy.RSSMPolicy attribute) (world_models.RSSMPolicy attribute) HDF5Dataset (class in world_models.datasets.video_datasets) HEADS (torchwm.DiTConfig attribute) (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) (world_models.DiTConfig attribute) height (world_models.configs.wm_config.WMVAEConfig attribute) hidden_size (torchwm.IRISActor attribute) (torchwm.IRISCritic attribute) (torchwm.IRISPolicy attribute) (world_models.configs.wm_config.WMControllerConfig attribute) (world_models.configs.wm_config.WMMDNRNNConfig attribute) (world_models.controller.iris_policy.IRISActor attribute) (world_models.controller.iris_policy.IRISCritic attribute) (world_models.controller.iris_policy.IRISPolicy attribute) (world_models.IRISActor attribute) (world_models.IRISCritic attribute) (world_models.IRISPolicy attribute) (world_models.models.controller.Controller attribute)

I

image_size (torchwm.DreamerConfig attribute) (torchwm.DynamicsModelConfig attribute) (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (torchwm.LatentActionModelConfig attribute) (torchwm.VideoTokenizerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.configs.DynamicsModelConfig attribute) (world_models.configs.genie_config.DynamicsModelConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.genie_config.LatentActionModelConfig attribute) (world_models.configs.genie_config.VideoTokenizerConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.configs.LatentActionModelConfig attribute) (world_models.configs.VideoTokenizerConfig attribute) (world_models.datasets.tinyworlds.TinyWorldsConfig attribute) (world_models.datasets.video_datasets.DatasetConfig attribute) (world_models.DreamerConfig attribute) (world_models.DynamicsModelConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.LatentActionModelConfig attribute) (world_models.training.train_genie.GenieConfig attribute) (world_models.VideoTokenizerConfig attribute) ImageFolderDataset (class in world_models.datasets.video_datasets) ImageNet (class in world_models.datasets.imagenet1k) ImageNetSubset (class in world_models.datasets.imagenet1k) imagination_horizon (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) imagine() (world_models.models.iris_transformer.IRISWorldModel method) imagine_horizon (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) imagine_rollout() (torchwm.IRISAgent method) (torchwm.ModularRSSM method) (torchwm.RSSM method) (world_models.IRISAgent method) (world_models.models.dreamer_rssm.RSSM method) (world_models.models.iris_agent.IRISAgent method) (world_models.models.modular_rssm.ModularRSSM method) (world_models.ModularRSSM method) (world_models.RSSM method) imagine_step() (torchwm.ModularRSSM method) (torchwm.RSSM method) (world_models.models.dreamer_rssm.RSSM method) (world_models.models.modular_rssm.ModularRSSM method) (world_models.ModularRSSM method) (world_models.RSSM method) img_channels (world_models.vision.VAE.ConvVAE.ConvVAEDecoder attribute) (world_models.vision.VAE.ConvVAE.ConvVAEEncoder attribute) IMG_SIZE (torchwm.DiTConfig attribute) (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) (world_models.DiTConfig attribute) import_path (torchwm.ModelSpec attribute) (world_models.api.ModelSpec attribute) (world_models.ModelSpec attribute) in_channels (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (torchwm.LatentActionModelConfig attribute) (torchwm.VideoTokenizerConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.genie_config.LatentActionModelConfig attribute) (world_models.configs.genie_config.VideoTokenizerConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.configs.LatentActionModelConfig attribute) (world_models.configs.VideoTokenizerConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.LatentActionModelConfig attribute) (world_models.training.train_genie.GenieConfig attribute) (world_models.VideoTokenizerConfig attribute)

include_reward (world_models.configs.wm_config.WMMDNRNNConfig attribute) infer_actions() (torchwm.Genie method) (world_models.Genie method) (world_models.models.genie.Genie method) info (torchwm.Episode attribute) (world_models.Episode attribute) (world_models.memory.planet_memory.Episode attribute) init_distributed() (in module world_models.utils.jepa_utils) init_hidden() (torchwm.ActorCriticNetwork method) (torchwm.IRISPolicy method) (torchwm.RewardTerminationModel method) (world_models.ActorCriticNetwork method) (world_models.controller.iris_policy.IRISPolicy method) (world_models.IRISPolicy method) (world_models.models.diffusion.actor_critic.ActorCriticNetwork method) (world_models.models.diffusion.reward_termination.RewardTerminationModel method) (world_models.RewardTerminationModel method) init_hidden_state() (torchwm.IRISActor method) (torchwm.IRISCritic method) (world_models.controller.iris_policy.IRISActor method) (world_models.controller.iris_policy.IRISCritic method) (world_models.IRISActor method) (world_models.IRISCritic method) init_model() (in module world_models.helpers.jepa_helper) init_opt() (in module world_models.helpers.jepa_helper) init_state() (torchwm.ModularRSSM method) (torchwm.RSSM method) (world_models.models.dreamer_rssm.RSSM method) (world_models.models.modular_rssm.BackboneBase method) (world_models.models.modular_rssm.GRUBackbone method) (world_models.models.modular_rssm.LSTMBackbone method) (world_models.models.modular_rssm.ModularRSSM method) (world_models.models.modular_rssm.TransformerBackbone method) (world_models.ModularRSSM method) (world_models.RSSM method) input_specs (torchwm.OperatorABC attribute) (world_models.inference.operators.base.OperatorABC attribute) (world_models.OperatorABC attribute) install_export_method() (in module world_models.export) interpolate_pos_encoding() (world_models.models.vit.VisionTransformer method) iqm_of_array() (in module world_models.benchmarks.metrics) IRISActor (class in torchwm) (class in world_models) (class in world_models.controller.iris_policy) IRISActorCriticExport (class in world_models.export) IRISAdapter (class in world_models.benchmarks.adapters) IRISAgent (class in torchwm) (class in world_models) (class in world_models.models.iris_agent) IRISConfig (class in torchwm) (class in world_models) (class in world_models.configs) (class in world_models.configs.iris_config) IRISCritic (class in torchwm) (class in world_models) (class in world_models.controller.iris_policy) IRISDecoder (class in torchwm) (class in world_models) (class in world_models.vision.iris_decoder) IRISEncoder (class in torchwm) (class in world_models) (class in world_models.vision.iris_encoder) IRISOnPolicyBuffer (class in torchwm) (class in world_models) (class in world_models.memory.iris_memory) IrisOperator (class in torchwm) (class in world_models) (class in world_models.inference.operators.iris_operator) IRISPolicy (class in torchwm) (class in world_models) (class in world_models.controller.iris_policy) IRISReplayBuffer (class in torchwm) (class in world_models) (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) is_moved_mujoco_error() (in module world_models.envs.robotics_env) is_ready() (world_models.datasets.diamond_dataset.ReplayBuffer method)

J

JEPAAgent (class in world_models.models.jepa_agent) JEPAConfig (class in torchwm) (class in world_models) (class in world_models.configs) (class in world_models.configs.jepa_config) JEPAOperator (class in torchwm) (class in world_models) (class in world_models.inference.operators.jepa_operator)

jit_compile_function() (in module world_models.utils.jit_utils) jit_compile_module() (in module world_models.utils.jit_utils) jsonl_filename (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute)

K

K (torchwm.RSSMPolicy attribute) (world_models.controller.rssm_policy.RSSMPolicy attribute) (world_models.RSSMPolicy attribute) kl_alpha (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute)

kl_loss_coeff (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute)

L

lambda_returns (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) latent_size (torchwm.RSSMPolicy attribute) (world_models.configs.wm_config.WMControllerConfig attribute) (world_models.configs.wm_config.WMMDNRNNConfig attribute) (world_models.configs.wm_config.WMVAEConfig attribute) (world_models.controller.rssm_policy.RSSMPolicy attribute) (world_models.models.controller.Controller attribute) (world_models.RSSMPolicy attribute) (world_models.vision.VAE.ConvVAE.ConvVAEDecoder attribute) (world_models.vision.VAE.ConvVAE.ConvVAEEncoder attribute) LatentActionModel (class in torchwm) (class in world_models) (class in world_models.models.latent_action_model) LatentActionModelConfig (class in torchwm) (class in world_models) (class in world_models.configs) (class in world_models.configs.genie_config) LatentSequenceDataset (class in world_models.datasets.wm_dataset) learning_rate (torchwm.DiamondConfig attribute) (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.configs.wm_config.WMMDNRNNConfig attribute) (world_models.configs.wm_config.WMVAEConfig attribute) (world_models.DiamondConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute) lineplot() (in module world_models.utils.utils) list_available_atari_envs() (in module world_models.envs) (in module world_models.envs.ale_atari_env) list_available_bsuite_ids() (in module world_models.envs) list_available_datasets() (world_models.datasets.tinyworlds.TinyWorldsDataLoader static method) list_env_backends() (in module torchwm) (in module world_models) (in module world_models.api) list_envs() (in module torchwm) (in module world_models) (in module world_models.api) list_gymnasium_robotics_envs() (in module world_models.envs) (in module world_models.envs.robotics_env) list_models() (in module torchwm) (in module world_models) (in module world_models.api) list_procgen_envs() (in module world_models.envs) (in module world_models.envs.procgen_env) load() (torchwm.IRISAgent method) (world_models.IRISAgent method) (world_models.models.iris_agent.IRISAgent method) load_checkpoint() (in module world_models.helpers.jepa_helper) (world_models.benchmarks.adapters.BaseAdapter method) (world_models.benchmarks.adapters.DiamondAdapter method) (world_models.benchmarks.adapters.DreamerAdapter method) (world_models.benchmarks.adapters.IRISAdapter method) (world_models.training.train_diamond.DiamondAgent method) (world_models.training.train_genie.GenieTrainer method)

load_memory() (in module world_models.utils.utils) load_next_buffer() (world_models.datasets.wm_dataset.RolloutDataset method) load_parameters() (in module world_models.training.train_controller) load_state_dict() (world_models.datasets.diamond_dataset.ReplayBuffer method) (world_models.utils.train_utils.EarlyStopping method) (world_models.utils.train_utils.ReduceLROnPlateau method) load_yml_config() (in module world_models.utils.utils) log() (world_models.utils.jepa_utils.CSVLogger method) (world_models.utils.logging_utils.MetricsLogger method) log_abs_det_jacobian() (torchwm.TanhBijector method) (world_models.TanhBijector method) (world_models.vision.dreamer_decoder.TanhBijector method) log_dir (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) log_file (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) log_interval (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) log_level (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) log_scalar() (world_models.utils.dreamer_utils.Logger method) log_scalars() (world_models.utils.dreamer_utils.Logger method) log_system_stats_freq (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) log_video() (world_models.utils.logging_utils.MetricsLogger method) log_video_freq (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) log_videos() (world_models.utils.dreamer_utils.Logger method) logdir (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.configs.wm_config.WMControllerConfig attribute) (world_models.configs.wm_config.WMMDNRNNConfig attribute) (world_models.configs.wm_config.WMVAEConfig attribute) (world_models.DreamerConfig attribute) Logger (class in world_models.utils.dreamer_utils) LR (torchwm.DiTConfig attribute) (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) (world_models.DiTConfig attribute) lr (world_models.utils.train_utils.ReduceLROnPlateau attribute) (world_models.utils.train_utils.ReduceLROnPlateau property) LSTMBackbone (class in world_models.models.modular_rssm)

M

main() (in module world_models.training.train_diamond) (in module world_models.training.train_genie) (in module world_models.training.train_iris) (in module world_models.training.train_jepa) (in module world_models.training.train_planet) (in module world_models.training.train_rssm) (in module world_models.training.train_world_model) main_from_cli() (in module world_models.training.train_jepa) make_atari_env() (in module world_models.envs) (in module world_models.envs.ale_atari_env) make_atari_vector_env() (in module world_models.envs) (in module world_models.envs.ale_atari_vector_env) make_brax_env() (in module world_models.envs) make_bsuite_env() (in module world_models.envs) make_cifar10() (in module world_models.datasets.cifar10) make_diamond_atari_env() (in module world_models.envs.diamond_atari) make_dmlab_env() (in module world_models.envs) (in module world_models.envs.dmlab) make_env() (in module torchwm) (in module world_models) (in module world_models.api) (in module world_models.envs) (in module world_models.models.dreamer) make_gym_env() (in module world_models.envs) (in module world_models.envs.gym_env) make_gymnasium_env_with_robotics_fallback() (in module world_models.envs.robotics_env) make_imagefolder() (in module world_models.datasets.imagenet1k) make_imagenet1k() (in module world_models.datasets.imagenet1k) make_mujoco_env() (in module world_models.envs) (in module world_models.envs.mujoco_env) make_mujoco_env_from_config() (in module world_models.envs) (in module world_models.envs.mujoco_env) make_nuplan_dataloader() (in module world_models.datasets.nuplan) make_procgen_env() (in module world_models.envs) (in module world_models.envs.procgen_env) make_robotics_env() (in module world_models.envs) (in module world_models.envs.robotics_env) make_transforms() (in module world_models.transforms.image) make_unity_mlagents_env() (in module world_models.envs) (in module world_models.envs.unity_env) make_world_model_env() (in module world_models.envs) map_raster (world_models.datasets.nuplan.NuPlanSample attribute) mask_prob_max (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute) mask_prob_min (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute) MaskCollator (class in world_models.masks.multiblock) (class in world_models.masks.random) maskgit_steps (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute) MaskGITSampler (class in world_models.models.dynamics_model) max_episode_length (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) max_episode_steps (world_models.controller.rollout_generator.RolloutGenerator attribute) (world_models.envs.BraxImageEnv property) (world_models.envs.BSuiteImageEnv property) (world_models.envs.dmlab.DMLabEnv property) (world_models.envs.DMLabEnv property) (world_models.envs.gym_env.GymImageEnv attribute) (world_models.envs.gym_env.GymImageEnv property) (world_models.envs.GymImageEnv attribute) (world_models.envs.GymImageEnv property) (world_models.envs.procgen_env.ProcgenImageEnv property) (world_models.envs.ProcgenImageEnv property) (world_models.envs.unity_env.UnityMLAgentsEnv attribute) (world_models.envs.unity_env.UnityMLAgentsEnv property) (world_models.envs.UnityMLAgentsEnv attribute) (world_models.envs.UnityMLAgentsEnv property) (world_models.utils.utils.TorchImageEnvWrapper property) max_noop (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) max_steps (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (torchwm.IRISOnPolicyBuffer attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.IRISOnPolicyBuffer attribute) (world_models.memory.iris_memory.IRISOnPolicyBuffer attribute) (world_models.training.train_genie.GenieConfig attribute) max_videos_to_save (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) max_workers (world_models.configs.wm_config.WMControllerConfig attribute) MDRNN (class in world_models.models.mdrnn) MDRNNCell (class in world_models.models.mdrnn) mean() (torchwm.SampleDist method) (world_models.SampleDist method) (world_models.vision.dreamer_decoder.SampleDist method) Memory (class in torchwm) (class in world_models) (class in world_models.memory.dreamer_memory) (class in world_models.memory.planet_memory) metadata (world_models.envs.WorldModelEnv attribute) MetricsLogger (class in world_models.utils.logging_utils) MLP (class in world_models.models.vit) mlp_ratio (torchwm.DynamicsModelConfig attribute) (torchwm.STTransformerConfig attribute) (world_models.configs.DynamicsModelConfig attribute) (world_models.configs.genie_config.DynamicsModelConfig attribute) (world_models.configs.genie_config.STTransformerConfig attribute) (world_models.configs.STTransformerConfig attribute) (world_models.DynamicsModelConfig attribute) (world_models.STTransformerConfig attribute) MLPDecoder (class in world_models.models.modular_rssm) MLPEncoder (class in world_models.models.modular_rssm) mode() (torchwm.SampleDist method) (world_models.SampleDist method) (world_models.vision.dreamer_decoder.SampleDist method) model_learning_rate (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) ModelPreset (class in world_models.configs.diamond_config) ModelSpec (class in torchwm) (class in world_models) (class in world_models.api)

ModularRSSM (class in torchwm) (class in world_models) (class in world_models.models.modular_rssm) module torchwm world_models world_models.api world_models.benchmarks world_models.benchmarks.adapters world_models.benchmarks.metrics world_models.benchmarks.reporting world_models.benchmarks.runner world_models.blocks world_models.blocks.mhsa world_models.blocks.st_transformer world_models.catalog world_models.configs world_models.configs.diamond_config world_models.configs.dit_config world_models.configs.dreamer_config world_models.configs.genie_config world_models.configs.iris_config world_models.configs.jepa_config world_models.configs.wm_config world_models.controller.iris_policy world_models.controller.rollout_generator world_models.controller.rssm_policy world_models.datasets.cifar10 world_models.datasets.diamond_dataset world_models.datasets.imagenet1k world_models.datasets.nuplan world_models.datasets.tinyworlds world_models.datasets.video_datasets world_models.datasets.wm_dataset world_models.envs world_models.envs.ale_atari_env world_models.envs.ale_atari_vector_env world_models.envs.diamond_atari world_models.envs.dmc world_models.envs.dmlab world_models.envs.gym_env world_models.envs.mujoco_env world_models.envs.procgen_env world_models.envs.robotics_env world_models.envs.unity_env world_models.envs.vector_env world_models.envs.wrappers world_models.export world_models.helpers.jepa_helper world_models.inference.operators.base world_models.inference.operators.dreamer_operator world_models.inference.operators.iris_operator world_models.inference.operators.jepa_operator world_models.inference.operators.planet_operator world_models.losses.convae_loss world_models.losses.gmm_loss world_models.masks world_models.masks.default 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.controller world_models.models.diffusion world_models.models.diffusion.actor_critic world_models.models.diffusion.DDPM world_models.models.diffusion.diamond_diffusion world_models.models.diffusion.DiT world_models.models.diffusion.reward_termination world_models.models.dreamer world_models.models.dreamer_rssm world_models.models.dynamics_model world_models.models.genie world_models.models.iris_agent world_models.models.iris_transformer world_models.models.jepa_agent world_models.models.latent_action_model world_models.models.mdrnn world_models.models.modular_rssm world_models.models.planet world_models.models.rssm world_models.models.vit world_models.training.rl_harness world_models.training.train_controller world_models.training.train_convvae world_models.training.train_diamond world_models.training.train_genie world_models.training.train_iris world_models.training.train_jepa world_models.training.train_mdn_rnn world_models.training.train_planet world_models.training.train_rssm world_models.training.train_world_model world_models.transforms.image world_models.utils.data_utils world_models.utils.dreamer_utils world_models.utils.jepa_utils world_models.utils.jit_utils world_models.utils.logging_utils world_models.utils.memory_utils world_models.utils.train_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.planet_decoder world_models.vision.planet_encoder world_models.vision.VAE.ConvVAE world_models.vision.video_tokenizer world_models.vision.vq_layer mujoco_binary_path (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) mujoco_camera (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) mujoco_frame_skip (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) mujoco_reset_noise_scale (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) mujoco_xml_path (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) mujoco_xml_string (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) MuJoCoImageEnv (class in world_models.envs) (class in world_models.envs.mujoco_env) MultiAgentBenchmarkRunner (class in world_models.benchmarks.runner) MultiblockMaskCollator (in module world_models.masks) MultiHeadAttention (in module torchwm) (in module world_models) MultiHeadSelfAttention (class in torchwm) (class in world_models) (class in world_models.blocks.mhsa)

N

N (torchwm.RSSMPolicy attribute) (world_models.controller.rssm_policy.RSSMPolicy attribute) (world_models.RSSMPolicy attribute) n_samples (world_models.configs.wm_config.WMControllerConfig attribute) name (torchwm.EnvBackendSpec attribute) (torchwm.ModelSpec attribute) (torchwm.SampleDist property) (world_models.api.EnvBackendSpec attribute) (world_models.api.ModelSpec attribute) (world_models.controller.rollout_generator.RolloutGenerator attribute) (world_models.EnvBackendSpec attribute) (world_models.ModelSpec attribute) (world_models.SampleDist property) (world_models.vision.dreamer_decoder.SampleDist property) no_gpu (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) noreload (world_models.configs.wm_config.WMMDNRNNConfig attribute) (world_models.configs.wm_config.WMVAEConfig attribute) normalize_frames_for_saving() (in module world_models.utils.utils) normalize_procgen_env_name() (in module world_models.envs) (in module world_models.envs.procgen_env) NormalizeActions (class in world_models.envs) (class in world_models.envs.wrappers) nosamples (world_models.configs.wm_config.WMVAEConfig attribute) num_buckets (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) num_conditioning_frames (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) num_epochs (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.configs.wm_config.WMMDNRNNConfig attribute) (world_models.configs.wm_config.WMVAEConfig attribute) (world_models.DiamondConfig attribute) num_frames (torchwm.DynamicsModelConfig attribute) (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (torchwm.LatentActionModelConfig attribute) (torchwm.STTransformerConfig attribute) (torchwm.VideoTokenizerConfig attribute) (world_models.configs.DynamicsModelConfig attribute) (world_models.configs.genie_config.DynamicsModelConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.genie_config.LatentActionModelConfig attribute) (world_models.configs.genie_config.STTransformerConfig attribute) (world_models.configs.genie_config.VideoTokenizerConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.configs.LatentActionModelConfig attribute) (world_models.configs.STTransformerConfig attribute) (world_models.configs.VideoTokenizerConfig attribute) (world_models.datasets.tinyworlds.TinyWorldsConfig attribute) (world_models.datasets.video_datasets.DatasetConfig attribute) (world_models.DynamicsModelConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.LatentActionModelConfig attribute) (world_models.STTransformerConfig attribute) (world_models.training.train_genie.GenieConfig attribute) (world_models.VideoTokenizerConfig attribute)

num_heads (torchwm.DynamicsModelConfig attribute) (torchwm.LatentActionModelConfig attribute) (torchwm.STTransformerConfig attribute) (torchwm.VideoTokenizerConfig attribute) (world_models.configs.DynamicsModelConfig attribute) (world_models.configs.genie_config.DynamicsModelConfig attribute) (world_models.configs.genie_config.LatentActionModelConfig attribute) (world_models.configs.genie_config.STTransformerConfig attribute) (world_models.configs.genie_config.VideoTokenizerConfig attribute) (world_models.configs.LatentActionModelConfig attribute) (world_models.configs.STTransformerConfig attribute) (world_models.configs.VideoTokenizerConfig attribute) (world_models.DynamicsModelConfig attribute) (world_models.LatentActionModelConfig attribute) (world_models.STTransformerConfig attribute) (world_models.VideoTokenizerConfig attribute) num_layers (torchwm.IRISActor attribute) (torchwm.IRISCritic attribute) (torchwm.IRISPolicy attribute) (world_models.controller.iris_policy.IRISActor attribute) (world_models.controller.iris_policy.IRISCritic attribute) (world_models.controller.iris_policy.IRISPolicy attribute) (world_models.IRISActor attribute) (world_models.IRISCritic attribute) (world_models.IRISPolicy attribute) num_patches_per_frame (torchwm.STTransformerConfig attribute) (world_models.configs.genie_config.STTransformerConfig attribute) (world_models.configs.STTransformerConfig attribute) (world_models.STTransformerConfig attribute) num_sampling_steps (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) num_test_files (world_models.datasets.wm_dataset.RolloutDataset attribute) num_units (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) num_workers (world_models.datasets.tinyworlds.TinyWorldsConfig attribute) (world_models.datasets.video_datasets.DatasetConfig attribute) NumPyDataset (class in world_models.datasets.video_datasets) NuPlanDataset (class in world_models.datasets.nuplan) NuPlanSample (class in world_models.datasets.nuplan)

O

obs_embed_size (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) obs_key (world_models.datasets.video_datasets.VideoDatasetConfig attribute) obs_shape (torchwm.IRISReplayBuffer attribute) (world_models.IRISReplayBuffer attribute) (world_models.memory.iris_memory.IRISReplayBuffer attribute) obs_size (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) obs_space (world_models.envs.RenderImage property) (world_models.envs.ResizeImage property) (world_models.envs.wrappers.RenderImage property) (world_models.envs.wrappers.ResizeImage property) ObsDict (class in world_models.envs) (class in world_models.envs.wrappers) observation_size (world_models.utils.utils.TorchImageEnvWrapper property) observation_space (world_models.envs.BraxImageEnv property) (world_models.envs.BSuiteImageEnv property) (world_models.envs.DeepMindControlEnv attribute) (world_models.envs.DeepMindControlEnv property) (world_models.envs.dmc.DeepMindControlEnv attribute) (world_models.envs.dmc.DeepMindControlEnv property) (world_models.envs.dmlab.DMLabEnv property) (world_models.envs.DMLabEnv property) (world_models.envs.gym_env.GymImageEnv attribute) (world_models.envs.gym_env.GymImageEnv property) (world_models.envs.GymImageEnv attribute) (world_models.envs.GymImageEnv property) (world_models.envs.mujoco_env.MuJoCoImageEnv property) (world_models.envs.MuJoCoImageEnv property) (world_models.envs.ObsDict property) (world_models.envs.procgen_env.ProcgenImageEnv property) (world_models.envs.ProcgenImageEnv property) (world_models.envs.RewardObs property) (world_models.envs.unity_env.UnityMLAgentsEnv attribute) (world_models.envs.unity_env.UnityMLAgentsEnv property) (world_models.envs.UnityMLAgentsEnv attribute) (world_models.envs.UnityMLAgentsEnv property) (world_models.envs.wrappers.ObsDict property) (world_models.envs.wrappers.RewardObs property)

ObservationDataset (class in world_models.datasets.wm_dataset) observations (torchwm.IRISOnPolicyBuffer attribute) (world_models.IRISOnPolicyBuffer attribute) (world_models.memory.iris_memory.IRISOnPolicyBuffer attribute) observe_rollout() (torchwm.ModularRSSM method) (torchwm.RSSM method) (world_models.models.dreamer_rssm.RSSM method) (world_models.models.modular_rssm.ModularRSSM method) (world_models.ModularRSSM method) (world_models.RSSM method) observe_step() (torchwm.ModularRSSM method) (torchwm.RSSM method) (world_models.models.dreamer_rssm.RSSM method) (world_models.models.modular_rssm.ModularRSSM method) (world_models.ModularRSSM method) (world_models.RSSM method) OneHotAction (class in world_models.envs) (class in world_models.envs.wrappers) operator_action_dim (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) operator_state_dim (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) OperatorABC (class in torchwm) (class in world_models) (class in world_models.inference.operators.base) optimize_memory_efficient_ops() (in module world_models.utils.memory_utils) output_size (torchwm.CNNFeatureExtractor attribute) (world_models.CNNFeatureExtractor attribute) (world_models.controller.iris_policy.CNNFeatureExtractor attribute) output_specs (torchwm.OperatorABC attribute) (world_models.inference.operators.base.OperatorABC attribute) (world_models.OperatorABC attribute)

P

p_mean (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) p_sample() (torchwm.DDPM method) (world_models.DDPM method) (world_models.models.diffusion.DDPM.DDPM method) p_std (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) parameter_count() (torchwm.DiT method) (torchwm.Genie method) (torchwm.IRISAgent method) (world_models.DiT method) (world_models.Genie method) (world_models.IRISAgent method) (world_models.models.diffusion.DiT.DiT method) (world_models.models.dreamer.Dreamer method) (world_models.models.dreamer.DreamerAgent method) (world_models.models.genie.Genie method) (world_models.models.iris_agent.IRISAgent method) (world_models.models.jepa_agent.JEPAAgent method) (world_models.training.train_diamond.DiamondAgent method) PATCH (torchwm.DiTConfig attribute) (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) (world_models.DiTConfig attribute) patch_size (torchwm.DynamicsModelConfig attribute) (torchwm.LatentActionModelConfig attribute) (torchwm.VideoTokenizerConfig attribute) (world_models.configs.DynamicsModelConfig attribute) (world_models.configs.genie_config.DynamicsModelConfig attribute) (world_models.configs.genie_config.LatentActionModelConfig attribute) (world_models.configs.genie_config.VideoTokenizerConfig attribute) (world_models.configs.LatentActionModelConfig attribute) (world_models.configs.VideoTokenizerConfig attribute) (world_models.DynamicsModelConfig attribute) (world_models.LatentActionModelConfig attribute) (world_models.VideoTokenizerConfig attribute) PatchEmbed (class in torchwm) (class in world_models) (class in world_models.models.diffusion.DiT) (class in world_models.models.vit) PatchUnEmbed (class in torchwm) (class in world_models) (class in world_models.models.diffusion.DiT) persistent_workers (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) pin_memory (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.datasets.video_datasets.DatasetConfig attribute) (world_models.DiamondConfig attribute) Planet (class in world_models.models.planet) PlaNetOperator (class in torchwm) (class in world_models) (class in world_models.inference.operators.planet_operator) planning_target (world_models.datasets.nuplan.NuPlanSample attribute) play() (torchwm.Genie method) (world_models.Genie method) (world_models.models.genie.Genie method)

plot_metrics() (in module world_models.utils.utils) policy (world_models.controller.rollout_generator.RolloutGenerator attribute) policy_loss() (world_models.models.diffusion.actor_critic.RLLoss method) poll() (torchwm.RSSMPolicy method) (world_models.controller.rssm_policy.RSSMPolicy method) (world_models.RSSMPolicy method) pop_size (world_models.configs.wm_config.WMControllerConfig attribute) postprocess() (torchwm.OperatorABC method) (world_models.inference.operators.base.OperatorABC method) (world_models.OperatorABC method) postprocess_img() (in module world_models.utils.utils) PPOTrainer (class in world_models.training.rl_harness) precompute_latents() (in module world_models.training.train_mdn_rnn) pred_reward() (torchwm.RecurrentStateSpaceModel method) (world_models.models.rssm.RecurrentStateSpaceModel method) (world_models.RecurrentStateSpaceModel method) predict() (torchwm.RewardTerminationModel method) (world_models.models.diffusion.reward_termination.RewardTerminationModel method) (world_models.RewardTerminationModel method) predict_next_tokens() (world_models.models.iris_transformer.IRISTransformer method) prefetch_iterator() (in module world_models.utils.data_utils) preprocess() (torchwm.DreamerOperator method) (torchwm.IrisOperator method) (torchwm.JEPAOperator method) (torchwm.OperatorABC method) (torchwm.PlaNetOperator method) (world_models.DreamerOperator method) (world_models.inference.operators.base.OperatorABC method) (world_models.inference.operators.dreamer_operator.DreamerOperator method) (world_models.inference.operators.iris_operator.IrisOperator method) (world_models.inference.operators.jepa_operator.JEPAOperator method) (world_models.inference.operators.planet_operator.PlaNetOperator method) (world_models.IrisOperator method) (world_models.JEPAOperator method) (world_models.OperatorABC method) (world_models.PlaNetOperator 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 (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) process() (torchwm.OperatorABC method) (world_models.inference.operators.base.OperatorABC method) (world_models.OperatorABC method) procgen_distribution_mode (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) procgen_num_levels (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) procgen_start_level (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) ProcgenImageEnv (class in world_models.envs) (class in world_models.envs.procgen_env)

Q

q_sample() (torchwm.DDPM method) (world_models.DDPM method) (world_models.models.diffusion.DDPM.DDPM method) qkv_bias (torchwm.DynamicsModelConfig attribute) (torchwm.STTransformerConfig attribute) (world_models.configs.DynamicsModelConfig attribute) (world_models.configs.genie_config.DynamicsModelConfig attribute) (world_models.configs.genie_config.STTransformerConfig attribute) (world_models.configs.STTransformerConfig attribute) (world_models.DynamicsModelConfig attribute) (world_models.STTransformerConfig attribute)

R

r (torchwm.Episode attribute) (world_models.Episode attribute) (world_models.memory.planet_memory.Episode attribute) RandomMaskCollator (in module world_models.masks) RecurrentStateSpaceModel (class in torchwm) (class in world_models) (class in world_models.models.rssm) recursive (world_models.datasets.video_datasets.VideoDatasetConfig attribute) ReduceLROnPlateau (class in world_models.utils.train_utils) register_buffers() (world_models.utils.dreamer_utils.TwoHotEncoder method) register_gymnasium_robotics_envs() (in module world_models.envs) (in module world_models.envs.robotics_env) render (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) render() (world_models.envs.BraxImageEnv method) (world_models.envs.BSuiteImageEnv method) (world_models.envs.DeepMindControlEnv method) (world_models.envs.dmc.DeepMindControlEnv method) (world_models.envs.dmlab.DMLabEnv method) (world_models.envs.DMLabEnv method) (world_models.envs.gym_env.GymImageEnv method) (world_models.envs.GymImageEnv method) (world_models.envs.mujoco_env.MuJoCoImageEnv method) (world_models.envs.MuJoCoImageEnv method) (world_models.envs.procgen_env.ProcgenImageEnv method) (world_models.envs.ProcgenImageEnv method) (world_models.envs.unity_env.UnityMLAgentsEnv method) (world_models.envs.UnityMLAgentsEnv method) (world_models.envs.WorldModelEnv method) (world_models.utils.utils.TorchImageEnvWrapper method) render_batch() (world_models.envs.vector_env.VectorizedEnv method) RenderImage (class in world_models.envs) (class in world_models.envs.wrappers) repeat_interleave_batch() (in module world_models.utils.jepa_utils) ReplayBuffer (class in torchwm) (class in world_models) (class in world_models.datasets.diamond_dataset) (class in world_models.memory.dreamer_memory) required (torchwm.TensorSpec attribute) (world_models.inference.operators.base.TensorSpec attribute) (world_models.TensorSpec attribute) ResBlock (class in world_models.models.diffusion.diamond_diffusion) reset() (torchwm.RSSMPolicy method) (world_models.controller.rssm_policy.RSSMPolicy method) (world_models.envs.BraxImageEnv method) (world_models.envs.BSuiteImageEnv method) (world_models.envs.DeepMindControlEnv method) (world_models.envs.diamond_atari.DiamondAtariWrapper method) (world_models.envs.dmc.DeepMindControlEnv method) (world_models.envs.dmlab.DMLabEnv method) (world_models.envs.DMLabEnv method) (world_models.envs.gym_env.GymImageEnv method) (world_models.envs.GymImageEnv method) (world_models.envs.mujoco_env.MuJoCoImageEnv method) (world_models.envs.MuJoCoImageEnv method) (world_models.envs.ObsDict method) (world_models.envs.OneHotAction method) (world_models.envs.procgen_env.ProcgenImageEnv method) (world_models.envs.ProcgenImageEnv method) (world_models.envs.RenderImage method) (world_models.envs.ResizeImage method) (world_models.envs.RewardObs method) (world_models.envs.TimeLimit method) (world_models.envs.unity_env.UnityMLAgentsEnv method) (world_models.envs.UnityMLAgentsEnv method) (world_models.envs.WorldModelEnv method) (world_models.envs.wrappers.ObsDict method) (world_models.envs.wrappers.OneHotAction method) (world_models.envs.wrappers.RenderImage method) (world_models.envs.wrappers.ResizeImage method) (world_models.envs.wrappers.RewardObs method) (world_models.envs.wrappers.TimeLimit method) (world_models.envs.wrappers.UUID method) (world_models.RSSMPolicy method) (world_models.utils.jepa_utils.AverageMeter method) (world_models.utils.utils.TorchImageEnvWrapper method)

reset_batch() (world_models.envs.vector_env.TorchVectorizedEnv method) (world_models.envs.vector_env.VectorizedEnv method) ResidualBlock (class in world_models.vision.iris_decoder) (class in world_models.vision.iris_encoder) ResizeImage (class in world_models.envs) (class in world_models.envs.wrappers) restore (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) restore_checkpoint() (world_models.models.dreamer.Dreamer method) reward_channels (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.diamond_config.ModelPreset attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) reward_clip (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) reward_cond_dim (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) reward_lstm_dim (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.diamond_config.ModelPreset attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) reward_res_blocks (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) RewardModel (class in torchwm) (class in world_models) RewardObs (class in world_models.envs) (class in world_models.envs.wrappers) rewards (torchwm.IRISOnPolicyBuffer attribute) (world_models.IRISOnPolicyBuffer attribute) (world_models.memory.iris_memory.IRISOnPolicyBuffer attribute) RewardTerminationLoss (class in world_models.models.diffusion.reward_termination) RewardTerminationModel (class in torchwm) (class in world_models) (class in world_models.models.diffusion.reward_termination) rho (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) RLEnvironmentDataset (class in world_models.datasets.video_datasets) RLLoss (class in world_models.models.diffusion.actor_critic) RMSNorm (class in torchwm) (class in world_models) rollout_eval() (world_models.controller.rollout_generator.RolloutGenerator method) rollout_eval_n() (world_models.controller.rollout_generator.RolloutGenerator method) rollout_n() (world_models.controller.rollout_generator.RolloutGenerator method) rollout_once() (world_models.controller.rollout_generator.RolloutGenerator method) rollout_prior() (torchwm.RecurrentStateSpaceModel method) (world_models.models.rssm.RecurrentStateSpaceModel method) (world_models.RecurrentStateSpaceModel method) RolloutDataset (class in world_models.datasets.wm_dataset) RolloutGenerator (class in world_models.controller.rollout_generator) root (world_models.datasets.wm_dataset.RolloutDataset attribute) ROOT_PATH (torchwm.DiTConfig attribute) (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) (world_models.DiTConfig attribute) RSSM (class in torchwm) (class in world_models) (class in world_models.models.dreamer_rssm) rssm (torchwm.RSSMPolicy attribute) (world_models.controller.rssm_policy.RSSMPolicy attribute) (world_models.RSSMPolicy attribute) RSSMPolicy (class in torchwm) (class in world_models) (class in world_models.controller.rssm_policy) run() (world_models.benchmarks.runner.BenchmarkRunner method) (world_models.envs.vector_env.SimWorker method) run_all() (world_models.benchmarks.runner.MultiAgentBenchmarkRunner method) run_training_pipeline() (in module world_models.training.train_world_model)

S

sample() (torchwm.DDPM method) (torchwm.DynamicsModel method) (torchwm.Memory method) (torchwm.ReplayBuffer method) (torchwm.SampleDist method) (world_models.datasets.diamond_dataset.ReplayBuffer method) (world_models.DDPM method) (world_models.DynamicsModel method) (world_models.Memory method) (world_models.memory.dreamer_memory.Memory method) (world_models.memory.dreamer_memory.ReplayBuffer method) (world_models.memory.planet_memory.Memory method) (world_models.models.diffusion.DDPM.DDPM method) (world_models.models.diffusion.diamond_diffusion.EulerSampler method) (world_models.models.dynamics_model.DynamicsModel method) (world_models.models.dynamics_model.MaskGITSampler method) (world_models.ReplayBuffer method) (world_models.SampleDist method) (world_models.vision.dreamer_decoder.SampleDist method) sample_interval (world_models.configs.wm_config.WMVAEConfig attribute) sample_next_tokens() (world_models.models.iris_transformer.IRISTransformer method) sample_noise_level() (world_models.models.diffusion.diamond_diffusion.EDMPreconditioner method) sample_random_action() (world_models.utils.utils.TorchImageEnvWrapper method) sample_sequence() (torchwm.IRISReplayBuffer method) (world_models.datasets.diamond_dataset.ReplayBuffer method) (world_models.IRISReplayBuffer method) (world_models.memory.iris_memory.IRISReplayBuffer method) sample_single() (torchwm.IRISReplayBuffer method) (world_models.IRISReplayBuffer method) (world_models.memory.iris_memory.IRISReplayBuffer method) sample_temperature (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute) SampleDist (class in torchwm) (class in world_models) (class in world_models.vision.dreamer_decoder) sampling_method (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) save() (torchwm.IRISAgent method) (world_models.IRISAgent method) (world_models.models.dreamer.Dreamer method) (world_models.models.iris_agent.IRISAgent method) save_checkpoint() (in module world_models.training.train_convvae) (in module world_models.training.train_mdn_rnn) (world_models.training.train_diamond.DiamondAgent method) (world_models.training.train_genie.GenieTrainer method) save_frames() (in module world_models.utils.utils) save_interval (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) save_pretrained() (torchwm.DiT method) (torchwm.Genie method) (world_models.DiT method) (world_models.Genie method) (world_models.models.diffusion.DiT.DiT method) (world_models.models.genie.Genie method) save_video() (in module world_models.utils.utils) scalar_freq (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) scenario_name (world_models.datasets.nuplan.NuPlanSample attribute) scheduler_factor (world_models.configs.wm_config.WMMDNRNNConfig attribute) (world_models.configs.wm_config.WMVAEConfig attribute) scheduler_patience (world_models.configs.wm_config.WMMDNRNNConfig attribute) (world_models.configs.wm_config.WMVAEConfig attribute) seed (torchwm.DiamondConfig attribute) (torchwm.DreamerConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DiamondConfig attribute) (world_models.DreamerConfig attribute) seed_steps (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) SelectAction (class in world_models.envs) (class in world_models.envs.wrappers) SelfAttentionBlock (class in world_models.vision.iris_encoder) seq_len (torchwm.IRISReplayBuffer attribute) (world_models.configs.wm_config.WMMDNRNNConfig attribute) (world_models.datasets.wm_dataset.SequenceDataset attribute) (world_models.IRISReplayBuffer attribute) (world_models.memory.iris_memory.IRISReplayBuffer attribute) seq_to_batch() (torchwm.ModularRSSM method) (torchwm.RSSM method) (world_models.models.dreamer_rssm.RSSM method) (world_models.models.modular_rssm.ModularRSSM method) (world_models.ModularRSSM method) (world_models.RSSM method) SequenceDataset (class in world_models.datasets.diamond_dataset) (class in world_models.datasets.wm_dataset) setup_logging() (in module world_models.utils.logging_utils) shape (torchwm.TensorSpec attribute) (world_models.inference.operators.base.TensorSpec attribute) (world_models.TensorSpec attribute) shuffle (world_models.datasets.video_datasets.DatasetConfig attribute) sigma_data (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) sigma_max (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) sigma_min (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) sign (torchwm.TanhBijector property) (world_models.TanhBijector property) (world_models.vision.dreamer_decoder.TanhBijector property)

SimWorker (class in world_models.envs.vector_env) sinusoidal_time_embedding() (in module torchwm) (in module world_models) (in module world_models.models.diffusion.DiT) size (torchwm.Episode property) (torchwm.IRISReplayBuffer attribute) (torchwm.Memory attribute) (torchwm.Memory property) (world_models.Episode property) (world_models.IRISReplayBuffer attribute) (world_models.Memory attribute) (world_models.Memory property) (world_models.memory.iris_memory.IRISReplayBuffer attribute) (world_models.memory.planet_memory.Episode property) (world_models.memory.planet_memory.Memory attribute) (world_models.memory.planet_memory.Memory property) slave_routine() (in module world_models.training.train_controller) split (world_models.datasets.tinyworlds.TinyWorldsConfig attribute) state (world_models.envs.WorldModelEnv property) state_dict() (world_models.datasets.diamond_dataset.ReplayBuffer method) (world_models.utils.train_utils.EarlyStopping method) (world_models.utils.train_utils.ReduceLROnPlateau method) state_posterior() (torchwm.RecurrentStateSpaceModel method) (world_models.models.rssm.RecurrentStateSpaceModel method) (world_models.RecurrentStateSpaceModel method) state_prior() (torchwm.RecurrentStateSpaceModel method) (world_models.models.rssm.RecurrentStateSpaceModel method) (world_models.RecurrentStateSpaceModel method) state_size (torchwm.RSSMPolicy attribute) (world_models.controller.rssm_policy.RSSMPolicy attribute) (world_models.RSSMPolicy attribute) step() (world_models.envs.ActionRepeat method) (world_models.envs.BraxImageEnv method) (world_models.envs.BSuiteImageEnv method) (world_models.envs.DeepMindControlEnv method) (world_models.envs.diamond_atari.DiamondAtariWrapper method) (world_models.envs.dmc.DeepMindControlEnv method) (world_models.envs.dmlab.DMLabEnv method) (world_models.envs.DMLabEnv method) (world_models.envs.gym_env.GymImageEnv method) (world_models.envs.GymImageEnv method) (world_models.envs.mujoco_env.MuJoCoImageEnv method) (world_models.envs.MuJoCoImageEnv method) (world_models.envs.NormalizeActions method) (world_models.envs.ObsDict method) (world_models.envs.OneHotAction method) (world_models.envs.procgen_env.ProcgenImageEnv method) (world_models.envs.ProcgenImageEnv 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.unity_env.UnityMLAgentsEnv method) (world_models.envs.UnityMLAgentsEnv method) (world_models.envs.WorldModelEnv method) (world_models.envs.wrappers.ActionRepeat method) (world_models.envs.wrappers.NormalizeActions method) (world_models.envs.wrappers.ObsDict method) (world_models.envs.wrappers.OneHotAction method) (world_models.envs.wrappers.RenderImage method) (world_models.envs.wrappers.ResizeImage method) (world_models.envs.wrappers.RewardObs method) (world_models.envs.wrappers.SelectAction method) (world_models.envs.wrappers.TimeLimit 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.train_utils.EarlyStopping method) (world_models.utils.train_utils.ReduceLROnPlateau method) (world_models.utils.utils.TorchImageEnvWrapper method) step_batch() (world_models.envs.vector_env.TorchVectorizedEnv method) (world_models.envs.vector_env.VectorizedEnv method) steps (torchwm.IRISReplayBuffer attribute) (world_models.IRISReplayBuffer attribute) (world_models.memory.iris_memory.IRISReplayBuffer attribute) STMLP (class in world_models.blocks.st_transformer) stoch_size (torchwm.DreamerConfig attribute) (torchwm.ModularRSSM property) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) (world_models.models.modular_rssm.BackboneBase attribute) (world_models.models.modular_rssm.ModularRSSM property) (world_models.ModularRSSM property) stop (world_models.utils.train_utils.EarlyStopping attribute) (world_models.utils.train_utils.EarlyStopping property) StreamingVideoWriter (class in world_models.utils.utils) STSpatialAttention (class in world_models.blocks.st_transformer) STTemporalAttention (class in world_models.blocks.st_transformer) STTransformer (class in torchwm) (class in world_models) (class in world_models.blocks.st_transformer) STTransformerBlock (class in world_models.blocks.st_transformer) STTransformerConfig (class in torchwm) (class in world_models) (class in world_models.configs) (class in world_models.configs.genie_config) summary() (torchwm.DiT method) (torchwm.Genie method) (torchwm.IRISAgent method) (world_models.DiT method) (world_models.Genie method) (world_models.IRISAgent method) (world_models.models.diffusion.DiT.DiT method) (world_models.models.dreamer.Dreamer method) (world_models.models.dreamer.DreamerAgent method) (world_models.models.genie.Genie method) (world_models.models.iris_agent.IRISAgent method) (world_models.models.jepa_agent.JEPAAgent method) (world_models.training.train_diamond.DiamondAgent method) sweep_train() (in module world_models.training.train_jepa) symexp() (in module world_models.utils.dreamer_utils) symlog() (in module world_models.utils.dreamer_utils) symlog_range (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute)

T

t (torchwm.Episode attribute) T (torchwm.RSSMPolicy attribute) (world_models.controller.rssm_policy.RSSMPolicy attribute) t (world_models.Episode attribute) (world_models.memory.planet_memory.Episode attribute) T (world_models.RSSMPolicy attribute) TanhBijector (class in torchwm) (class in world_models) (class in world_models.vision.dreamer_decoder) target_return (world_models.configs.wm_config.WMControllerConfig attribute) td_lambda (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) TensorBoardMetrics (class in world_models.utils.utils) TensorSpec (class in torchwm) (class in world_models) (class in world_models.inference.operators.base) terminals (torchwm.IRISOnPolicyBuffer attribute) (world_models.IRISOnPolicyBuffer attribute) (world_models.memory.iris_memory.IRISOnPolicyBuffer attribute) terminate() (torchwm.Episode method) (world_models.Episode method) (world_models.memory.planet_memory.Episode method) terminate_on_life_loss (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) test (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) test_episodes (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) test_epoch() (in module world_models.training.train_convvae) test_interval (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) test_trained_model() (in module world_models.training.train_world_model) time_limit (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.configs.wm_config.WMControllerConfig attribute) (world_models.DreamerConfig attribute) TimeLimit (class in world_models.envs) (class in world_models.envs.wrappers) TimestepEmbedding (class in world_models.models.diffusion.diamond_diffusion) TIMESTEPS (torchwm.DiTConfig attribute) (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) (world_models.DiTConfig attribute) TinyWorldsConfig (class in world_models.datasets.tinyworlds) TinyWorldsDataLoader (class in world_models.datasets.tinyworlds) TinyWorldsDataset (class in world_models.datasets.tinyworlds) to() (torchwm.OperatorABC method) (world_models.inference.operators.base.OperatorABC method) (world_models.OperatorABC method) (world_models.utils.dreamer_utils.TwoHotEncoder method) to_dict() (torchwm.JEPAConfig method) (world_models.configs.jepa_config.JEPAConfig method) (world_models.configs.JEPAConfig method) (world_models.configs.wm_config.WMControllerConfig method) (world_models.configs.wm_config.WMMDNRNNConfig method) (world_models.JEPAConfig method) to_latent() (in module world_models.training.train_mdn_rnn) to_nested_dict() (torchwm.JEPAConfig method) (world_models.configs.jepa_config.JEPAConfig method) (world_models.configs.JEPAConfig method) (world_models.JEPAConfig method) to_tensor_obs() (in module world_models.utils.utils) to_train_dict() (torchwm.JEPAConfig method) (world_models.configs.jepa_config.JEPAConfig method) (world_models.configs.JEPAConfig method) (world_models.JEPAConfig method) tokenizer_decoder_depth (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute) tokenizer_decoder_dim (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute)

tokenizer_embedding_dim (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute) tokenizer_encoder_depth (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute) tokenizer_encoder_dim (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute) tokenizer_num_heads (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) tokenizer_vocab_size (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute) TorchImageEnvWrapper (class in world_models.utils.utils) TorchVectorizedEnv (class in world_models.envs.vector_env) torchwm module total_steps (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) train (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.datasets.wm_dataset.RolloutDataset attribute) (world_models.DreamerConfig attribute) train() (in module world_models.training.train_planet) (torchwm.DiT class method) (world_models.DiT class method) (world_models.models.diffusion.DiT.DiT class method) (world_models.models.dreamer.DreamerAgent method) (world_models.models.jepa_agent.JEPAAgent method) (world_models.models.planet.Planet method) (world_models.training.rl_harness.PPOTrainer method) (world_models.training.train_diamond.DiamondAgent method) (world_models.training.train_genie.GenieTrainer method) (world_models.training.train_iris.IRISTrainer method) train_batch_size (world_models.configs.wm_config.WMVAEConfig attribute) train_controller() (in module world_models.training.train_controller) train_convae() (in module world_models.training.train_convvae) train_diamond() (in module world_models.training.train_diamond) train_epoch() (in module world_models.training.train_convvae) (world_models.training.train_iris.IRISTrainer method) train_mdn_rnn() (in module world_models.training.train_mdn_rnn) train_one_batch() (world_models.models.dreamer.Dreamer method) train_rssm() (in module world_models.training.train_rssm) train_seq_len (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) train_step() (world_models.training.rl_harness.PPOTrainer method) (world_models.training.train_genie.GenieTrainer method) training_step() (torchwm.Genie method) (world_models.Genie method) (world_models.models.genie.Genie method) training_steps_per_epoch (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) transform (world_models.datasets.wm_dataset.RolloutDataset attribute) TransformerBackbone (class in world_models.models.modular_rssm) TransformerBlock (class in world_models.models.diffusion.DiT) (class in world_models.models.modular_rssm) trunc_normal_() (in module world_models.utils.jepa_utils) TwoHotEncoder (class in world_models.utils.dreamer_utils)

U

u (torchwm.Episode attribute) (world_models.Episode attribute) (world_models.memory.planet_memory.Episode attribute) unity_base_port (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) unity_behavior_name (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) unity_file_name (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) unity_no_graphics (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) unity_quality_level (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) unity_time_scale (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) unity_worker_id (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) UnityMLAgentsEnv (class in world_models.envs) (class in world_models.envs.unity_env)

UpBlock (class in world_models.models.diffusion.diamond_diffusion) update() (world_models.utils.jepa_utils.AverageMeter method) (world_models.utils.utils.TensorBoardMetrics method) update_actor_critic() (torchwm.IRISAgent method) (world_models.IRISAgent method) (world_models.models.iris_agent.IRISAgent method) update_autoencoder() (torchwm.IRISAgent method) (world_models.IRISAgent method) (world_models.models.iris_agent.IRISAgent method) update_steps (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) update_transformer() (torchwm.IRISAgent method) (world_models.IRISAgent method) (world_models.models.iris_agent.IRISAgent method) UpsampleBlock (class in world_models.vision.iris_decoder) use_amp (torchwm.DiamondConfig attribute) (torchwm.DreamerConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DiamondConfig attribute) (world_models.DreamerConfig attribute) use_disc_model (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) use_ema (torchwm.VideoTokenizerConfig attribute) (world_models.configs.genie_config.VideoTokenizerConfig attribute) (world_models.configs.VideoTokenizerConfig attribute) (world_models.VideoTokenizerConfig attribute) UUID (class in world_models.envs.wrappers)

V

validate() (world_models.training.train_genie.GenieTrainer method) validate_mapping() (torchwm.OperatorABC class method) (world_models.inference.operators.base.OperatorABC class method) (world_models.OperatorABC class method) value_learning_rate (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) value_loss() (world_models.models.diffusion.actor_critic.RLLoss method) (world_models.models.dreamer.Dreamer method) ValueModel (class in torchwm) (class in world_models) VectorizedEnv (class in world_models.envs.vector_env) VectorQuantizer (class in torchwm) (class in world_models) (class in world_models.vision.vq_layer) VectorQuantizerEMA (class in torchwm) (class in world_models) (class in world_models.vision.vq_layer) video_format (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) video_fps (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) video_paths (world_models.datasets.video_datasets.VideoDatasetBase attribute) VideoDataset (class in world_models.training.train_genie) VideoDatasetBase (class in world_models.datasets.video_datasets) VideoDatasetConfig (class in world_models.datasets.video_datasets)

VideoFolderDataset (class in world_models.datasets.video_datasets) VideoTokenizer (class in torchwm) (class in world_models) (class in world_models.vision.video_tokenizer) VideoTokenizerConfig (class in torchwm) (class in world_models) (class in world_models.configs) (class in world_models.configs.genie_config) 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) ViTEncoder (class in world_models.models.modular_rssm) vocab_size (torchwm.DynamicsModelConfig attribute) (torchwm.LatentActionModelConfig attribute) (torchwm.VideoTokenizerConfig attribute) (world_models.configs.DynamicsModelConfig attribute) (world_models.configs.genie_config.DynamicsModelConfig attribute) (world_models.configs.genie_config.LatentActionModelConfig attribute) (world_models.configs.genie_config.VideoTokenizerConfig attribute) (world_models.configs.LatentActionModelConfig attribute) (world_models.configs.VideoTokenizerConfig attribute) (world_models.DynamicsModelConfig attribute) (world_models.LatentActionModelConfig attribute) (world_models.VideoTokenizerConfig attribute)

W

wandb_api_key (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) wandb_entity (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) wandb_project (torchwm.DreamerConfig attribute) (world_models.configs.dreamer_config.DreamerConfig attribute) (world_models.configs.DreamerConfig attribute) (world_models.DreamerConfig attribute) warmup() (world_models.models.planet.Planet method) warmup_steps (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute) WarmupCosineSchedule (class in world_models.utils.jepa_utils) weight_decay (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.training.train_genie.GenieConfig attribute) weight_decay_actor (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) weight_decay_diffusion (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) weight_decay_reward (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) WIDTH (torchwm.DiTConfig attribute) (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) width (world_models.configs.wm_config.WMVAEConfig attribute) WIDTH (world_models.DiTConfig attribute) window_attention_heads (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (torchwm.LatentActionModelConfig attribute) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.genie_config.LatentActionModelConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.configs.LatentActionModelConfig attribute) (world_models.GenieConfig attribute) (world_models.GenieSmallConfig attribute) (world_models.LatentActionModelConfig attribute) (world_models.training.train_genie.GenieConfig attribute) WMControllerConfig (class in world_models.configs.wm_config) WMMDNRNNConfig (class in world_models.configs.wm_config) WMVAEConfig (class in world_models.configs.wm_config) WORKDIR (torchwm.DiTConfig attribute) (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) (world_models.DiTConfig attribute) world_model_loss() (world_models.models.dreamer.Dreamer method) world_models module world_models.api module world_models.benchmarks module world_models.benchmarks.adapters module world_models.benchmarks.metrics module world_models.benchmarks.reporting module world_models.benchmarks.runner module world_models.blocks module world_models.blocks.mhsa module world_models.blocks.st_transformer module world_models.catalog module world_models.configs module world_models.configs.diamond_config module world_models.configs.dit_config module world_models.configs.dreamer_config module world_models.configs.genie_config module world_models.configs.iris_config module world_models.configs.jepa_config module world_models.configs.wm_config module world_models.controller.iris_policy module world_models.controller.rollout_generator module world_models.controller.rssm_policy module world_models.datasets.cifar10 module world_models.datasets.diamond_dataset module world_models.datasets.imagenet1k module world_models.datasets.nuplan module world_models.datasets.tinyworlds module world_models.datasets.video_datasets module world_models.datasets.wm_dataset module world_models.envs module world_models.envs.ale_atari_env module world_models.envs.ale_atari_vector_env module world_models.envs.diamond_atari module world_models.envs.dmc module world_models.envs.dmlab module world_models.envs.gym_env module world_models.envs.mujoco_env module

world_models.envs.procgen_env module world_models.envs.robotics_env module world_models.envs.unity_env module world_models.envs.vector_env module world_models.envs.wrappers module world_models.export module world_models.helpers.jepa_helper module world_models.inference.operators.base module world_models.inference.operators.dreamer_operator module world_models.inference.operators.iris_operator module world_models.inference.operators.jepa_operator module world_models.inference.operators.planet_operator module world_models.losses.convae_loss module world_models.losses.gmm_loss module world_models.masks module world_models.masks.default 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.controller module world_models.models.diffusion module world_models.models.diffusion.actor_critic module world_models.models.diffusion.DDPM module world_models.models.diffusion.diamond_diffusion module world_models.models.diffusion.DiT module world_models.models.diffusion.reward_termination module world_models.models.dreamer module world_models.models.dreamer_rssm module world_models.models.dynamics_model module world_models.models.genie module world_models.models.iris_agent module world_models.models.iris_transformer module world_models.models.jepa_agent module world_models.models.latent_action_model module world_models.models.mdrnn module world_models.models.modular_rssm module world_models.models.planet module world_models.models.rssm module world_models.models.vit module world_models.training.rl_harness module world_models.training.train_controller module world_models.training.train_convvae module world_models.training.train_diamond module world_models.training.train_genie module world_models.training.train_iris module world_models.training.train_jepa module world_models.training.train_mdn_rnn module world_models.training.train_planet module world_models.training.train_rssm module world_models.training.train_world_model module world_models.transforms.image module world_models.utils.data_utils module world_models.utils.dreamer_utils module world_models.utils.jepa_utils module world_models.utils.jit_utils module world_models.utils.logging_utils module world_models.utils.memory_utils module world_models.utils.train_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.planet_decoder module world_models.vision.planet_encoder module world_models.vision.VAE.ConvVAE module world_models.vision.video_tokenizer module world_models.vision.vq_layer module WorldModelEnv (class in world_models.envs) write_frame() (world_models.utils.utils.StreamingVideoWriter method)

X

x (torchwm.Episode attribute) (world_models.Episode attribute) (world_models.memory.planet_memory.Episode attribute)