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,

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.

Attributes:

peft_type: ‘lora’ or ‘dora’. rank: Adapter rank. alpha: Adapter scaling factor. model: TensorDict module that writes PEFT conditioning features. If omitted, a parameter-free ObsProcessor is built from actor in_keys. 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 frozen-prior nucleus. sampling_mode: Sampling strategy: ‘nucleus’ or ‘prior_constraint’. prior_top_p: Nucleus threshold for the frozen prior when sampling_mode=’prior_constraint’. kl_coeff: KL divergence loss coefficient for training (0 = disabled). clear_peft: If True, zero adapter residuals once after the RL reference is pinned. This lets RLFT keep an SFT prior constraint while starting the active student from the base prior. 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. 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

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,

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,

_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. 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,

_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,

_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>,

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. 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: Optional checkpoint for an external expert policy. loss: Supervised loss over PEFT latent logits. 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,

_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>,

pretrained_modules=<factory>,

)