protomotions.agents.supervised.masked_mimic_model module

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)

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_inference(tensordict)

get_inference_in_keys()

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).