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)

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,

)

load(

checkpoint,

load_env=True,

load_training_state=True,

)

fit()

Main training loop for the agent.

Executes the complete training process including: 1. Experience buffer setup (auto-registers keys from model outputs) 2. Environment rollouts (data collection) 3. Model optimization 4. Periodic evaluation 5. Checkpoint saving 6. Logging and metrics

The loop runs for max_epochs epochs, where each epoch collects num_steps of experience from num_envs parallel environments, then performs multiple optimization steps on the collected data.

Note

This is the main entry point for training. Call setup() before fit().

actor_step(batch_dict)

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)

Critic MSE loss against computed returns.

perform_optimization_step(

batch_dict,

batch_idx,

)

Run one PEFT PPO minibatch update.