protomotions.agents.peft.actor module#

Discrete-prior PEFT actor built around a pretrained discrete latent prior.

This module owns observation routing, FSQ encode/decode helpers, and actor rollout outputs. Token-level adapter generation lives in prior_with_peft; the pretrained prior model itself lives in the Supervised latent-prior package.

class protomotions.agents.peft.actor.DiscretePriorPEFTActor( config, pretrained_prior_model, mimic_target_poses_dim=0, )[source]#

Bases: <Mock object at 0x7fd694b1b9d0>[]

Actor that wraps DiscretePriorWithPEFT and the frozen decoder for action generation.

The actor owns the public observation contract. Task observations flow through config.peft.model into config.peft.condition_key; frozen prior context keys are discovered from the pretrained prior and appended to in_keys internally. The target encoder is optional and is used only for SFT batches that include mimic_target_poses.

__init__( config, pretrained_prior_model, mimic_target_poses_dim=0, )[source]#

optional_full_checkpoint_state_prefixes()[source]#

Frozen target encoder state is present only for SFT/checkpointing flows.

adapter_state_dict()[source]#

Return only PEFT adapter weights, excluding the frozen prior.

load_adapter_state_dict( state_dict, strict=True, )[source]#

Load adapter-only state, ignoring frozen-prior and critic entries.

init_peft(warmup_obs=None)[source]#

train(mode=True)[source]#

quantize(z)[source]#

fsq_codes_to_fsq_indices(codes)[source]#

fsq_indices_to_codes(indices)[source]#

fsq_indices_to_prior_tokens(fsq_indices)[source]#

prior_tokens_to_fsq_indices(prior_tokens)[source]#

one_hot_prior_tokens(prior_tokens)[source]#

perturb_tokens(tokens, *, rate, mode)[source]#

Apply SFT token-noise augmentation to GPC prior tokens.

Teacher forcing on perturbed target tokens reduces exact-sequence memorization and makes the adapter less brittle when RLFT samples drift off the expert trajectory.

predict_target_prior_tokens(tensordict)[source]#: Encode target poses to GPC prior tokens for SFT.

build_prior_input(tensordict, tokens=None)[source]#

forward(input_dict)[source]#

Teacher-forced forward -> logits (B, num_prior_tokens, prior_token_vocab_size).

get_action_and_logp(tensordict)[source]#: Rollout step: generate action + per-token log-probs for PPO.