protomotions.agents.masked_mimic.config module¶

class protomotions.agents.masked_mimic.config.KLDScheduleConfig(init_kld_coeff=0.0001, end_kld_coeff=0.01, start_epoch=3000, end_epoch=6000)[source]¶

Bases: ConfigBuilder

Configuration for KL divergence scheduling in VAE training.

init_kld_coeff: float = 0.0001¶

end_kld_coeff: float = 0.01¶

start_epoch: int = 3000¶

end_epoch: int = 6000¶

__init__(init_kld_coeff=0.0001, end_kld_coeff=0.01, start_epoch=3000, end_epoch=6000)¶

class protomotions.agents.masked_mimic.config.VaeNoiseType(value)[source]¶

Bases: Enum

NORMAL = 'normal'¶

UNIFORM = 'uniform'¶

ZEROS = 'zeros'¶

classmethod from_str(value)[source]¶

Create enum from string, case-insensitive.

class protomotions.agents.masked_mimic.config.VaeConfig(kld_schedule=<factory>, vae_latent_dim=64, vae_noise_type=VaeNoiseType.NORMAL)[source]¶

Bases: ConfigBuilder

Configuration for VAE-specific parameters.

kld_schedule: KLDScheduleConfig¶

vae_latent_dim: int = 64¶

vae_noise_type: VaeNoiseType = 'normal'¶

__init__(kld_schedule=<factory>, vae_latent_dim=64, vae_noise_type=VaeNoiseType.NORMAL)¶

class protomotions.agents.masked_mimic.config.FeedForwardModelConfig(_target_='protomotions.agents.masked_mimic.model.FeedForwardModel', in_keys=<factory>, out_keys=<factory>, trunk=<factory>)[source]¶

Bases: BaseModelConfig

Configuration for FeedForwardModel.

trunk: SequentialModuleConfig¶

__init__(_target_='protomotions.agents.masked_mimic.model.FeedForwardModel', in_keys=<factory>, out_keys=<factory>, trunk=<factory>)¶

class protomotions.agents.masked_mimic.config.MaskedMimicModelConfig(_target_='protomotions.agents.masked_mimic.model.MaskedMimicModel', in_keys=<factory>, out_keys=<factory>, encoder=<factory>, prior=<factory>, trunk=<factory>, vae=<factory>, optimizer=<factory>)[source]¶

Bases: BaseModelConfig

Configuration for MaskedMimic Model (VAE-based imitation learning).

encoder: MultiOutputModuleConfig¶

prior: MultiOutputModuleConfig¶

trunk: SequentialModuleConfig¶

vae: VaeConfig¶

optimizer: OptimizerConfig¶

__init__(_target_='protomotions.agents.masked_mimic.model.MaskedMimicModel', in_keys=<factory>, out_keys=<factory>, encoder=<factory>, prior=<factory>, trunk=<factory>, vae=<factory>, optimizer=<factory>)¶

class protomotions.agents.masked_mimic.config.MaskedMimicAgentConfig(batch_size, training_max_steps, _target_='protomotions.agents.masked_mimic.agent.MaskedMimic', model=<factory>, num_steps=32, gradient_clip_val=0.0, fail_on_bad_grads=False, check_grad_mag=True, gamma=0.99, bounds_loss_coef=0.0, task_reward_w=1.0, num_mini_epochs=1, training_early_termination=None, save_epoch_checkpoint_every=1000, save_last_checkpoint_every=10, max_episode_length_manager=None, evaluator=<factory>, normalize_rewards=True, normalized_reward_clamp_value=5.0, expert_model_path=None)[source]¶

Bases: BaseAgentConfig

Main configuration class for MaskedMimic Agent.

model: MaskedMimicModelConfig | FeedForwardModelConfig¶

expert_model_path: str | None = None¶

__init__(batch_size, training_max_steps, _target_='protomotions.agents.masked_mimic.agent.MaskedMimic', model=<factory>, num_steps=32, gradient_clip_val=0.0, fail_on_bad_grads=False, check_grad_mag=True, gamma=0.99, bounds_loss_coef=0.0, task_reward_w=1.0, num_mini_epochs=1, training_early_termination=None, save_epoch_checkpoint_every=1000, save_last_checkpoint_every=10, max_episode_length_manager=None, evaluator=<factory>, normalize_rewards=True, normalized_reward_clamp_value=5.0, expert_model_path=None)¶