protomotions.agents.peft.prior_agent module#

RLFT agent for PEFT adapters on a frozen discrete-token GPC prior.

class protomotions.agents.peft.prior_agent.DiscretePriorPEFTRLFTAgent(fabric, env, config, root_dir=None)[source]#

Bases: DiscretePriorPEFTSetupMixin, FineTuningAgent

PPO/RLFT for a PEFT-adapted discrete GPC prior.

This class is intentionally RL-only: it owns PPO ratio loss, critic loss, RLFT checkpoint warm-start from SFT, and the frozen KL/sampling reference. Supervised PEFT training lives in sft_agent.py and uses the generic supervised loop.

__init__(
fabric,
env,
config,
root_dir=None,
)[source]#
load(
checkpoint,
load_env=True,
load_training_state=True,
*,
for_inference=False,
)[source]#
actor_step(batch_dict)[source]#

Compute actor loss and perform policy update.

Computes PPO clipped surrogate objective plus optional bounds loss and extra algorithm-specific losses.

Parameters:

batch_dict (Dict) – Minibatch containing obs, actions, old neglogp, advantages.

Returns:

  • actor_loss: Total actor loss for backprop

  • log_dict: Dictionary of actor metrics for logging

Return type:

Tuple of (actor_loss, log_dict) where

critic_step(batch_dict)[source]#

Critic MSE loss against computed returns.

perform_optimization_step(
batch_dict,
batch_idx,
)[source]#

Run one PEFT PPO minibatch update.