protomotions.agents.supervised.masked_mimic_config module#

MaskedMimic configs for the generic supervised agent.

class protomotions.agents.supervised.masked_mimic_config.KLDScheduleConfig( init_kld_coeff=0.0001, end_kld_coeff=0.01, start_epoch=3000, end_epoch=6000, )[source]#

Bases: object

KL coefficient schedule for the MaskedMimic VAE posterior loss.

Attributes:: init_kld_coeff: Initial KL divergence coefficient. end_kld_coeff: Final KL divergence coefficient. start_epoch: Epoch to start KL coefficient annealing. end_epoch: Epoch to end KL coefficient annealing.

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.supervised.masked_mimic_config.VAENoiseType(value)[source]#

Bases: Enum

Noise distribution used for MaskedMimic latent sampling.

NORMAL = 'normal'#

UNIFORM = 'uniform'#

ZEROS = 'zeros'#

classmethod from_str(value)[source]#

Create enum from a case-insensitive string.

class protomotions.agents.supervised.masked_mimic_config.MaskedMimicVAEConfig( kld_schedule=<factory>, vae_latent_dim=64, vae_noise_type=VAENoiseType.NORMAL, )[source]#

Bases: object

VAE settings for the MaskedMimic learned-prior student.

Attributes:: kld_schedule: KL divergence annealing schedule. vae_latent_dim: Dimension of VAE latent space. vae_noise_type: Type of latent noise: normal, uniform, or zeros.

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.supervised.masked_mimic_config.MaskedMimicModelConfig( _target_='protomotions.agents.supervised.masked_mimic_model.MaskedMimicModel', in_keys=<factory>, out_keys=<factory>, encoder=<factory>, prior=<factory>, trunk=<factory>, vae=<factory>, optimizer=<factory>, )[source]#

Bases: BaseModelConfig

MaskedMimic VAE learned-prior model configuration.

Attributes:: in_keys: Input keys. out_keys: Output keys. encoder: Privileged encoder network. prior: Deployable prior network. trunk: Latent-to-action decoder trunk. vae: MaskedMimic VAE settings. optimizer: Optimizer settings for supervised training.

encoder: ModuleContainerConfig#

prior: ModuleContainerConfig#

trunk: ModuleContainerConfig#

vae: MaskedMimicVAEConfig#

optimizer: OptimizerConfig#

__init__( _target_='protomotions.agents.supervised.masked_mimic_model.MaskedMimicModel', in_keys=<factory>, out_keys=<factory>, encoder=<factory>, prior=<factory>, trunk=<factory>, vae=<factory>, optimizer=<factory>, )#

class protomotions.agents.supervised.masked_mimic_config.MaskedMimicSupervisedAgentConfig( batch_size, training_max_steps, _target_='protomotions.agents.supervised.agent.SupervisedAgent', 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, save_inference_checkpoint=False, max_episode_length_manager=None, evaluator=<factory>, normalize_rewards=True, normalized_reward_clamp_value=5.0, reward_norm_ema_decay=None, expert_model_path=None, rollout_actor=RolloutActor.STUDENT, loss=<factory>, )[source]#

Bases: SupervisedAgentConfig

MaskedMimic preset for SupervisedAgent.

The training loop is generic supervised imitation. The student is the MaskedMimic VAE learned-prior model, and the default supervision target is the privileged action produced by that model.

Attributes:: batch_size: Training batch size. training_max_steps: Maximum training steps. model: MaskedMimic model configuration. num_steps: Environment steps per update. gradient_clip_val: Max gradient norm. 0=disabled. fail_on_bad_grads: Fail on NaN/Inf gradients. check_grad_mag: Log gradient magnitude. gamma: Discount factor. bounds_loss_coef: Action bounds loss. 0 for tanh outputs. task_reward_w: Task reward weight. num_mini_epochs: Mini-epochs per update. training_early_termination: Stop early at this step. None=disabled. save_epoch_checkpoint_every: Save epoch_xxx.ckpt every N epochs. save_last_checkpoint_every: Save last.ckpt every K epochs. save_inference_checkpoint: Also save inference_<name>.ckpt without optimizer or other training-only state. max_episode_length_manager: Episode length curriculum. evaluator: Evaluation config. normalize_rewards: Normalize rewards. normalized_reward_clamp_value: Clamp normalized rewards to [-val, val]. reward_norm_ema_decay: EMA decay for reward normalization (None = Welford). Set to e.g. 0.99 to track non-stationary reward distributions. expert_model_path: Optional checkpoint for an external expert policy. rollout_actor: Policy used for collecting rollout actions. loss: Supervision loss over MaskedMimic outputs.

model: MaskedMimicModelConfig#

loss: SupervisionLossConfig#

__init__( batch_size, training_max_steps, _target_='protomotions.agents.supervised.agent.SupervisedAgent', 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, save_inference_checkpoint=False, max_episode_length_manager=None, evaluator=<factory>, normalize_rewards=True, normalized_reward_clamp_value=5.0, reward_norm_ema_decay=None, expert_model_path=None, rollout_actor=RolloutActor.STUDENT, loss=<factory>, )#