protomotions.agents.peft.utils.adapter_state module

State-dict filtering for PEFT adapter checkpoints.

High-level actor and agent code should talk in terms of “save adapter” and “load adapter”. This module owns the low-level key rules that decide which entries are adapter/task state and which entries belong to the frozen prior, critic, or training-only scaffolding.

An adapter checkpoint holds everything that cannot be rebuilt from the frozen prior checkpoint: the active (student) adapter and task state, plus the state of the RLFT reference policy. The reference policy is a frozen copy of the student networks: the adapter-wrapped prior, the FiLM input normalizer and the conditioning network. Its frozen transformer weights are rebuilt from the configured prior instead of stored.

protomotions.agents.peft.utils.adapter_state.strip_actor_prefix(key)

Accept both model-level keys and actor-local keys.

protomotions.agents.peft.utils.adapter_state.is_adapter_state_key(key)

Return True for the active adapter and task keys of an adapter PEFT checkpoint.

protomotions.agents.peft.utils.adapter_state.is_reference_state_key(key)

Return True for the reference-policy keys an adapter PEFT checkpoint stores.

These are the reference prior’s adapter tensors, the reference’s FiLM input normalizer and its conditioning network. The reference prior’s frozen weights are not included because they equal the configured frozen prior.

protomotions.agents.peft.utils.adapter_state.reference_state_from_policy_state(state_dict)

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

Accepts full and adapter checkpoints, with or without the _actor. prefix. The checkpoint’s own reference, its frozen prior weights and any critic or discriminator state are left out.

protomotions.agents.peft.utils.adapter_state.reference_lacks_conditioning_copy(state_dict)

True for a file saved before the reference kept its own conditioning network: it stores reference state and a student conditioning network, but no reference_peft_model.* keys. A partial copy is a key mismatch that validation reports, not such a file; a conditioning network that saves no state leaves nothing to copy, so a file without student network keys is not one either.

exception protomotions.agents.peft.utils.adapter_state.ReferenceWithoutConditioningError

Bases: RuntimeError

The refusal of a file saved before the reference kept its own conditioning network; the load entry points that know the file’s path re-raise it with the path in front.

protomotions.agents.peft.utils.adapter_state.refuse_reference_without_conditioning(state_dict)

Refuse a file saved before the reference kept its own conditioning network.

Such a reference read the student’s live conditioning network, so the student network saved in the same file is the one it read at save time; MIGRATION_SCRIPT copies it under the reference prefix into a new file. Nothing fills the copy in on load.

protomotions.agents.peft.utils.adapter_state.is_adapter_checkpoint_key(key)

Return True for any key that belongs in an adapter PEFT checkpoint.

protomotions.agents.peft.utils.adapter_state.contains_adapter_state(state_dict)

Return True when a checkpoint contains at least one adapter key.

protomotions.agents.peft.utils.adapter_state.build_adapter_state_dict(module)

Build an adapter-only state dict from a live actor module.

protomotions.agents.peft.utils.adapter_state.build_reference_state_dict(module)

Build the adapter reference-policy state from a live actor module.

The result is empty when no reference policy has been built yet.

protomotions.agents.peft.utils.adapter_state.validate_adapter_checkpoint(module, state_dict, *, strict=True)

Validate student and reference together, before either network is changed.

protomotions.agents.peft.utils.adapter_state.load_adapter_state_dict(module, state_dict, *, strict=True)

Load adapter/task keys into module without touching frozen prior state.

protomotions.agents.peft.utils.adapter_state.load_reference_state_dict(module, state_dict)

Restore the reference policy stored in an adapter checkpoint.

Returns False, and changes nothing, when state_dict stores no reference. Otherwise the reference prior is rebuilt from the module’s current frozen prior, then its adapter tensors, FiLM input normalizer and conditioning network are overwritten with the stored ones, so it matches the saved reference exactly.