protomotions.agents.common.autoencoder package

Shared encoder-bottleneck-decoder TensorDict modules.

This file intentionally stays narrow. AutoEncoder is useful for models that really are encoder -> bottleneck -> decoder, such as FSQ trackers and small reconstruction-style students. Autoregressive GPC priors are not autoencoders and should use their own BaseModel implementation instead.

class protomotions.agents.common.autoencoder.AutoEncoder(*args, **kwargs)

Bases: BaseModel

Generic encoder-bottleneck-decoder module.

Subclasses customize only the bottleneck behavior. For example, an FSQ tracker quantizes the encoder latent before decoding. Models with a different shape, such as causal token priors, should not inherit this class.

supports_log_internals = False

__init__(config)

config: AutoEncoderConfig

bottleneck(latent, tensordict)

internal_logs(latent, tensordict)

predict_latent(tensordict)

decode(tensordict, latent)

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