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.

The reference is a frozen copy of the policy a training run starts from. This module holds its prior and FiLM input normalizer; the actor holds its copy of the conditioning network, and the RLFT agent decides when to pin it.

protomotions.agents.peft.prior_with_peft.frozen_copy(module)[source]#

Return a frozen copy of module for the reference policy.

The copy is in eval mode, takes no gradients and never updates its normalizers. Its running statistics share the live Lightning Fabric instead of copying it: a Fabric can hold process handles that cannot be copied, and frozen statistics never synchronize anyway.

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 0x7a9277a2ca10>[]

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 reference policy (only used when sampling_mode=”prior_constraint”). Controls how conservatively the reference 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.

optional_full_checkpoint_state_prefixes()[source]#

State that is optional when loading full PEFT training checkpoints.

The reference prior and FiLM input normalizer held here are declared by the actor, which lists every reference network in one place.

property reference_ready: bool#
ensure_reference_modules()[source]#

Build the reference prior and FiLM input normalizer if missing.

Inside an actor, call DiscretePriorPEFTActor.ensure_reference_modules, which also builds the reference’s copy of the conditioning network.

mark_reference_loaded()[source]#

Mark the reference state just loaded as ready to use.

Inside an actor, call DiscretePriorPEFTActor.mark_reference_loaded.

reset_reference()[source]#

Drop the reference prior and FiLM input normalizer.

Inside an actor, call DiscretePriorPEFTActor.reset_reference, which also drops the reference’s copy of the conditioning network. Called here alone, the copy stays behind, and a last.ckpt saved afterwards fails to load.

require_reference()[source]#
capture_reference()[source]#

Pin the KL/sampling reference to the current adapters and FiLM normalizer.

This replaces any reference pinned or loaded earlier. Inside an actor, call DiscretePriorPEFTActor.capture_reference, which also pins the reference’s copy of the conditioning network.

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

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