protomotions.agents.common.discrete_latent module#

Reusable helpers for discrete latent targets, decoders, and GPC priors.

class protomotions.agents.common.discrete_latent.DiscreteLatentQuantizer(*args, **kwargs)[source]#

Bases: Protocol

Interface required by discrete latent adapters.

FSQ-style quantizers expose scalar-code dimensions with FSQ terminology so callers do not confuse them with autoregressive prior tokens.

num_fsq_levels: int#

num_fsq_scalars: int#

quantize(latent)[source]#

codes_to_indices(codes)[source]#

indices_to_codes(indices)[source]#

__init__(*args, **kwargs)#

class protomotions.agents.common.discrete_latent.FSQTokenization( *, num_fsq_levels, num_fsq_scalars, fsq_scalars_per_prior_token, )[source]#

Bases: <Mock object at 0x7fd696f00490>[]

Pack flat FSQ scalar indices into categorical prior tokens.

FSQ produces one small discrete index per latent scalar. The GPC prior consumes fewer categorical tokens by packing several FSQ scalar indices into a mixed-radix vocabulary entry.

__init__( *, num_fsq_levels, num_fsq_scalars, fsq_scalars_per_prior_token, )[source]#

fsq_indices_to_prior_tokens(fsq_indices)[source]#

Pack (batch, num_fsq_scalars) FSQ indices into prior tokens.

prior_tokens_to_fsq_indices(prior_tokens)[source]#

Unpack prior tokens back to flat FSQ scalar indices.

one_hot_prior_tokens(prior_tokens)[source]#

One-hot encode prior-token indices using the derived vocabulary.

class protomotions.agents.common.discrete_latent.DiscreteLatentDecoder( *, decoder, quantizer, latent_key='latent', decoder_out_key=None, freeze=True, )[source]#

Bases: <Mock object at 0x7fd696bd26d0>[]

Frozen FSQ-code-to-action decoder path.

__init__( *, decoder, quantizer, latent_key='latent', decoder_out_key=None, freeze=True, )[source]#

property num_fsq_levels: int#

property num_fsq_scalars: int#

freeze()[source]#

train(mode=True)[source]#

indices_to_codes(indices)[source]#

decode(tensordict, codes)[source]#

class protomotions.agents.common.discrete_latent.DiscreteLatentTargetEncoder(*args, **kwargs)[source]#

Bases: ProtoMotionsTensorDictModule

Encode target observations to FSQ indices, then pack prior tokens.

__init__( *, encoder, quantizer, tokenization, target_key, encoder_out_key=None, freeze=True, )[source]#

freeze()[source]#

train(mode=True)[source]#

compute_model_loss( tensordict, current_epoch, zero_loss, log_prefix='model', )[source]#

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

protomotions.agents.common.discrete_latent.make_discrete_latent_decoder(module, *, freeze=True)[source]#

Adapt a module exposing decoder/quantizer to DiscreteLatentDecoder.

protomotions.agents.common.discrete_latent.load_pretrained_discrete_latent_decoder(config, device)[source]#

Load the frozen latent-index-to-action path from a pretrained module.

protomotions.agents.common.discrete_latent.make_discrete_latent_target_encoder( module, *, tokenization, target_key, freeze=True, )[source]#

Adapt a module exposing encoder/quantizer to a target-label module.

protomotions.agents.common.discrete_latent.load_pretrained_discrete_latent_target_encoder( config, *, tokenization, target_key, device, )[source]#

Load the frozen observation-to-target-latent path from a pretrained module.