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

kl_loss(tensordict)[source]¶

Compute KL divergence between encoder and prior.

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