protomotions.agents.supervised.masked_mimic_model module

Latent decoding

The prior outputs its Gaussian mean and log-variance. The encoder outputs a mean residual (added to the prior mean) and the posterior log-variance. The decoder receives the sampled latent and its configured observations. Inference runs only the prior and decoder; it never requires encoder inputs.

MaskedMimicVAEConfig(latent_geometry=LatentGeometry.SPHERE) projects sampled prior and posterior latents to unit length before decoding. The Gaussian parameters remain unprojected for the KL regularizer; this is not a KL between spherical densities. An exactly zero sample maps to the first unit basis vector. Gaussian sampling remains the default. prior_mean_unit_norm_coeff optionally penalizes (||prior_mean|| - 1)^2 for spherical models.

Decoder observation statistics are recorded once on the posterior pass during training, not twice for the prior and posterior. Deterministic goal features must be computed from deployable inputs; masking and token construction belong in the configured prior modules. The model does not infer or reconstruct hidden goals.

MaskedMimic VAE student for the generic supervised agent.

MaskedMimic is trained by SupervisedAgent, but the architecture itself is not generic supervised-agent machinery. It is a VAE learned-prior student:

• prior predicts a deployable latent distribution from sparse observations.

• encoder predicts a privileged residual posterior during training.

• trunk decodes latent samples to actions.

Keeping this model named after MaskedMimic makes experiment files searchable and keeps the VAE-specific KL/noise logic out of the generic supervised loop.

class protomotions.agents.supervised.masked_mimic_model.MaskedMimicModel(*args, **kwargs)

Bases: BaseModel

MaskedMimic learned-prior VAE student.

The non-privileged path is used at inference: prior -> latent sample -> trunk. During supervised training, the privileged encoder adds a residual posterior path that produces privileged_action for the imitation loss and a KL term against the prior distribution.

__init__(config)

build_training_targets(rollout)

Derive model-owned targets from [time, env, …] rollout tensors.

BaseAgent calls this once per rollout under no_grad, after agent preprocessing and before flattening/shuffling destroys time adjacency. Inputs are borrowed experience-buffer views: do not mutate them. Return additional [time, env, …] targets/masks for compute_model_loss; the agent owns batching, not their meaning. The default gathers child targets; overrides should call super() to preserve those results.

rollout_state_specs()

Local model-owned per-env rollout state declarations.

forward(tensordict, log_internals=False)

Forward pass through the model.

Parameters:

tensordict (MockTensorDict) – TensorDict containing observations.

Returns:

TensorDict with model outputs added.

Return type:

MockTensorDict

forward_deployable(tensordict)

The deployable forward: by default the same graph the policy runs.

A policy family whose exported pass differs from its training forward (an actor that samples exploration noise) overrides this; everything else inherits it, so the exporter, evaluator and teacher call one name without an isinstance check on the policy class.

property deployable_in_keys: list

The input keys the deployable graph reads; in_keys by default.

materialize(tensordict)

Deployment initializes the prior; checkpoint loading supplies the posterior.

Training supplies all observations and must initialize both paths before DDP. A partially supplied posterior is still an input-contract error.

kl_loss(tensordict)

compute_model_loss(

tensordict,

current_epoch,

zero_loss,

log_prefix='model',

)

Optional module-owned auxiliary loss for agent optimization loops.

Most modules do not own an auxiliary loss. Models that do, such as a VAE-backed policy head, override this and return (loss, log_dict).

Temporal objectives

temporal_losses lists TemporalLossConfig entries. Each compares a declared model output (prediction_key, default key) against the rollout value of key at the previous step, or at the next step when forward=True. Targets are held fixed during optimization. For the Metatron recipe, use TemporalLossConfig("privileged_action", 0.05), TemporalLossConfig("solution_features", 0.001). These previous-step terms do not compare against the action selected by an expert-driven rollout. A next-step objective is also supported, but requires the model to produce the configured prediction key; it is not enabled in this recipe. All terms average squared feature errors over valid transitions, with no timestep scaling, and are logged as model/temporal_{previous|next}_{prediction_key}.

The model’s build_training_targets hook constructs targets on time-by-env tensors before flattening and shuffling. A reset after time t invalidates the t-to-t+1 pair. Missing neighbors at buffer edges are excluded. Targets are stop-gradient rollout values, not recomputed neighboring minibatch rows. Model rollout state, including latent noise, is replayed from the collected batch.