protomotions.agents.common.autoregressive module#

Common autoregressive modules for categorical token sequences.

protomotions.agents.common.autoregressive.sampling_log_probs(logits, p=0.9, temperature=1.0)[source]#

Log-probabilities of the actual temperature/top-p sampling distribution.

protomotions.agents.common.autoregressive.prior_constrained_sampling_log_probs( logits, prior_logits, p=0.99, temperature=1.0, overlap_threshold=0.001, )[source]#

Log-probs after constraining model support to the prior top-p nucleus.

protomotions.agents.common.autoregressive.nucleus_sampling(logits, p=0.9, temperature=1.0)[source]#

Sample categorical indices from the top-p nucleus.

protomotions.agents.common.autoregressive.nucleus_sampling_prior_constraint( logits, prior_logits, p=0.99, temperature=1.0, overlap_threshold=0.001, )[source]#

Sample from logits while restricting support to the prior top-p nucleus.

If the model assigns effectively no probability mass to the prior nucleus, sample from the prior nucleus instead of falling back to unconstrained model probabilities. That keeps prior-constraint mode active after policy drift.

protomotions.agents.common.autoregressive.kl_divergence_categorical(logits, prior_logits, reduction='mean')[source]#

KL divergence between categorical distributions parameterized by logits.

protomotions.agents.common.autoregressive.kl_divergence_from_log_probs(log_p, log_q, reduction='mean')[source]#

KL divergence for already transformed categorical log-probabilities.

protomotions.agents.common.autoregressive.kl_divergence_sampling_distribution( logits, prior_logits, *, p=0.9, temperature=1.0, prior_constraint=False, reduction='mean', )[source]#

KL between the actual transformed token sampling distributions.

protomotions.agents.common.autoregressive.generate_causal_mask(num_target, num_context=0, device=None)[source]#

Build a float causal attention mask with an optional context prefix.

protomotions.agents.common.autoregressive.resolve_discrete_autoregressive_config( config, *, num_tokens, vocab_size, )[source]#

Return a copy of config with token count and vocabulary resolved.

class protomotions.agents.common.autoregressive.DiscreteAutoregressiveTransformer(*args, **kwargs)[source]#

Bases: ProtoMotionsTensorDictModule

Categorical autoregressive transformer with configurable projections.

__init__(config)[source]#

config: DiscreteAutoregressiveTransformerConfig#

muon_adam_fallback_modules()[source]#: Categorical input/output projections should use the auxiliary Adam path.

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

encode_context(tensordict)[source]#

encode_tokens(tokens)[source]#

decode_logits(hidden)[source]#

add_positions( sequence, pos_emb=None, )[source]#

run_transformer( sequence, num_target, transformer=None, transformer_kwargs=None, )[source]#

forward_from_tokens( context, tokens, *, transformer=None, pos_emb=None, transformer_kwargs=None, )[source]#

teacher_force(tensordict)[source]#

forward(tensordict)[source]#

Run teacher forcing when tokens are supplied, otherwise generate tokens.

next_logits_from_context( context, token_indices=None, *, transformer=None, pos_emb=None, transformer_kwargs=None, )[source]#

Return logits for the next token after an optional prefix.

generate_from_context( context, num_tokens, *, temperature=1.0, top_p=0.9, prior_constraint=None, transformer=None, pos_emb=None, transformer_kwargs=None, )[source]#