protomotions.agents.peft package

Parameter-efficient fine-tuning workflows for pretrained motion priors.

The PEFT package contains the discrete-prior PEFT actor/agent stack used by GPC SFT and RLFT experiments, plus adapter wrappers for transformer layers.

Subpackages

• protomotions.agents.peft.utils package
  • Submodules
    • protomotions.agents.peft.utils.adapter_state module
      • strip_actor_prefix()
      • is_adapter_state_key()
      • contains_adapter_state()
      • build_adapter_state_dict()
      • load_adapter_state_dict()
    • protomotions.agents.peft.utils.frozen_prior_checkpoint module
      • reload_frozen_base_prior_from_checkpoint()
    • protomotions.agents.peft.utils.frozen_prior_contract module
      • require_frozen_prior_attr()
      • resolve_frozen_prior_input_keys()
    • protomotions.agents.peft.utils.model_state module
      • has_reference_state()
      • load_compatible_peft_model_state()
      • optional_full_checkpoint_state_prefixes()
      • is_adapter_only_model_state()

Submodules

• protomotions.agents.peft.actor module
  • DiscretePriorPEFTActor
    • __init__()
    • optional_full_checkpoint_state_prefixes()
    • adapter_state_dict()
    • load_adapter_state_dict()
    • init_peft()
    • train()
    • quantize()
    • fsq_codes_to_fsq_indices()
    • fsq_indices_to_codes()
    • fsq_indices_to_prior_tokens()
    • prior_tokens_to_fsq_indices()
    • one_hot_prior_tokens()
    • perturb_tokens()
    • predict_target_prior_tokens()
    • build_prior_input()
    • forward()
    • get_action_and_logp()
• protomotions.agents.peft.adapters module
  • LoRALayer
    • __init__()
    • forward()
  • TransformerLayerWithLoRA
    • __init__()
    • forward()
  • TransformerLayerWithDoRA
    • __init__()
    • forward()
  • TransformerEncoderWithConditioning
    • __init__()
    • forward()
  • freeze_base_and_enable_peft()
  • set_peft_layers_train_mode()
  • inject_transformer_peft()
• protomotions.agents.peft.config module
  • TransformerPEFTConfig
    • peft_type
    • rank
    • alpha
    • __init__()
• protomotions.agents.peft.model module
  • DiscretePriorPEFTModel
    • __init__()
    • optional_full_checkpoint_state_prefixes()
    • forward()
    • forward_rollout()
• protomotions.agents.peft.prior_agent module
  • DiscretePriorPEFTRLFTAgent
    • __init__()
    • load()
    • fit()
    • actor_step()
    • critic_step()
    • perform_optimization_step()
• protomotions.agents.peft.prior_amp_agent module
  • DiscretePriorPEFTRLFTAMPAgent
    • __init__()
• protomotions.agents.peft.prior_amp_config module
  • DiscretePriorPEFTRLFTAMPModelConfig
    • discriminator
    • discriminator_optimizer
    • disc_critic
    • disc_critic_optimizer
    • __init__()
  • DiscretePriorPEFTRLFTAMPAgentConfig
    • model
    • amp_parameters
    • reference_obs_components
    • __init__()
• protomotions.agents.peft.prior_amp_model module
  • DiscretePriorPEFTRLFTAMPModel
    • __init__()
    • optional_full_checkpoint_state_prefixes()
    • experience_buffer_keys()
    • forward_rollout()
• protomotions.agents.peft.prior_config module
  • default_peft_model_config()
  • DiscretePriorPEFTConfig
    • model
    • condition_key
    • temperature
    • top_p
    • sampling_mode
    • prior_top_p
    • kl_coeff
    • clear_peft
    • m_clamp
    • film_input_norm
    • film_input_norm_clamp
    • resolve_model()
    • __init__()
  • DiscretePriorPEFTActorConfig
    • in_keys
    • out_keys
    • peft
    • __init__()
  • DiscretePriorPEFTBaseModelConfig
    • actor
    • actor_optimizer
    • __init__()
  • DiscretePriorPEFTRLFTModelConfig
    • critic
    • critic_optimizer
    • __init__()
  • DiscretePriorPEFTSFTModelConfig
    • token_perturb_rate
    • token_perturb_mode
    • __init__()
  • DiscretePriorPEFTRLFTAgentConfig
    • pretrained_modules
    • model
    • save_inference_checkpoint
    • num_mini_epochs
    • gradient_clip_val
    • entropy_coef
    • target_kl
    • __init__()
  • DiscretePriorPEFTSFTAgentConfig
    • pretrained_modules
    • model
    • rollout_actor
    • loss
    • save_inference_checkpoint
    • num_mini_epochs
    • gradient_clip_val
    • __init__()
• protomotions.agents.peft.prior_setup module
  • DiscretePriorPEFTSetupMixin
    • require_reward_norm_on_load
    • pretrained
    • create_model()
    • get_inference_state_dict()
    • load_adapter_checkpoint()
• protomotions.agents.peft.prior_with_peft module
  • DiscretePriorWithPEFT
    • __init__()
    • task_c_dim
    • init_peft()
    • reference_state_prefixes()
    • optional_full_checkpoint_state_prefixes()
    • reference_ready
    • ensure_reference_modules()
    • mark_reference_loaded()
    • require_reference()
    • capture_reference()
    • train()
    • forward()
• protomotions.agents.peft.sft_agent module
  • DiscretePriorPEFTSFTAgent
    • __init__()
    • has_critic
    • create_model()
    • create_optimizers()
    • get_state_dict()
• protomotions.agents.peft.sft_model module
  • DiscretePriorPEFTSFTModel
    • collect_expert_rollout()
    • materialize()
    • forward()