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, cyclical=False, )[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. cyclical: Repeat the inclusive start/end ramp after end_epoch, restarting at start_epoch on end_epoch + 1.

init_kld_coeff: float = 0.0001#

end_kld_coeff: float = 0.01#

start_epoch: int = 3000#

end_epoch: int = 6000#

cyclical: bool = False#

__init__( init_kld_coeff=0.0001, end_kld_coeff=0.01, start_epoch=3000, end_epoch=6000, cyclical=False, )#

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

Bases: Enum

Geometry of samples passed to the VAE decoder.

GAUSSIAN = 'gaussian'#

SPHERE = 'sphere'#

class protomotions.agents.supervised.masked_mimic_config.MaskedMimicVAEConfig( kld_schedule=<factory>, vae_latent_dim=64, vae_noise_type=VAENoiseType.NORMAL, latent_geometry=LatentGeometry.GAUSSIAN, prior_mean_unit_norm_coeff=0.0, )[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. latent_geometry: Decoder latent geometry: gaussian or sphere. prior_mean_unit_norm_coeff: Penalty on the learned prior mean’s distance from unit norm.

kld_schedule: KLDScheduleConfig | None#

vae_latent_dim: int = 64#

vae_noise_type: VAENoiseType = 'normal'#

latent_geometry: LatentGeometry = 'gaussian'#

prior_mean_unit_norm_coeff: float = 0.0#

__init__( kld_schedule=<factory>, vae_latent_dim=64, vae_noise_type=VAENoiseType.NORMAL, latent_geometry=LatentGeometry.GAUSSIAN, prior_mean_unit_norm_coeff=0.0, )#

class protomotions.agents.supervised.masked_mimic_config.TemporalLossConfig( key, coefficient, prediction_key=None, forward=False, )[source]#

Bases: object

Stop-gradient MSE between a model output and an adjacent rollout value.

prediction_key (default key) is compared against key from the previous step, or from the next step when forward is set. Both keys must be declared model outputs. Pairs that cross an episode reset are excluded.

key: str#

coefficient: float#

prediction_key: str | None = None#

forward: bool = False#

property target_key: str#

property valid_key: str#

property direction: str#

__init__( key, coefficient, prediction_key=None, forward=False, )#

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>, temporal_losses=<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. temporal_losses: Stop-gradient MSE terms against adjacent rollout outputs. optimizer: Optimizer settings for supervised training.

encoder: ModuleContainerConfig#

prior: ModuleContainerConfig#

trunk: ModuleContainerConfig#

vae: MaskedMimicVAEConfig#

temporal_losses: List[TemporalLossConfig]#

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>, temporal_losses=<factory>, optimizer=<factory>, )#

class protomotions.agents.supervised.masked_mimic_config.MaskedMimicSupervisedAgentConfig( batch_size, training_max_steps, training_max_iterations=None, _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>, expert_rollout_probability=0.0, )[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. training_max_iterations: Maximum number of complete rollout and optimization iterations. When set, this takes precedence over training_max_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: External expert checkpoint that labels rollouts; requires rollout_actor=student. rollout_actor: Model forward used for collecting rollout actions. loss: Supervision loss over MaskedMimic outputs. expert_rollout_probability: Per-episode probability that the external expert drives the episode.

model: MaskedMimicModelConfig#

loss: SupervisionLossConfig#

apply_inference_defaults(env)[source]#

Student alone (the supervised cleanup), driven by the learned prior.

Training evaluates the posterior (it has the full future); inference has only the masked goal, so the prior is the action path.

__init__( batch_size, training_max_steps, training_max_iterations=None, _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>, expert_rollout_probability=0.0, )#