Index A | B | C | D | E | F | G | H | I | K | L | M | N | O | P | R | S | T | U | V | W | X A action() (ContinuousActionWrapper method) (DiscreteActionWrapper method) (ExplorationWrapper method) (MultiBinaryActionWrapper method) action_space (GymnaxCompatibility property) actions (Batch attribute) (FrozenTraj attribute) (RLData attribute) actions_differentiable (Beta property) (Discrete property) (GMM property) (Multibinary property) (PolicyOutput property) (TanhGaussian property) activation_switch() (in module amago.nets.utils) Actor (class in amago.nets.actor_critic) AdamWRel (class in amago.utils) add_activation_log() (in module amago.nets.utils) add_common_cli() (in module amago.cli_utils) add_exploration_noise() (BilevelEpsilonGreedy method) (ExplorationWrapper method) add_goal_info() (RelabelInfoWrapper method) add_score() (ReturnHistory method) (SpecialMetricHistory method) add_timestep() (Trajectory method) Agent (class in amago.agent) agent_type (Experiment attribute) ALE (class in amago.envs.builtin.ale_retro) ALEAction (class in amago.envs.builtin.ale_retro) AlreadyVectorizedEnv (class in amago.envs.env_utils) always_load_latest (Experiment attribute) always_save_latest (Experiment attribute) amago module amago.agent module amago.cli_utils module amago.envs module amago.envs.amago_env module amago.envs.builtin module amago.envs.builtin.alchemy module amago.envs.builtin.ale_retro module amago.envs.builtin.babyai module amago.envs.builtin.gymnax_envs module amago.envs.builtin.half_cheetah_v4_vel module amago.envs.builtin.mazerunner module amago.envs.builtin.metaworld_ml module amago.envs.builtin.popgym_envs module amago.envs.builtin.procgen_envs module amago.envs.builtin.tmaze module amago.envs.builtin.toy_gym module amago.envs.builtin.xland_minigrid module amago.envs.env_utils module amago.envs.exploration module amago.experiment module amago.hindsight module amago.loading module amago.nets module amago.nets.actor_critic module amago.nets.cnn module amago.nets.ff module amago.nets.goal_embedders module amago.nets.policy_dists module amago.nets.traj_encoders module amago.nets.transformer module amago.nets.tstep_encoders module amago.nets.utils module amago.utils module amago_warning() (in module amago.utils) AMAGOEnv (class in amago.envs.amago_env) AmagoWarning as_input() (Timestep method) as_input_sequence() (Trajectory method) async_env_mp_context (Experiment attribute) AtariAMAGOWrapper (class in amago.envs.builtin.ale_retro) AtariGame (class in amago.envs.builtin.ale_retro) AttentionLayer (class in amago.nets.transformer) avg_over_accelerate() (in module amago.utils) B Batch (class in amago.loading) batch_size (Experiment attribute) batched_envs (Timestep attribute) batches_per_update (Experiment attribute) Beta (class in amago.nets.policy_dists) BilevelEpsilonGreedy (class in amago.envs.exploration) bin_dist_to_raw_vals() (NCriticsTwoHot method) binary_filter() (in module amago.agent) blank_img (CNN property) C Cache (class in amago.nets.transformer) cached_training_mask() (FlexAttention method) call_async() (AlreadyVectorizedEnv method) (DummyAsyncVectorEnv method) call_async_env() (in module amago.utils) call_wait() (AlreadyVectorizedEnv method) (DummyAsyncVectorEnv method) caster() (Experiment method) check_configured() (RLDataset method) ckpt_base_dir (Experiment attribute) ckpt_interval (Experiment attribute) close() (RetroArcade method) CNN (class in amago.nets.cnn) CNNTstepEncoder (class in amago.nets.tstep_encoders) collect_new_training_data() (Experiment method) compute_loss() (Experiment method) configure() (RLDataset method) configure_from_experiment() (MixtureOfDatasets method) (RLDataset method) console_cabinets (RetroArcade attribute) ContinuousActionWrapper (class in amago.envs.env_utils) conv_forward() (CNN method) (DrQCNN method) (GridworldCNN method) (IMPALAishCNN method) (NatureishCNN method) count_params() (in module amago.utils) create_experiment_from_cli() (in module amago.cli_utils) create_obs() (POPGym method) create_observation() (MultitaskMetaBabyAI method) create_reset_version() (Timestep method) critic_loss_weight (Experiment attribute) current_eps() (BilevelEpsilonGreedy method) current_timestep (SequenceWrapper property) D dataset (Experiment attribute) delete() (DiskTrajDataset method) (MixtureOfDatasets method) (RLDataset method) delete_buffer_from_disk() (Experiment method) denormalize_values() (InputNorm method) DEVICE (Experiment property) Discrete (class in amago.nets.policy_dists) discrete_action_n (RetroArcade attribute) DiscreteActionWrapper (class in amago.envs.env_utils) DiscreteLikeContinuous (class in amago.nets.policy_dists) DiskTrajDataset (class in amago.loading) dloader_workers (Experiment attribute) dones (Batch attribute) (FrozenTraj attribute) (RLData attribute) DoNothingDataset (class in amago.loading) DrQCNN (class in amago.nets.cnn) DrQv2Aug (class in amago.nets.cnn) DummyAsyncVectorEnv (class in amago.envs.env_utils) E edit_actor_mask() (Experiment method) edit_critic_mask() (Experiment method) emb_dim (CNNTstepEncoder property) (FFTrajEncoder property) (FFTstepEncoder property) (GRUTrajEncoder property) (MambaTrajEncoder property) (TformerTrajEncoder property) (TrajEncoder property) (Transformer property) (TstepEncoder property) entropy() (DiscreteLikeContinuous method) env_mode (Experiment attribute) env_name (AMAGOEnv property) (AtariAMAGOWrapper property) (Metaworld property) (MultiDomainPOPGymAMAGO property) (ProcgenAMAGO property) (RetroAMAGOWrapper property) (SequenceWrapper property) EnvCreator (class in amago.envs.amago_env) epochs (Experiment attribute) EpsilonGreedy (class in amago.envs.exploration) evaluate_test() (Experiment method) evaluate_val() (Experiment method) exp_filter() (in module amago.agent) Experiment (class in amago.experiment) exploration_wrapper_type (EnvCreator attribute) (Experiment attribute) ExplorationWrapper (class in amago.envs.exploration) extend_box_obs_space_by() (in module amago.envs.env_utils) F FFBlock (class in amago.nets.ff) FFGoalEmb (class in amago.nets.goal_embedders) FFTrajEncoder (class in amago.nets.traj_encoders) FFTstepEncoder (class in amago.nets.tstep_encoders) finish_active_traj() (SequenceWrapper method) FixedPosEmb (class in amago.nets.transformer) FlashAttention (class in amago.nets.transformer) flex_attention() (FlexAttention method) flex_attention_inf() (FlexAttention method) FlexAttention (class in amago.nets.transformer) force_reset_train_envs_every (Experiment attribute) forward() (Actor method) (Agent method) (AttentionLayer method) (Beta method) (CNN method) (Discrete method) (DrQv2Aug method) (FFBlock method) (FFGoalEmb method) (FFTrajEncoder method) (FixedPosEmb method) (FlashAttention method) (FlexAttention method) (GMM method) (GRUTrajEncoder method) (InputNorm method) (LearnablePosEmb method) (MambaTrajEncoder method) (MLP method) (Multibinary method) (MultiTaskAgent method) (NCritics method) (NCriticsTwoHot method) (Normalization method) (PolicyOutput method) (PopArtLayer method) (SelfAttention method) (SigmaReparam method) (SlowAdaptiveRational method) (TanhGaussian method) (TformerTrajEncoder method) (TokenGoalEmb method) (TrajEncoder method) (Transformer method) (TransformerLayer method) (TstepEncoder method) (VanillaAttention method) frame() (TwoAttemptMTProcgen method) freeze() (Trajectory method) from_dict() (FrozenTraj class method) FrozenTraj (class in amago.hindsight) G gammas_as_input_seq() (in module amago.nets.actor_critic) generate_task() (RoomKeyDoor method) get_actions() (Agent method) get_constant_schedule_with_warmup() (in module amago.utils) get_description() (DiskTrajDataset method) (DoNothingDataset method) (MixtureOfDatasets method) (RLDataset method) get_env_funcs() (KEpisodeMetaworld method) get_grad_norm() (in module amago.utils) get_obs() (KEpisodeMetaworld method) (XLandMinigridVectorizedGym method) get_path_to_trajs() (in module amago.loading) gin_as_wandb_config() (in module amago.utils) GMM (class in amago.nets.policy_dists) go_back_to_start() (MazeRunnerGymEnv method) goal_sequence() (MazeRunnerGymEnv method) grad_clip (Experiment attribute) GridworldCNN (class in amago.nets.cnn) GRUTrajEncoder (class in amago.nets.traj_encoders) GymnaxCompatibility (class in amago.envs.builtin.gymnax_envs) H HalfCheetahV4_MetaVelocity (class in amago.envs.builtin.half_cheetah_v4_vel) HalfCheetahV4LogVelocity (class in amago.envs.builtin.half_cheetah_v4_vel) hard_sync_targets() (Agent method) has_dset_edit_rights (Experiment attribute) I IMPALAishCNN (class in amago.nets.cnn) inference_forward() (Transformer method) init_checkpoints() (Experiment method) init_dloaders() (Experiment method) init_dsets() (Experiment method) init_envs() (Experiment method) init_hidden_state() (MambaTrajEncoder method) (TformerTrajEncoder method) (TrajEncoder method) init_logger() (Experiment method) init_model() (Experiment method) init_new_env() (SymbolicAlchemy method) init_optimizer() (Experiment method) inner_forward() (CNNTstepEncoder method) (FFTstepEncoder method) (TstepEncoder method) inner_reset() (AMAGOEnv method) inner_step() (AMAGOEnv method) input_dimension (Beta property) (Discrete property) (GMM property) (Multibinary property) (PolicyOutput property) (TanhGaussian property) InputNorm (class in amago.nets.utils) interact() (Experiment method) is_discrete (Beta property) (Discrete property) (GMM property) (Multibinary property) (PolicyOutput property) (TanhGaussian property) K KEpisodeMetaworld (class in amago.envs.builtin.metaworld_ml) kv_cache_mask_mod() (FlexAttention method) kv_cache_score_mod() (FlexAttention method) L l2_coeff (Experiment attribute) learn() (Experiment method) LearnablePosEmb (class in amago.nets.transformer) learning_rate (Experiment attribute) load_checkpoint() (Experiment method) load_checkpoint_from_path() (Experiment method) load_traj_from_disk() (in module amago.loading) log() (Experiment method) log_interval (Experiment attribute) log_prefix (SpecialMetricHistory attribute) log_prob() (DiscreteLikeContinuous method) log_to_wandb (Experiment attribute) logits (DiscreteLikeContinuous property) lr_warmup_steps (Experiment attribute) M make_action_rep() (AMAGOEnv method) make_env (EnvCreator attribute) make_experiment_collect_only() (in module amago.cli_utils) make_experiment_learn_only() (in module amago.cli_utils) make_obs() (MetaFrozenLake method) (MultiDomainPOPGym method) make_train_env (Experiment attribute) make_val_env (Experiment attribute) MambaTrajEncoder (class in amago.nets.traj_encoders) masked_avg() (in module amago.utils) masked_stats() (InputNorm method) max_seq_len (Experiment attribute) MazeRunnerAMAGOEnv (class in amago.envs.builtin.mazerunner) MazeRunnerGymEnv (class in amago.envs.builtin.mazerunner) MetaFrozenLake (class in amago.envs.builtin.toy_gym) Metaworld (class in amago.envs.builtin.metaworld_ml) mixed_precision (Experiment attribute) MixtureOfDatasets (class in amago.loading) MLP (class in amago.nets.ff) module amago amago.agent amago.cli_utils amago.envs amago.envs.amago_env amago.envs.builtin amago.envs.builtin.alchemy amago.envs.builtin.ale_retro amago.envs.builtin.babyai amago.envs.builtin.gymnax_envs amago.envs.builtin.half_cheetah_v4_vel amago.envs.builtin.mazerunner amago.envs.builtin.metaworld_ml amago.envs.builtin.popgym_envs amago.envs.builtin.procgen_envs amago.envs.builtin.tmaze amago.envs.builtin.toy_gym amago.envs.builtin.xland_minigrid amago.envs.env_utils amago.envs.exploration amago.experiment amago.hindsight amago.loading amago.nets amago.nets.actor_critic amago.nets.cnn amago.nets.ff amago.nets.goal_embedders amago.nets.policy_dists amago.nets.traj_encoders amago.nets.transformer amago.nets.tstep_encoders amago.nets.utils amago.utils mt_names (MultiDomainPOPGym attribute) Multibinary (class in amago.nets.policy_dists) MultiBinaryActionWrapper (class in amago.envs.env_utils) MultiDomainPOPGym (class in amago.envs.builtin.popgym_envs) MultiDomainPOPGymAMAGO (class in amago.envs.builtin.popgym_envs) Multigammas (class in amago.agent) MultiTaskAgent (class in amago.agent) MultitaskMetaBabyAI (class in amago.envs.builtin.babyai) N NatureishCNN (class in amago.nets.cnn) NCritics (class in amago.nets.actor_critic) NCriticsTwoHot (class in amago.nets.actor_critic) new_reset_keys() (XLandMinigridVectorizedGym method) new_ruleset_keys() (XLandMinigridVectorizedGym method) NoOpRelabeler (class in amago.hindsight) Normalization (class in amago.nets.ff) normalize_values() (InputNorm method) (PopArtLayer method) O obs (Batch attribute) (FrozenTraj attribute) (RLData attribute) (Timestep attribute) obs() (RoomKeyDoor method) observation_space (GymnaxCompatibility property) on_end_of_collection() (DiskTrajDataset method) (MixtureOfDatasets method) (RLDataset method) P padded_sampling (Experiment attribute) parallel_actors (Experiment attribute) pick_new_game() (ALE method) policy (Experiment property) policy_metrics() (Experiment method) PolicyOutput (class in amago.nets.policy_dists) PopArtLayer (class in amago.nets.actor_critic) POPGym (class in amago.envs.builtin.popgym_envs) POPGymAMAGO (class in amago.envs.builtin.popgym_envs) preprocess_seq() (Transformer method) prev_action (Timestep attribute) probs (DiscreteLikeContinuous property) ProcgenAMAGO (class in amago.envs.builtin.procgen_envs) R random_maze() (in module amago.envs.builtin.mazerunner) random_slice() (RLData method) random_traj_length() (SequenceWrapper method) raw_vals_to_labels() (NCriticsTwoHot method) read_latest_policy() (Experiment method) ready_for_training (DiskTrajDataset property) (DoNothingDataset property) (MixtureOfDatasets property) (RLDataset property) relabel() (NoOpRelabeler method) (Relabeler method) Relabeler (class in amago.hindsight) RelabelInfoWrapper (class in amago.envs.builtin.mazerunner) render() (AlreadyVectorizedEnv method) (AMAGOEnv method) (AtariGame method) (DummyAsyncVectorEnv method) (MazeRunnerGymEnv method) (MetaFrozenLake method) (RetroArcade method) (RoomKeyDoor method) (XLandMinigridVectorizedGym method) reset() (ALE method) (AlreadyVectorizedEnv method) (AMAGOEnv method) (AtariGame method) (BilevelEpsilonGreedy method) (ContinuousActionWrapper method) (DiscreteActionWrapper method) (DummyAsyncVectorEnv method) (GymnaxCompatibility method) (HalfCheetahV4_MetaVelocity method) (HalfCheetahV4LogVelocity method) (KEpisodeMetaworld method) (MazeRunnerGymEnv method) (MetaFrozenLake method) (MultiDomainPOPGym method) (MultitaskMetaBabyAI method) (POPGym method) (RelabelInfoWrapper method) (RetroArcade method) (RoomKeyDoor method) (SequenceWrapper method) (SymbolicAlchemy method) (TformerHiddenState method) (TwoAttemptMTProcgen method) (XLandMinigridVectorizedGym method) reset_hidden_state() (GRUTrajEncoder method) (MambaTrajEncoder method) (TformerTrajEncoder method) (TrajEncoder method) reset_if_done() (XLandMinigridVectorizedGym method) reset_same_task() (RoomKeyDoor method) reset_stats() (SequenceWrapper method) RetroAMAGOWrapper (class in amago.envs.builtin.ale_retro) RetroArcade (class in amago.envs.builtin.ale_retro) retry_load_checkpoint() (in module amago.utils) return_history (ExplorationWrapper property) ReturnHistory (class in amago.envs.amago_env) reward (Timestep attribute) reward_fn() (TMazeAlt method) reward_scales (KEpisodeMetaworld attribute) rews (Batch attribute) (FrozenTraj attribute) (RLData attribute) rl2s (Batch attribute) (FrozenTraj attribute) (RLData attribute) RLData (class in amago.loading) RLData_pad_collate() (in module amago.loading) RLDataset (class in amago.loading) roll_back() (Cache method) RoomKeyDoor (class in amago.envs.builtin.toy_gym) rsample() (DiscreteLikeContinuous method) run_name (Experiment attribute) S sample() (DiscreteLikeContinuous method) sample_actions (Experiment attribute) sample_new_env() (MultitaskMetaBabyAI method) sample_random_trajectory() (DiskTrajDataset method) (DoNothingDataset method) (MixtureOfDatasets method) (RLDataset method) sample_target_velocity() (HalfCheetahV4_MetaVelocity method) save_checkpoint() (Experiment method) save_every_high (EnvCreator attribute) save_every_low (EnvCreator attribute) save_finished_trajs() (SequenceWrapper method) save_new_trajs_to (DiskTrajDataset property) (DoNothingDataset property) (MixtureOfDatasets property) (RLDataset property) save_to_disk() (Trajectory method) save_trajs_as (EnvCreator attribute) (Experiment attribute) save_trajs_to (EnvCreator attribute) screen() (AtariGame method) (RetroArcade method) SelfAttention (class in amago.nets.transformer) SequenceWrapper (class in amago.envs.amago_env) sigma (InputNorm property) (PopArtLayer property) SigmaReparam (class in amago.nets.transformer) SlidingWindowFlexAttention (class in amago.nets.transformer) SlowAdaptiveRational (class in amago.nets.utils) soft_reset() (MetaFrozenLake method) soft_sync_targets() (Agent method) softplus_bounded_positive() (in module amago.nets.policy_dists) space_convert() (in module amago.envs.env_utils) SpecialMetricHistory (class in amago.envs.amago_env) split_batch() (in module amago.utils) split_batched_timestep() (in module amago.hindsight) split_dict() (in module amago.utils) stack_list_array_dicts() (in module amago.utils) stagger_traj_file_lengths (Experiment attribute) start() (Experiment method) start_collecting_at_epoch (Experiment attribute) start_learning_at_epoch (Experiment attribute) start_plotting() (MazeRunnerGymEnv method) step() (AdamWRel method) (ALE method) (AlreadyVectorizedEnv method) (AMAGOEnv method) (AtariGame method) (BilevelEpsilonGreedy method) (DummyAsyncVectorEnv method) (GymnaxCompatibility method) (HalfCheetahV4_MetaVelocity method) (HalfCheetahV4LogVelocity method) (KEpisodeMetaworld method) (MazeRunnerGymEnv method) (MetaFrozenLake method) (MultiDomainPOPGym method) (MultitaskMetaBabyAI method) (POPGym method) (RelabelInfoWrapper method) (RetroArcade method) (RoomKeyDoor method) (SequenceWrapper method) (SymbolicAlchemy method) (TMazeAlt method) (TwoAttemptMTProcgen method) (XLandMinigridVectorizedGym method) step_count (SequenceWrapper property) str2action (ALEAction attribute) success_history (ExplorationWrapper property) suggested_max_seq_len (XLandMinigridVectorizedGym property) sum_over_accelerate() (in module amago.utils) summary() (Experiment method) switch_agent() (in module amago.cli_utils) switch_async_mode() (in module amago.cli_utils) switch_exploration() (in module amago.cli_utils) switch_traj_encoder() (in module amago.cli_utils) switch_tstep_encoder() (in module amago.cli_utils) SymbolicAlchemy (class in amago.envs.builtin.alchemy) symexp() (in module amago.nets.utils) symlog() (in module amago.nets.utils) T tanh_bounded_positive() (in module amago.nets.policy_dists) TanhGaussian (class in amago.nets.policy_dists) terminal (Timestep attribute) TformerHiddenState (class in amago.nets.transformer) TformerTrajEncoder (class in amago.nets.traj_encoders) time_idx (Timestep attribute) time_idxs (Batch attribute) (FrozenTraj attribute) (RLData attribute) Timestep (class in amago.hindsight) TMazeAlt (class in amago.envs.builtin.tmaze) TMazeAltActive (class in amago.envs.builtin.tmaze) TMazeAltPassive (class in amago.envs.builtin.tmaze) to() (Batch method) (PopArtLayer method) to_dict() (FrozenTraj method) to_discrete() (ALEAction method) TokenGoalEmb (class in amago.nets.goal_embedders) total_frames (SequenceWrapper property) total_frames_by_env_name (SequenceWrapper property) total_return (Trajectory property) train_batches_per_epoch (Experiment attribute) train_step() (Experiment method) train_timesteps_per_epoch (Experiment attribute) trainable_params (Agent property) training_forward() (Transformer method) traj_encoder_type (Experiment attribute) traj_save_len (Experiment attribute) Trajectory (class in amago.hindsight) TrajEncoder (class in amago.nets.traj_encoders) Transformer (class in amago.nets.transformer) TransformerLayer (class in amago.nets.transformer) tstep_encoder_type (Experiment attribute) TstepEncoder (class in amago.nets.tstep_encoders) TwoAttemptMTProcgen (class in amago.envs.builtin.procgen_envs) U update() (TformerHiddenState method) update_dset_weights() (MixtureOfDatasets method) update_stats() (InputNorm method) (PopArtLayer method) use_config() (in module amago.cli_utils) V val_interval (Experiment attribute) val_timesteps_per_epoch (Experiment attribute) VanillaAttention (class in amago.nets.transformer) VanillaFlexAttention (class in amago.nets.transformer) velocity_reward_term() (HalfCheetahV4_MetaVelocity method) verbose (Experiment attribute) W wandb_entity (Experiment attribute) wandb_group_name (Experiment attribute) wandb_project (Experiment attribute) weight_init() (in module amago.nets.cnn) WORDS (MultitaskMetaBabyAI attribute) write_latest_policy() (Experiment method) X x_axis_metrics() (Experiment method) XLandMinigridVectorizedGym (class in amago.envs.builtin.xland_minigrid)