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_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_size (torchwm.IRISActor attribute) (torchwm.IRISReplayBuffer 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.utils.utils.TorchImageEnvWrapper property) action_space (world_models.envs.BraxImageEnv 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.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.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_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) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig 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) 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) 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) (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_size (torchwm.DiamondConfig 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.genie_config.GenieConfig attribute) (world_models.configs.genie_config.GenieSmallConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.configs.GenieSmallConfig attribute) (world_models.datasets.tinyworlds.TinyWorldsConfig attribute) (world_models.datasets.video_datasets.DatasetConfig attribute) (world_models.DiamondConfig 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) (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) (in module world_models.benchmarks.metrics) bootstrap_iqm_ci() (in module world_models.benchmarks.metrics) bottle() (in module world_models.utils.utils) BraxImageEnv (class in world_models.envs) buffer_capacity (torchwm.IRISReplayBuffer property) (world_models.IRISReplayBuffer property) (world_models.memory.iris_memory.IRISReplayBuffer property) 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) CHANNELS (torchwm.DiTConfig attribute) (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) (world_models.DiTConfig 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.gym_env.GymImageEnv method) (world_models.envs.GymImageEnv method) (world_models.envs.mujoco_env.MuJoCoImageEnv method) (world_models.envs.MuJoCoImageEnv method) (world_models.envs.unity_env.UnityMLAgentsEnv method) (world_models.envs.UnityMLAgentsEnv method) (world_models.envs.vector_env.VectorizedEnv method) (world_models.utils.utils.StreamingVideoWriter method) (world_models.utils.utils.TorchImageEnvWrapper method) 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_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) (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)

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) 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_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

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_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) 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) deter_size (torchwm.ModularRSSM property) (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) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) DiamondAdapter (class in world_models.benchmarks) (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) 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) 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) 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) (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) DreamerV1Adapter (class in world_models.benchmarks.adapters) DreamerV2Adapter (class in world_models.benchmarks.adapters) 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) 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

EDMPreconditioner (class in world_models.models.diffusion.diamond_diffusion) 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_mixed_precision() (in module world_models.utils.memory_utils) 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_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) 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) 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() (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.benchmarks.BaseAdapter method) (world_models.benchmarks.DiamondAdapter method) (world_models.benchmarks.DreamerAdapter method) (world_models.benchmarks.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) export_csv() (in module world_models.benchmarks.reporting) export_latex() (in module world_models.benchmarks.reporting) export_markdown() (in module world_models.benchmarks.reporting) extensions (world_models.datasets.video_datasets.VideoDatasetConfig 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) flush() (world_models.utils.dreamer_utils.Logger 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.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.Genie 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.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.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.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.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) FreezeParameters (class in world_models.utils.dreamer_utils)

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.transforms) generate() (torchwm.Genie method) (world_models.Genie method) (world_models.models.genie.Genie method) 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_state() (torchwm.RecurrentStateSpaceModel method) (world_models.models.rssm.RecurrentStateSpaceModel method) (world_models.RecurrentStateSpaceModel method) 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_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) gpu_timer() (in module world_models.utils.jepa_utils) grad_logger() (in module world_models.utils.jepa_utils) GRUBackbone (class in world_models.models.modular_rssm) GymImageEnv (class in world_models.envs) (class in world_models.envs.gym_env)

H

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) hidden_size (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)

I

image_size (torchwm.DynamicsModelConfig attribute) (torchwm.GenieConfig attribute) (torchwm.GenieSmallConfig attribute) (torchwm.LatentActionModelConfig 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.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.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_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_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) 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) interpolate_pos_encoding() (world_models.models.vit.VisionTransformer method) iqm_of_array() (in module world_models.benchmarks) (in module world_models.benchmarks.metrics) IRISActor (class in torchwm) (class in world_models) (class in world_models.controller.iris_policy) IRISAdapter (class in world_models.benchmarks) (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)

L

lambda_returns (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig 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) 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.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_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) 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.benchmarks.BaseAdapter method) (world_models.benchmarks.DiamondAdapter method) (world_models.benchmarks.DreamerAdapter method) (world_models.benchmarks.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_state_dict() (world_models.datasets.diamond_dataset.ReplayBuffer method) load_yml_config() (in module world_models.utils.utils) log() (world_models.utils.jepa_utils.CSVLogger method) log_abs_det_jacobian() (torchwm.TanhBijector method) (world_models.TanhBijector method) (world_models.vision.dreamer_decoder.TanhBijector method) log_interval (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig attribute) log_scalar() (world_models.utils.dreamer_utils.Logger method) log_scalars() (world_models.utils.dreamer_utils.Logger method) log_videos() (world_models.utils.dreamer_utils.Logger method) Logger (class in world_models.utils.dreamer_utils) LR (torchwm.DiTConfig attribute) (world_models.configs.dit_config.DiTConfig attribute) (world_models.configs.DiTConfig attribute) (world_models.DiTConfig attribute) LSTMBackbone (class in world_models.models.modular_rssm)

M

main() (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) 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_cifar10() (in module world_models.datasets.cifar10) make_diamond_atari_env() (in module world_models.envs.diamond_atari) 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_robotics_env() (in module world_models.envs) (in module world_models.envs.robotics_env) make_transforms() (in module world_models.transforms.transforms) make_unity_mlagents_env() (in module world_models.envs) (in module world_models.envs.unity_env) 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) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.GenieConfig attribute) (world_models.training.train_genie.GenieConfig attribute) MaskGITSampler (class in world_models.models.dynamics_model) max_episode_steps (world_models.envs.BraxImageEnv 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.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) 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) 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) 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.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.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.tinyworlds world_models.datasets.video_datasets 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.gym_env world_models.envs.mujoco_env world_models.envs.robotics_env world_models.envs.unity_env world_models.envs.vector_env world_models.envs.wrappers 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.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.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.modular_rssm world_models.models.planet world_models.models.rssm world_models.models.vit world_models.training.rl_harness world_models.training.train_diamond world_models.training.train_genie world_models.training.train_iris world_models.training.train_jepa world_models.training.train_planet world_models.training.train_rssm world_models.transforms.transforms 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.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.video_tokenizer world_models.vision.vq_layer MuJoCoImageEnv (class in world_models.envs) (class in world_models.envs.mujoco_env) MultiAgentBenchmarkRunner (class in world_models.benchmarks) (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

name (torchwm.EnvBackendSpec attribute) (torchwm.ModelSpec attribute) (torchwm.SampleDist property) (world_models.api.EnvBackendSpec attribute) (world_models.api.ModelSpec attribute) (world_models.EnvBackendSpec attribute) (world_models.ModelSpec attribute) (world_models.SampleDist property) (world_models.vision.dreamer_decoder.SampleDist property) normalize_frames_for_saving() (in module world_models.utils.utils) NormalizeActions (class in world_models.envs) (class in world_models.envs.wrappers) 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.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_workers (world_models.datasets.tinyworlds.TinyWorldsConfig attribute) (world_models.datasets.video_datasets.DatasetConfig attribute) NumPyDataset (class in world_models.datasets.video_datasets)

O

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.DeepMindControlEnv attribute) (world_models.envs.DeepMindControlEnv property) (world_models.envs.dmc.DeepMindControlEnv attribute) (world_models.envs.dmc.DeepMindControlEnv 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.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)

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)

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) 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) pin_memory (world_models.datasets.video_datasets.DatasetConfig attribute) Planet (class in world_models.models.planet) PlaNetOperator (class in torchwm) (class in world_models) (class in world_models.inference.operators.planet_operator)

play() (torchwm.Genie method) (world_models.Genie method) (world_models.models.genie.Genie method) plot_metrics() (in module world_models.utils.utils) 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) postprocess_img() (in module world_models.utils.utils) PPOTrainer (class in world_models.training.rl_harness) 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_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.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)

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) 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() (world_models.envs.BraxImageEnv method) (world_models.envs.DeepMindControlEnv method) (world_models.envs.dmc.DeepMindControlEnv 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.unity_env.UnityMLAgentsEnv method) (world_models.envs.UnityMLAgentsEnv 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) 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.DeepMindControlEnv method) (world_models.envs.diamond_atari.DiamondAtariWrapper method) (world_models.envs.dmc.DeepMindControlEnv 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.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.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_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) RolloutGenerator (class in world_models.controller.rollout_generator) 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) RSSMPolicy (class in torchwm) (class in world_models) (class in world_models.controller.rssm_policy) run() (world_models.benchmarks.BenchmarkRunner method) (world_models.benchmarks.runner.BenchmarkRunner method) (world_models.envs.vector_env.SimWorker method) run_all() (world_models.benchmarks.MultiAgentBenchmarkRunner method) (world_models.benchmarks.runner.MultiAgentBenchmarkRunner method)

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_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) (world_models.configs.genie_config.GenieConfig attribute) (world_models.configs.GenieConfig attribute) (world_models.GenieConfig 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() (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_video() (in module world_models.utils.utils) seed (torchwm.DiamondConfig attribute) (world_models.configs.diamond_config.DiamondConfig attribute) (world_models.configs.DiamondConfig attribute) (world_models.DiamondConfig 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.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) setup_logging() (in module world_models.utils.logging_utils) 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) split (world_models.datasets.tinyworlds.TinyWorldsConfig attribute) state_dict() (world_models.datasets.diamond_dataset.ReplayBuffer 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) step() (world_models.envs.ActionRepeat method) (world_models.envs.BraxImageEnv method) (world_models.envs.DeepMindControlEnv method) (world_models.envs.diamond_atari.DiamondAtariWrapper method) (world_models.envs.dmc.DeepMindControlEnv 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.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.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.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.ModularRSSM property) (world_models.models.modular_rssm.BackboneBase attribute) (world_models.models.modular_rssm.ModularRSSM property) (world_models.ModularRSSM 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) 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)

T

t (torchwm.Episode attribute) (world_models.Episode attribute) (world_models.memory.planet_memory.Episode attribute) TanhBijector (class in torchwm) (class in world_models) (class in world_models.vision.dreamer_decoder) TensorBoardMetrics (class in world_models.utils.utils) 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) 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() (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.JEPAConfig method) to_tensor_obs() (in module world_models.utils.utils) 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 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_diamond() (in module world_models.training.train_diamond) train_epoch() (world_models.training.train_iris.IRISTrainer method) train_one_batch() (world_models.models.dreamer.Dreamer method) train_rssm() (in module world_models.training.train_rssm) 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) 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) 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_transformer() (torchwm.IRISAgent method) (world_models.IRISAgent method) (world_models.models.iris_agent.IRISAgent method) UpsampleBlock (class in world_models.vision.iris_decoder) 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) 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_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

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) (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) 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.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.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.tinyworlds module world_models.datasets.video_datasets 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.gym_env module world_models.envs.mujoco_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.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.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.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.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_diamond module world_models.training.train_genie module world_models.training.train_iris module world_models.training.train_jepa module world_models.training.train_planet module world_models.training.train_rssm module world_models.transforms.transforms 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.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.video_tokenizer module world_models.vision.vq_layer module write_frame() (world_models.utils.utils.StreamingVideoWriter method)

X

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