protomotions.agents.masked_mimic.model module#

class protomotions.agents.masked_mimic.model.FeedForwardModel(*args, **kwargs)[source]#

Bases: BaseModel

Simple feedforward model for masked mimic without VAE.

__init__(config)[source]#

forward(tensordict)[source]#

Forward pass computing action.

Parameters:: tensordict (MockTensorDict) – TensorDict containing observations.
Returns:: TensorDict with action added.
Return type:: MockTensorDict

class protomotions.agents.masked_mimic.model.MaskedMimicModel(*args, **kwargs)[source]#

Bases: BaseModel

MaskedMimic model architecture with Variational Autoencoder (VAE).

Combines a prior network (acting on sparse observations) and an encoder network (acting on full/expert observations) to learn a latent space for motion control.

The prior learns to predict the expert’s latent distribution from sparse data, enabling the agent to perform robust control when full state info is unavailable during inference.

Parameters:: config (MaskedMimicModelConfig) – Configuration for the MaskedMimic model.

config: MaskedMimicModelConfig#

__init__(config)[source]#

Initialize the MaskedMimic model components.

static reparameterization(mean, std, vae_noise)[source]#: Reparameterization trick: z = mu + std * noise

forward(tensordict)[source]#

Forward pass through MaskedMimic model.

Always computes both prior and encoder for consistency and ONNX compatibility. Expects vae_noise to be provided in tensordict (generated by agent).

Parameters:: tensordict (MockTensorDict) – TensorDict containing observations and vae_noise.
Returns:: TensorDict with action and all VAE outputs.
Return type:: MockTensorDict

forward_inference(tensordict)[source]#

Inference-only forward pass (prior path only, no encoder).

Use this for ONNX export and deployment. Only computes the action from the prior network, excluding the encoder which is only needed during training.

Parameters:: tensordict (MockTensorDict) – TensorDict containing prior observations and vae_noise.
Returns:: TensorDict with action only.
Return type:: MockTensorDict

get_inference_in_keys()[source]#

Get input keys needed for inference (prior + trunk only, no encoder).

kl_loss(tensordict)[source]#

Compute KL divergence between encoder and prior.

Parameters:: tensordict (MockTensorDict) – TensorDict containing prior and encoder outputs.
Returns:: KL divergence tensor.