protomotions.agents.peft.prior_with_peft module#

PEFT adapter attached to a pretrained autoregressive GPC prior.

This module owns PEFT-specific state: adapter injection, task conditioning, and the frozen KL/sampling reference. The generic autoregressive token loop stays on the common prior transformer.

class protomotions.agents.peft.prior_with_peft.DiscretePriorWithPEFT(
prior,
conditioning_dim=None,
rank=4,
alpha=0.5,
peft_type='dora',
temperature=1.0,
top_p=0.8,
sampling_mode='nucleus',
prior_top_p=0.99,
condition_key='task_cond',
film_input_norm=False,
film_input_norm_clamp=5.0,
)[source]#

Bases: <Mock object at 0x7fd694b56c90>[]

Attach conditioned LoRA/DoRA layers to a frozen prior transformer.

The wrapper consumes the prior’s context keys plus one configured PEFT condition key. It builds task_c for the adapter layers and delegates token teacher-forcing/generation to DiscreteAutoregressiveTransformer. Higher-level actors own observation preprocessing, decoding, and target encoding.

__init__(
prior,
conditioning_dim=None,
rank=4,
alpha=0.5,
peft_type='dora',
temperature=1.0,
top_p=0.8,
sampling_mode='nucleus',
prior_top_p=0.99,
condition_key='task_cond',
film_input_norm=False,
film_input_norm_clamp=5.0,
)[source]#

Initialize PEFT wrapper.

Parameters:

  • prior – Base TransformerPrior model to wrap.

  • conditioning_dim (int | None) – Optional dimension of the full PEFT conditioning vector, i.e. condition_key features plus frozen-prior context observations. If omitted, it is inferred from warmup_obs in init_peft().

  • rank (int) – LoRA/DoRA rank.

  • alpha (float) – LoRA/DoRA scaling factor.

  • peft_type (str) – “lora” or “dora”.

  • temperature (float) – Sampling temperature for the PEFT actor.

  • top_p (float) – Nucleus sampling threshold for the PEFT actor.

  • sampling_mode (str) – Sampling strategy - “nucleus” or “prior_constraint”.

  • prior_top_p (float) – Nucleus threshold for the frozen prior (only used when sampling_mode=”prior_constraint”). Controls how conservatively the frozen prior caps the PEFT actor’s distribution.

  • condition_key (str) – TensorDict key carrying PEFT conditioning features.

  • film_input_norm (bool) – Normalize PEFT conditioning before FiLM layers.

  • film_input_norm_clamp (float) – Clamp value for normalized conditioning.

property task_c_dim#
init_peft(warmup_obs=None)[source]#

Materialize prior context shape, then install PEFT adapters.

warmup_obs should contain the raw observation keys expected by the frozen prior. It is only used when the pretrained prior’s context dimension has not already been materialized by a previous forward pass.

static reference_state_prefixes(prefix='')[source]#
optional_full_checkpoint_state_prefixes()[source]#

State that is optional when loading full PEFT training checkpoints.

property reference_ready: bool#
ensure_reference_modules()[source]#
mark_reference_loaded()[source]#
require_reference()[source]#
capture_reference()[source]#

Pin the complete KL/sampling reference to the current PEFT policy.

train(mode=True)[source]#
forward(input_dict)[source]#

Teacher-force the PEFT-wrapped prior and return token logits.