protomotions.agents.fine_tuning package#

Reusable fine-tuning agent utilities.

class protomotions.agents.fine_tuning.FineTuningAgent(fabric, env, config, root_dir=None)[source]#

Bases: PretrainedModulesMixin, PPO

PPO base for models that capture frozen pretrained modules during setup.

pretrained is populated before create_model(), consumed by the concrete model, then dropped after model reset. Surviving references live through the model that captured them. Subclasses should put adapter setup or other initialization that depends on those captured modules in _post_create_model_hook.

abstractmethod create_model()[source]#

Build the trainable model. self.pretrained is already populated.

class protomotions.agents.fine_tuning.FineTuningAgentConfig(
batch_size,
training_max_steps,
_target_='protomotions.agents.fine_tuning.agent.FineTuningAgent',
model=<factory>,
num_steps=32,
gradient_clip_val=0.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=1,
training_early_termination=None,
save_epoch_checkpoint_every=1000,
save_last_checkpoint_every=10,
save_inference_checkpoint=False,
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.005,
l2c2=<factory>,
adaptive_lr=<factory>,
advantage_normalization=<factory>,
pretrained_modules=<factory>,
)[source]#

Bases: PPOAgentConfig

Base config for agents that train on top of frozen pretrained modules.

Attributes:

batch_size: Training batch size. training_max_steps: Maximum training steps. model: Model config supplied by concrete fine-tuning agents. num_steps: Environment steps per update. gradient_clip_val: Max gradient norm. 0=disabled. 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. num_mini_epochs: Mini-epochs per update. 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. save_inference_checkpoint: Also save inference_<name>.ckpt without optimizer or other training-only state. 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 for learnable std exploration. l2c2: L2C2 settings. adaptive_lr: Adaptive learning rate settings. advantage_normalization: Advantage normalization settings. pretrained_modules: Frozen lower-stage modules keyed by name. Each entry is loaded before create_model() runs.

__init__(
batch_size,
training_max_steps,
_target_='protomotions.agents.fine_tuning.agent.FineTuningAgent',
model=<factory>,
num_steps=32,
gradient_clip_val=0.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=1,
training_early_termination=None,
save_epoch_checkpoint_every=1000,
save_last_checkpoint_every=10,
save_inference_checkpoint=False,
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.005,
l2c2=<factory>,
adaptive_lr=<factory>,
advantage_normalization=<factory>,
pretrained_modules=<factory>,
)#
model: BaseModelConfig#
pretrained_modules: Dict[str, PretrainedModelConfig]#
class protomotions.agents.fine_tuning.PretrainedModulesMixin[source]#

Bases: object

Load configured frozen modules through the BaseAgent setup lifecycle.

The mixin is intentionally independent of the training algorithm. PPO-style fine-tuning and supervised fine-tuning can both populate self.pretrained before create_model() and then drop that temporary handle after the model captures the modules it needs.

pretrained: Dict[str, <Mock object at 0x7fd69604b7d0>[]]#

Submodules#