protomotions.agents.peft.prior_config module

Configuration objects for discrete-token GPC prior PEFT.

SFT and RLFT get separate public configs instead of one config hidden behind a training_mode flag, so the config you instantiate determines which training loop runs. These configs assume the prior emits discrete autoregressive tokens (FSQ code indices), which is why the actor head is a token classifier and SFT supervises with cross-entropy over token indices rather than continuous action regression.

The public actor config mirrors the normal actor/critic style: actor.in_keys names the task observations, actor.peft.model builds the conditioning tensor, and actor.peft.condition_key names the tensor consumed by adapter layers. The frozen prior context is discovered from the loaded prior checkpoint.

protomotions.agents.peft.prior_config.default_peft_model_config(in_keys, condition_key='task_cond')

Build the default task-observation-to-PEFT-condition module.

class protomotions.agents.peft.prior_config.DiscretePriorPEFTConfig(

peft_type='dora',

rank=4,

alpha=0.5,

model=None,

condition_key='task_cond',

temperature=1.0,

top_p=0.8,

sampling_mode='nucleus',

prior_top_p=0.99,

kl_coeff=0.0,

clear_peft=False,

reference_checkpoint=None,

m_clamp=None,

film_input_norm=False,

film_input_norm_clamp=5.0,

)

Bases: TransformerPEFTConfig

Discrete-prior PEFT adapter and conditioning-network configuration.

model is the task-side TensorDict module that writes condition_key. Adapter layers receive that condition tensor plus the frozen prior context discovered from the checkpoint; they do not read task or terrain-specific observation keys directly. model must not write the frozen prior’s context keys: the frozen prior and the RLFT reference read those straight from the observations.

Attributes:

peft_type: ‘lora’ or ‘dora’. rank: Adapter rank. alpha: Adapter scaling factor. model: TensorDict module that writes PEFT conditioning features. It may read the frozen prior’s context keys but must not write them. If omitted, an ObsProcessor is built from actor in_keys; it has no trainable weights, but its running observation normalizer keeps updating during training. The RLFT reference policy keeps its own frozen copy of this module. condition_key: TensorDict key produced by model and consumed by PEFT. temperature: Sampling temperature during generation. top_p: Student nucleus sampling threshold. Ignored when sampling_mode=’prior_constraint’, where prior_top_p controls the reference policy’s nucleus. sampling_mode: Sampling strategy: ‘nucleus’ or ‘prior_constraint’. prior_top_p: Nucleus threshold for the reference policy when sampling_mode=’prior_constraint’. The reference is pinned at the start of a new run; a resume keeps it. kl_coeff: Coefficient of the KL penalty between the student and the reference policy during RLFT (0 = disabled). clear_peft: If True, zero the student’s adapter residuals once, right after a new RLFT run pins its reference. This lets RLFT keep an SFT policy as the reference while starting the student from the base prior. Never applies on resume or inference, and is skipped with a warning when nothing reads a reference (nucleus sampling with kl_coeff=0). reference_checkpoint: RLFT only. A full last.ckpt or adapter inference_last.ckpt from an SFT or RLFT run whose trained policy (student) becomes the reference for a new run, instead of the policy the run starts from. Its adapters, FiLM input normalizer and conditioning network are attached to the configured frozen prior; any reference stored in that file is ignored. Ignored on resume, which keeps its saved reference, and ignored with a warning when nothing reads a reference (nucleus sampling with kl_coeff=0). m_clamp: If set, clamp DoRA m parameters to [-m_clamp, m_clamp] after each actor optimizer step. film_input_norm: Normalize the full PEFT conditioning vector before FiLM gamma/beta layers. Its statistics keep updating during SFT and RLFT, and the RLFT reference policy keeps a frozen copy. Enable in both SFT and RLFT so the SFT-learned stats load into RLFT. film_input_norm_clamp: Clamp value for normalized PEFT conditioning.

model: ModuleContainerConfig | None = None

condition_key: str = 'task_cond'

temperature: float = 1.0

top_p: float = 0.8

sampling_mode: str = 'nucleus'

prior_top_p: float = 0.99

kl_coeff: float = 0.0

clear_peft: bool = False

reference_checkpoint: str | None = None

m_clamp: float | None = None

film_input_norm: bool = False

film_input_norm_clamp: float = 5.0

resolve_model(actor_in_keys)

__init__(

peft_type='dora',

rank=4,

alpha=0.5,

model=None,

condition_key='task_cond',

temperature=1.0,

top_p=0.8,

sampling_mode='nucleus',

prior_top_p=0.99,

kl_coeff=0.0,

clear_peft=False,

reference_checkpoint=None,

m_clamp=None,

film_input_norm=False,

film_input_norm_clamp=5.0,

)

class protomotions.agents.peft.prior_config.DiscretePriorPEFTActorConfig(

_target_='protomotions.agents.peft.actor.DiscretePriorPEFTActor',

in_keys=<factory>,

out_keys=<factory>,

peft=<factory>,

)

Bases: object

Task-facing actor config for a PEFT-adapted discrete GPC prior.

in_keys should contain only observations needed by peft.model. The pretrained prior’s context keys are resolved from the checkpoint at model construction time.

Attributes:

in_keys: Task-specific PEFT observation keys. The frozen prior context keys are discovered from the loaded pretrained prior and appended at runtime. out_keys: Actor rollout output keys.

in_keys: List[str]

out_keys: List[str]

peft: DiscretePriorPEFTConfig

__init__(

_target_='protomotions.agents.peft.actor.DiscretePriorPEFTActor',

in_keys=<factory>,

out_keys=<factory>,

peft=<factory>,

)

class protomotions.agents.peft.prior_config.DiscretePriorPEFTBaseModelConfig(

_target_='protomotions.agents.base_agent.model.BaseModel',

in_keys=<factory>,

out_keys=<factory>,

actor=<factory>,

actor_optimizer=<factory>,

)

Bases: BaseModelConfig

Shared actor config for discrete-prior PEFT models.

Attributes:

in_keys: Input keys. out_keys: Output keys.

actor: DiscretePriorPEFTActorConfig

actor_optimizer: OptimizerConfig

__init__(

_target_='protomotions.agents.base_agent.model.BaseModel',

in_keys=<factory>,

out_keys=<factory>,

actor=<factory>,

actor_optimizer=<factory>,

)

class protomotions.agents.peft.prior_config.DiscretePriorPEFTRLFTModelConfig(

_target_='protomotions.agents.peft.model.DiscretePriorPEFTModel',

in_keys=<factory>,

out_keys=<factory>,

actor=<factory>,

actor_optimizer=<factory>,

critic=None,

critic_optimizer=<factory>,

)

Bases: DiscretePriorPEFTBaseModelConfig

RLFT model config: discrete-prior PEFT actor plus task critic.

Attributes:

in_keys: Input keys. out_keys: Output keys. critic: Task critic config. RLFT requires this to be set.

critic: MLPWithConcatConfig | None = None

critic_optimizer: OptimizerConfig

__init__(

_target_='protomotions.agents.peft.model.DiscretePriorPEFTModel',

in_keys=<factory>,

out_keys=<factory>,

actor=<factory>,

actor_optimizer=<factory>,

critic=None,

critic_optimizer=<factory>,

)

class protomotions.agents.peft.prior_config.DiscretePriorPEFTSFTModelConfig(

_target_='protomotions.agents.peft.sft_model.DiscretePriorPEFTSFTModel',

in_keys=<factory>,

out_keys=<factory>,

actor=<factory>,

actor_optimizer=<factory>,

token_perturb_rate=0.0,

token_perturb_mode='replace',

)

Bases: DiscretePriorPEFTBaseModelConfig

SFT model config for supervised discrete-token labels.

Attributes:

in_keys: Input keys. out_keys: Output keys. token_perturb_rate: Probability of perturbing teacher-forced input tokens. token_perturb_mode: Token perturbation mode: replace, neighbor, or mixed.

token_perturb_rate: float = 0.0

token_perturb_mode: str = 'replace'

__init__(

_target_='protomotions.agents.peft.sft_model.DiscretePriorPEFTSFTModel',

in_keys=<factory>,

out_keys=<factory>,

actor=<factory>,

actor_optimizer=<factory>,

token_perturb_rate=0.0,

token_perturb_mode='replace',

)

class protomotions.agents.peft.prior_config.DiscretePriorPEFTRLFTAgentConfig(

batch_size,

training_max_steps,

training_max_iterations=None,

_target_='protomotions.agents.peft.prior_agent.DiscretePriorPEFTRLFTAgent',

model=<factory>,

num_steps=32,

gradient_clip_val=25.0,

fail_on_bad_grads=False,

check_grad_mag=True,

gamma=0.99,

bounds_loss_coef=0.0,

task_reward_w=1.0,

num_mini_epochs=2,

training_early_termination=None,

save_epoch_checkpoint_every=1000,

save_last_checkpoint_every=10,

save_inference_checkpoint=True,

max_episode_length_manager=None,

evaluator=<factory>,

normalize_rewards=True,

normalized_reward_clamp_value=5.0,

reward_norm_ema_decay=None,

tau=0.95,

e_clip=0.2,

clip_critic_loss=True,

actor_clip_frac_threshold=0.6,

entropy_coef=0.0,

l2c2=<factory>,

adaptive_lr=<factory>,

advantage_normalization=<factory>,

pretrained_modules=<factory>,

target_kl=None,

)

Bases: FineTuningAgentConfig

RLFT config for PPO fine-tuning of a discrete-prior PEFT adapter.

Attributes:

batch_size: Training batch size. training_max_steps: Maximum training steps. training_max_iterations: Maximum number of complete rollout and optimization iterations. When set, this takes precedence over training_max_steps. num_steps: Environment steps per update. fail_on_bad_grads: Fail on NaN/Inf gradients. check_grad_mag: Log gradient magnitude. gamma: Discount factor. bounds_loss_coef: Action bounds loss. 0 for tanh outputs. task_reward_w: Task reward weight. training_early_termination: Stop early at this step. None=disabled. save_epoch_checkpoint_every: Save epoch_xxx.ckpt every N epochs. save_last_checkpoint_every: Save last.ckpt every K epochs. max_episode_length_manager: Episode length curriculum. evaluator: Evaluation config. normalize_rewards: Normalize rewards. normalized_reward_clamp_value: Clamp normalized rewards to [-val, val]. reward_norm_ema_decay: EMA decay for reward normalization (None = Welford). Set to e.g. 0.99 to track non-stationary reward distributions. tau: GAE lambda for advantage estimation. e_clip: PPO clipping parameter epsilon. clip_critic_loss: Clip critic loss similar to actor. actor_clip_frac_threshold: Skip actor update if clip_frac > threshold (e.g., 0.5). entropy_coef: Entropy bonus coefficient added to PPO actor loss. Keep 0.0 for SFT-warmed PEFT unless deliberately exploring. l2c2: L2C2 settings. adaptive_lr: Adaptive learning rate settings. advantage_normalization: Advantage normalization settings. pretrained_modules: Frozen modules keyed by name. PEFT expects a whole prior model under ‘prior’. target_kl: If set, skip actor updates once minibatch actor/kl exceeds target_kl * 1.5 for the current rollout update.

pretrained_modules: Dict[str, PretrainedModelConfig]

model: DiscretePriorPEFTRLFTModelConfig

save_inference_checkpoint: bool = True

num_mini_epochs: int = 2

gradient_clip_val: float = 25.0

entropy_coef: float = 0.0

target_kl: float | None = None

__init__(

batch_size,

training_max_steps,

training_max_iterations=None,

_target_='protomotions.agents.peft.prior_agent.DiscretePriorPEFTRLFTAgent',

model=<factory>,

num_steps=32,

gradient_clip_val=25.0,

fail_on_bad_grads=False,

check_grad_mag=True,

gamma=0.99,

bounds_loss_coef=0.0,

task_reward_w=1.0,

num_mini_epochs=2,

training_early_termination=None,

save_epoch_checkpoint_every=1000,

save_last_checkpoint_every=10,

save_inference_checkpoint=True,

max_episode_length_manager=None,

evaluator=<factory>,

normalize_rewards=True,

normalized_reward_clamp_value=5.0,

reward_norm_ema_decay=None,

tau=0.95,

e_clip=0.2,

clip_critic_loss=True,

actor_clip_frac_threshold=0.6,

entropy_coef=0.0,

l2c2=<factory>,

adaptive_lr=<factory>,

advantage_normalization=<factory>,

pretrained_modules=<factory>,

target_kl=None,

)

class protomotions.agents.peft.prior_config.DiscretePriorPEFTSFTAgentConfig(

batch_size,

training_max_steps,

training_max_iterations=None,

_target_='protomotions.agents.peft.sft_agent.DiscretePriorPEFTSFTAgent',

model=<factory>,

num_steps=32,

gradient_clip_val=25.0,

fail_on_bad_grads=False,

check_grad_mag=True,

gamma=0.99,

bounds_loss_coef=0.0,

task_reward_w=1.0,

num_mini_epochs=2,

training_early_termination=None,

save_epoch_checkpoint_every=1000,

save_last_checkpoint_every=10,

save_inference_checkpoint=True,

max_episode_length_manager=None,

evaluator=<factory>,

normalize_rewards=True,

normalized_reward_clamp_value=5.0,

reward_norm_ema_decay=None,

expert_model_path=None,

rollout_actor=RolloutActor.EXPERT,

loss=<factory>,

expert_rollout_probability=0.0,

pretrained_modules=<factory>,

)

Bases: SupervisedAgentConfig

SFT config for supervised discrete-token PEFT training.

Attributes:

batch_size: Training batch size. training_max_steps: Maximum training steps. training_max_iterations: Maximum number of complete rollout and optimization iterations. When set, this takes precedence over training_max_steps. num_steps: Environment steps per update. fail_on_bad_grads: Fail on NaN/Inf gradients. check_grad_mag: Log gradient magnitude. gamma: Discount factor. bounds_loss_coef: Action bounds loss. 0 for tanh outputs. task_reward_w: Task reward weight. training_early_termination: Stop early at this step. None=disabled. save_epoch_checkpoint_every: Save epoch_xxx.ckpt every N epochs. save_last_checkpoint_every: Save last.ckpt every K epochs. max_episode_length_manager: Episode length curriculum. evaluator: Evaluation config. normalize_rewards: Normalize rewards. normalized_reward_clamp_value: Clamp normalized rewards to [-val, val]. reward_norm_ema_decay: EMA decay for reward normalization (None = Welford). Set to e.g. 0.99 to track non-stationary reward distributions. expert_model_path: External expert checkpoint that labels rollouts; requires rollout_actor=student. loss: Supervised loss over PEFT latent logits. expert_rollout_probability: Per-episode probability that the external expert drives the episode. pretrained_modules: Frozen modules keyed by name. SFT expects a whole prior model under ‘prior’.

pretrained_modules: Dict[str, PretrainedModelConfig]

model: DiscretePriorPEFTSFTModelConfig

rollout_actor: RolloutActor = 'expert'

loss: SupervisionLossConfig

save_inference_checkpoint: bool = True

num_mini_epochs: int = 2

gradient_clip_val: float = 25.0

__init__(

batch_size,

training_max_steps,

training_max_iterations=None,

_target_='protomotions.agents.peft.sft_agent.DiscretePriorPEFTSFTAgent',

model=<factory>,

num_steps=32,

gradient_clip_val=25.0,

fail_on_bad_grads=False,

check_grad_mag=True,

gamma=0.99,

bounds_loss_coef=0.0,

task_reward_w=1.0,

num_mini_epochs=2,

training_early_termination=None,

save_epoch_checkpoint_every=1000,

save_last_checkpoint_every=10,

save_inference_checkpoint=True,

max_episode_length_manager=None,

evaluator=<factory>,

normalize_rewards=True,

normalized_reward_clamp_value=5.0,

reward_norm_ema_decay=None,

expert_model_path=None,

rollout_actor=RolloutActor.EXPERT,

loss=<factory>,

expert_rollout_probability=0.0,

pretrained_modules=<factory>,

)