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,

)

Bases: <Mock object at 0x7a927c173e90>[]

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 and passed to the frozen prior as they are; the PEFT model may read them but must not write them. 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,

)

optional_full_checkpoint_state_prefixes()

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

The reference policy’s networks are present only once a reference is pinned or loaded. Their prefixes come from the student-to-reference table in adapter_state, which also drives checkpoint validation.

adapter_state_dict()

Return only PEFT adapter weights, excluding the frozen prior.

load_adapter_state_dict(

state_dict,

strict=True,

)

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

This loads the active adapter and task state. When the checkpoint also stores the RLFT reference policy, it is restored too, and the returned reference_loaded entry is True; resume and inference use it, while a new training run pins its own. Otherwise any earlier reference is dropped, so the new weights never run against a reference that came with other weights.

reference_state_dict()

Return the reference policy’s state for an adapter checkpoint.

That is its adapter tensors, its FiLM input normalizer and its copy of the conditioning network. Empty until a reference policy has been pinned. The reference’s frozen prior weights are left out because they equal the configured prior.

property reference_ready: bool

True once a reference policy has been pinned or loaded.

ensure_reference_modules()

Build the reference’s frozen networks if missing, without pinning them.

capture_reference()

Pin the reference to the current student policy, replacing any earlier one.

capture_reference_from(model_state)

Pin the reference to the trained (student) policy saved in a checkpoint.

model_state comes from a full or an adapter checkpoint, and any reference stored in it is ignored. The policy’s adapters, FiLM input normalizer and conditioning network are attached to this actor’s configured frozen prior, after their names and shapes are checked.

mark_reference_loaded()

Mark the reference state just loaded into the networks as ready to use.

reset_reference()

Drop the reference policy until one is loaded or pinned again.

Every load that restores no reference calls this, so new weights never keep an earlier load’s reference. SFT, which never uses a reference, calls it after every load, and RLFT training start calls it before a new run pins its own and whenever the run reads none.

init_peft(warmup_obs=None)

train(mode=True)

quantize(z)

fsq_codes_to_fsq_indices(codes)

fsq_indices_to_codes(indices)

fsq_indices_to_prior_tokens(fsq_indices)

prior_tokens_to_fsq_indices(prior_tokens)

one_hot_prior_tokens(prior_tokens)

perturb_tokens(tokens, *, rate, mode)

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)

Encode target poses to GPC prior tokens for SFT.

build_prior_input(tensordict, tokens=None)

forward(input_dict)

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

get_action_and_logp(tensordict)

Rollout step: generate action + per-token log-probs for PPO.