protomotions.agents.common package

Submodules

• protomotions.agents.common.autoencoder package
  • AutoEncoder
    • supports_log_internals
    • __init__()
    • config
    • bottleneck()
    • internal_logs()
    • predict_latent()
    • decode()
    • forward()
• 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
    • __init__()
    • config
    • muon_adam_fallback_modules()
    • compute_model_loss()
    • encode_context()
    • encode_tokens()
    • decode_logits()
    • add_positions()
    • run_transformer()
    • forward_from_tokens()
    • teacher_force()
    • forward()
    • next_logits_from_context()
    • generate_from_context()
• protomotions.agents.common.common module
  • get_params()
  • weight_init()
  • weight_init_trainable()
  • NormObsBase
    • running_obs_norm
    • num_out
    • __init__()
    • build_norm()
    • forward()
  • apply_module_operations()
  • ObsProcessor
    • __init__()
    • config
    • forward()
  • ModuleContainer
    • __init__()
    • config
    • forward()
• protomotions.agents.common.config module
  • NormObsBaseConfig
    • normalize_obs
    • norm_clamp_value
    • norm_ema_decay
    • __init__()
  • ModuleOperationConfig
    • __init__()
  • ModuleOperationForwardConfig
    • __init__()
  • ModuleOperationPermuteConfig
    • new_order
    • __init__()
  • ModuleOperationReshapeConfig
    • new_shape
    • __init__()
  • ModuleOperationSqueezeConfig
    • squeeze_dim
    • __init__()
  • ModuleOperationUnsqueezeConfig
    • unsqueeze_dim
    • __init__()
  • ModuleOperationExpandConfig
    • expand_shape
    • __init__()
  • ModuleOperationSphereProjectionConfig
    • __init__()
  • ObsProcessorConfig
    • in_keys
    • out_keys
    • module_operations
    • __init__()
  • MLPLayerConfig
    • units
    • activation
    • use_layer_norm
    • __init__()
  • MLPWithConcatConfig
    • num_out
    • layers
    • in_keys
    • out_keys
    • output_activation
    • module_operations
    • __init__()
  • ModuleContainerConfig
    • models
    • in_keys
    • out_keys
    • __init__()
  • PretrainedModelConfig
    • checkpoint_path
    • module_path
    • module_config_path
    • module_config
    • config_path
    • state_dict_key
    • strict
    • freeze
    • __init__()
  • TransformerConfig
    • in_keys
    • out_keys
    • input_and_mask_mapping
    • transformer_token_size
    • latent_dim
    • num_heads
    • ff_size
    • num_layers
    • dropout
    • activation
    • output_activation
    • __init__()
  • DiscreteAutoregressiveTransformerConfig
    • in_keys
    • out_keys
    • context_key
    • token_key
    • logits_key
    • generated_tokens_key
    • logprob_key
    • context_embedding_key
    • token_embedding_key
    • hidden_key
    • context_encoder
    • token_encoder
    • output_head
    • d_model
    • num_heads
    • num_layers
    • ff_size
    • dropout
    • activation
    • num_tokens
    • vocab_size
    • max_seq_len
    • __init__()
• protomotions.agents.common.discrete_latent module
  • DiscreteLatentQuantizer
    • num_fsq_levels
    • num_fsq_scalars
    • quantize()
    • codes_to_indices()
    • indices_to_codes()
    • __init__()
  • FSQTokenization
    • __init__()
    • fsq_indices_to_prior_tokens()
    • prior_tokens_to_fsq_indices()
    • one_hot_prior_tokens()
  • DiscreteLatentDecoder
    • __init__()
    • num_fsq_levels
    • num_fsq_scalars
    • freeze()
    • train()
    • indices_to_codes()
    • decode()
  • DiscreteLatentTargetEncoder
    • __init__()
    • freeze()
    • train()
    • compute_model_loss()
  • 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
    • __init__()
    • bound()
    • round_ste()
    • quantize()
    • codes_to_indices()
    • indices_to_codes()
    • calculate_perplexity()
  • FSQAutoEncoder
    • supports_log_internals
    • config
    • __init__()
    • num_fsq_levels
    • num_fsq_scalars
    • L
    • half_width
    • half_L
    • bound()
    • round_ste()
    • codes_to_indices()
    • indices_to_codes()
    • quantize()
    • calculate_perplexity()
    • bottleneck()
    • internal_logs()
• protomotions.agents.common.fsq_config module
  • FSQAutoEncoderConfig
    • num_fsq_levels
    • num_fsq_scalars
    • __init__()
• protomotions.agents.common.latent module
  • compute_discrete_latent_ppo_loss()
• protomotions.agents.common.mlp module
  • build_mlp()
  • MLPWithConcat
    • norm
    • mlp
    • in_keys
    • out_keys
    • __init__()
    • config
    • forward()
• protomotions.agents.common.pretrained module
  • freeze_module()
  • get_path()
  • load_pretrained_model_module()
• protomotions.agents.common.supervision module
  • SupervisionLossType
    • MSE
    • DISCRETE_CROSS_ENTROPY
    • DISCRETE_KL
    • CONTINUOUS_GAUSSIAN_KL
  • SupervisionLossConfig
    • loss_type
    • prediction_key
    • target_key
    • prediction_logvar_key
    • target_logvar_key
    • label_smoothing
    • weight
    • log_prefix
    • enabled
    • extra
    • __init__()
  • compute_supervision_loss()
• protomotions.agents.common.transformer module
  • Transformer
    • input_models
    • sequence_pos_encoder
    • seqTransEncoder
    • in_keys
    • out_keys
    • __init__()
    • forward()