protomotions.agents package

Subpackages

• protomotions.agents.amp package
  • Submodules
    • protomotions.agents.amp.agent module
      • AMP
    • protomotions.agents.amp.config module
      • AMPParametersConfig
      • DiscriminatorConfig
      • AMPModelConfig
      • AMPAgentConfig
    • protomotions.agents.amp.model module
      • Discriminator
      • AMPModelComponentsMixin
      • AMPModel
• protomotions.agents.ase package
  • Submodules
    • protomotions.agents.ase.agent module
      • ASE
    • protomotions.agents.ase.config module
      • ASEParametersConfig
      • ASEDiscriminatorEncoderConfig
      • ASEModelConfig
      • ASEAgentConfig
    • protomotions.agents.ase.model module
      • ASEDiscriminatorEncoder
      • ASEModel
• protomotions.agents.base_agent package
  • Submodules
    • protomotions.agents.base_agent.agent module
      • BaseAgent
    • protomotions.agents.base_agent.config module
      • MaxEpisodeLengthManagerConfig
      • OptimizerConfig
      • MuonWithAuxAdamConfig
      • BaseModelConfig
      • BaseAgentConfig
    • protomotions.agents.base_agent.model module
      • RolloutStateSpec
      • ProtoMotionsTensorDictModule
      • BaseModel
• protomotions.agents.callbacks package
  • Submodules
    • protomotions.agents.callbacks.slurm_autoresume_srun module
      • wandb_run_exists()
      • AutoResumeCallbackSrun
• protomotions.agents.common package
  • Submodules
    • protomotions.agents.common.autoencoder package
      • AutoEncoder
    • protomotions.agents.common.autoregressive module
      • sampling_log_probs()
      • prior_constrained_sampling_log_probs()
      • nucleus_sampling()
      • nucleus_sampling_prior_constraint()
      • kl_divergence_categorical()
      • kl_divergence_from_log_probs()
      • kl_divergence_sampling_distribution()
      • generate_causal_mask()
      • resolve_discrete_autoregressive_config()
      • DiscreteAutoregressiveTransformer
    • protomotions.agents.common.common module
      • get_params()
      • weight_init()
      • weight_init_trainable()
      • NormObsBase
      • apply_module_operations()
      • ObsProcessor
      • ModuleContainer
    • protomotions.agents.common.config module
      • NormObsBaseConfig
      • ModuleOperationConfig
      • ModuleOperationForwardConfig
      • ModuleOperationPermuteConfig
      • ModuleOperationReshapeConfig
      • ModuleOperationSqueezeConfig
      • ModuleOperationUnsqueezeConfig
      • ModuleOperationExpandConfig
      • ModuleOperationSphereProjectionConfig
      • ObsProcessorConfig
      • MLPLayerConfig
      • MLPWithConcatConfig
      • ModuleContainerConfig
      • PretrainedModelConfig
      • TransformerConfig
      • DiscreteAutoregressiveTransformerConfig
    • protomotions.agents.common.discrete_latent module
      • DiscreteLatentQuantizer
      • FSQTokenization
      • DiscreteLatentDecoder
      • DiscreteLatentTargetEncoder
      • make_discrete_latent_decoder()
      • load_pretrained_discrete_latent_decoder()
      • make_discrete_latent_target_encoder()
      • load_pretrained_discrete_latent_target_encoder()
    • protomotions.agents.common.fsq module
      • FiniteScalarQuantizer
      • FSQAutoEncoder
    • protomotions.agents.common.fsq_config module
      • FSQAutoEncoderConfig
    • protomotions.agents.common.latent module
      • compute_discrete_latent_ppo_loss()
    • protomotions.agents.common.mlp module
      • build_mlp()
      • MLPWithConcat
    • protomotions.agents.common.pretrained module
      • freeze_module()
      • get_path()
      • load_pretrained_model_module()
    • protomotions.agents.common.supervision module
      • SupervisionLossType
      • SupervisionLossConfig
      • compute_supervision_loss()
    • protomotions.agents.common.transformer module
      • Transformer
• protomotions.agents.evaluators package
  • MimicEvaluator
    • __init__()
    • cleanup_after_evaluation()
    • evaluate_episode()
    • initialize_eval()
    • motion_lib
    • motion_manager
    • process_eval_results()
    • run_evaluation()
  • SmoothnessCalculator
    • __init__()
    • compute_normalized_jerk_from_pos()
    • compute_smoothness_metrics()
  • MotionMetrics
    • __init__()
    • compute_finite_difference_jitter_reduce_each_motion()
    • compute_jitter_reduce_each_motion()
    • compute_rotation_jitter_reduce_each_motion()
    • copy_from()
    • copy_from_motion_ids()
    • get_unfilled_mask()
    • jitter_mean_reduce_each_motion()
    • max_mean_reduce()
    • max_reduce_each_motion()
    • mean_max_reduce()
    • mean_mean_reduce()
    • mean_min_reduce()
    • mean_reduce_each_motion()
    • merge_from()
    • min_mean_reduce()
    • min_reduce_each_motion()
    • ops_mean_reduce()
    • reset()
    • rotation_jitter_mean_reduce_each_motion()
    • to()
    • update()
  • AggregateMetric
    • compute()
  • SmoothnessAggregateMetric
    • __init__()
    • compute()
  • ActionSmoothnessAggregateMetric
    • __init__()
    • compute()
  • Submodules
    • protomotions.agents.evaluators.base_evaluator module
      • BaseEvaluator
    • protomotions.agents.evaluators.config module
      • EvaluatorConfig
      • MotionWeightsRulesConfig
      • MimicEvaluatorConfig
    • protomotions.agents.evaluators.metrics module
      • MotionMetrics
    • protomotions.agents.evaluators.mimic_evaluator module
      • MimicEpisodeContext
      • MimicEvaluator
    • protomotions.agents.evaluators.smoothness_calculator module
      • SmoothnessCalculator
• protomotions.agents.fine_tuning package
  • FineTuningAgent
    • create_model()
  • FineTuningAgentConfig
    • __init__()
    • model
    • pretrained_modules
  • PretrainedModulesMixin
    • pretrained
  • Submodules
    • protomotions.agents.fine_tuning.agent module
      • FineTuningAgent
    • protomotions.agents.fine_tuning.config module
      • FineTuningAgentConfig
    • protomotions.agents.fine_tuning.pretrained_modules module
      • PretrainedModulesMixin
• protomotions.agents.mimic package
  • Submodules
• protomotions.agents.peft package
  • Subpackages
    • protomotions.agents.peft.utils package
      • Submodules
  • Submodules
    • protomotions.agents.peft.actor module
      • DiscretePriorPEFTActor
    • protomotions.agents.peft.adapters module
      • LoRALayer
      • TransformerLayerWithLoRA
      • TransformerLayerWithDoRA
      • TransformerEncoderWithConditioning
      • freeze_base_and_enable_peft()
      • set_peft_layers_train_mode()
      • inject_transformer_peft()
    • protomotions.agents.peft.config module
      • TransformerPEFTConfig
    • protomotions.agents.peft.model module
      • DiscretePriorPEFTModel
    • protomotions.agents.peft.prior_agent module
      • DiscretePriorPEFTRLFTAgent
    • protomotions.agents.peft.prior_amp_agent module
      • DiscretePriorPEFTRLFTAMPAgent
    • protomotions.agents.peft.prior_amp_config module
      • DiscretePriorPEFTRLFTAMPModelConfig
      • DiscretePriorPEFTRLFTAMPAgentConfig
    • protomotions.agents.peft.prior_amp_model module
      • DiscretePriorPEFTRLFTAMPModel
    • protomotions.agents.peft.prior_config module
      • default_peft_model_config()
      • DiscretePriorPEFTConfig
      • DiscretePriorPEFTActorConfig
      • DiscretePriorPEFTBaseModelConfig
      • DiscretePriorPEFTRLFTModelConfig
      • DiscretePriorPEFTSFTModelConfig
      • DiscretePriorPEFTRLFTAgentConfig
      • DiscretePriorPEFTSFTAgentConfig
    • protomotions.agents.peft.prior_setup module
      • DiscretePriorPEFTSetupMixin
    • protomotions.agents.peft.prior_with_peft module
      • DiscretePriorWithPEFT
    • protomotions.agents.peft.sft_agent module
      • DiscretePriorPEFTSFTAgent
    • protomotions.agents.peft.sft_model module
      • DiscretePriorPEFTSFTModel
• protomotions.agents.ppo package
  • Submodules
    • protomotions.agents.ppo.agent module
      • PPO
    • protomotions.agents.ppo.config module
      • PPOActorConfig
      • PPOModelConfig
      • AdvantageNormalizationConfig
      • AdaptiveLRConfig
      • L2C2Config
      • PPOAgentConfig
    • protomotions.agents.ppo.model module
      • PPOActor
      • PPOModel
    • protomotions.agents.ppo.utils module
      • discount_values()
• protomotions.agents.supervised package
  • Submodules
    • protomotions.agents.supervised.agent module
      • SupervisedAgent
    • protomotions.agents.supervised.config module
      • RolloutActor
      • SupervisedAgentConfig
    • protomotions.agents.supervised.expert_utils module
      • get_expert_actor_in_keys()
      • get_expert_observation_keys()
      • get_expert_observation_components()
    • protomotions.agents.supervised.masked_mimic_config module
      • KLDScheduleConfig
      • VAENoiseType
      • MaskedMimicVAEConfig
      • MaskedMimicModelConfig
      • MaskedMimicSupervisedAgentConfig
    • protomotions.agents.supervised.masked_mimic_model module
      • MaskedMimicModel
• protomotions.agents.utils package
  • Submodules
    • protomotions.agents.utils.data module
      • swap_and_flatten01()
      • ExperienceBuffer
      • DictDataset
    • protomotions.agents.utils.metering module
      • Timer
      • TimeReport
      • TensorAverageMeter
      • TensorAverageMeterDict
    • protomotions.agents.utils.normalization module
      • RunningMeanStd
      • materialize_lazy_running_stats_from_state_dict()
      • combine_moments()
      • RewardRunningMeanStd
    • protomotions.agents.utils.replay_buffer module
      • ReplayBuffer
    • protomotions.agents.utils.step_tracker module
      • StepTracker
    • protomotions.agents.utils.training module
      • bounds_loss()
      • handle_model_grad_clipping()
      • aggregate_scalar_metrics()
      • get_activation_func()