# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""RLFT agent for PEFT adapters on a frozen discrete-token GPC prior."""
import logging
from pathlib import Path
from typing import Dict, Optional, Tuple
from lightning.fabric import Fabric
import torch
from torch import Tensor
from tensordict import TensorDict
from protomotions.agents.common.autoregressive import (
kl_divergence_sampling_distribution,
)
from protomotions.agents.common.latent import compute_discrete_latent_ppo_loss
from protomotions.agents.fine_tuning.agent import FineTuningAgent
from protomotions.agents.peft.prior_setup import DiscretePriorPEFTSetupMixin
from protomotions.agents.peft.utils.adapter_state import (
ReferenceWithoutConditioningError,
)
from protomotions.agents.peft.utils.model_state import has_reference_state
log = logging.getLogger(__name__)
[docs]
class DiscretePriorPEFTRLFTAgent(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.
"""
[docs]
def __init__(self, fabric: Fabric, env, config, root_dir: Optional[Path] = None):
if config.model.critic is None:
raise ValueError(
"DiscretePriorPEFTRLFTAgent requires config.model.critic for RLFT."
)
super().__init__(fabric, env, config, root_dir=root_dir)
[docs]
def load(
self, checkpoint, load_env=True, load_training_state: bool = True,
*, for_inference: bool = False,
):
self._peft_loading_training_state = load_training_state
self._peft_resumed = False
try:
super().load(
checkpoint,
load_env=load_env,
load_training_state=load_training_state,
for_inference=for_inference,
)
except ReferenceWithoutConditioningError as error:
# The loaders see only the state dict; name the file here.
raise ReferenceWithoutConditioningError(f"{checkpoint}: {error}") from None
finally:
self._peft_loading_training_state = False
# Fit start needs to know: a resume keeps the reference saved in its
# checkpoint and never clears the adapters, while a new run pins its own.
self._peft_resumed = checkpoint is not None and load_training_state
def _validate_checkpoint(self, state_dict, *, load_training_state, for_inference):
super()._validate_checkpoint(
state_dict, load_training_state=load_training_state,
for_inference=for_inference,
)
peft = self.model._actor.prior_with_peft
if load_training_state and self._reads_reference():
if not has_reference_state(state_dict["model"]):
raise RuntimeError(
"This checkpoint has no saved reference policy, which is required "
"to resume reference-dependent training. Use a "
"checkpoint containing the reference; new RLFT training may instead "
"warm-start weights with load_training_state=False."
)
# Only RLFT samples against the reference. SFT inference replays its
# expert's tokens, so the SFT agent can load a file without one.
if for_inference and peft.sampling_mode == "prior_constraint":
if not has_reference_state(state_dict["model"]):
raise RuntimeError(
"This checkpoint has no saved reference policy, and "
"prior_constraint sampling needs one: it limits the policy to "
"tokens inside the reference policy's top-p set. Load "
"this run's last.ckpt instead if it was saved with a reference, "
"or sample without the constraint by running inference with "
"--overrides agent.model.actor.peft.sampling_mode=nucleus "
"(agent.model.actor.peft.top_p then sets the nucleus)."
)
def _before_fit(self):
self._prepare_rlft_prior_reference()
def _reads_reference(self) -> bool:
"""True when sampling or the KL penalty compares against the reference."""
actor = self.model._actor
return (
actor.prior_with_peft.sampling_mode == "prior_constraint"
or actor.kl_coeff > 0
)
def _prepare_rlft_prior_reference(self):
"""Pin the reference policy at training start, after all weight loading.
The reference is a frozen copy of the policy a run starts from. A new
run, fresh or warm-started from any checkpoint, ignores any reference
stored in that checkpoint: it pins the policy it just loaded, or the
trained policy in ``reference_checkpoint`` if that is set, and only
then applies ``clear_peft`` to the student. A resume keeps the reference
saved in its checkpoint, which load() already required, and never
clears the adapters. A run that never reads a reference holds none, so
it saves none either. Inference never reaches this, so it uses the
saved reference as-is.
The student's FiLM input normalizer and the conditioning network's keep
updating during training. The frozen prior's normalizers were frozen at
setup, and only the reference's own copies are frozen.
"""
actor = self.model._actor
peft_cfg = self.config.model.actor.peft
resumed = getattr(self, "_peft_resumed", False)
if not self._reads_reference():
# Never hold, or save again, a reference loaded from another run.
actor.reset_reference()
if not resumed:
self._warn_unused_reference_options()
return
if resumed:
log.info("Resumed run keeps the reference policy saved in its checkpoint.")
return
if actor.reference_ready:
log.info(
"New run: the reference stored in the loaded checkpoint is not "
"used; this run pins its own."
)
actor.reset_reference()
if peft_cfg.reference_checkpoint:
self._pin_reference_from_checkpoint(peft_cfg.reference_checkpoint)
else:
actor.capture_reference()
if peft_cfg.clear_peft:
actor.prior_with_peft.clear_peft()
def _pin_reference_from_checkpoint(self, path):
"""Pin the reference to the trained policy saved in another checkpoint."""
try:
state = torch.load(path, map_location=self.device, weights_only=False)
model_state = (
state["model"] if isinstance(state, dict) and "model" in state else state
)
self.model._actor.capture_reference_from(model_state)
except (OSError, RuntimeError, KeyError) as error:
raise RuntimeError(
f"agent.model.actor.peft.reference_checkpoint={path} cannot be used "
f"as the reference policy: {error}"
) from error
log.info("Pinned the reference to the trained policy in %s.", path)
def _warn_unused_reference_options(self):
"""Warn about configured options that only act on a pinned reference."""
peft_cfg = self.config.model.actor.peft
sampling_mode = self.model._actor.prior_with_peft.sampling_mode
if peft_cfg.clear_peft:
log.warning(
"clear_peft=True was skipped. With sampling_mode=%r and "
"kl_coeff=0 this run pins no reference policy, and clear_peft "
"only resets the adapters once a pinned reference keeps the "
"loaded policy; without one it would just discard that policy. "
"prior_constraint sampling or kl_coeff > 0 pins the loaded "
"policy as the reference and then clears the adapters. "
"Otherwise drop clear_peft.",
sampling_mode,
)
if peft_cfg.reference_checkpoint:
log.warning(
"reference_checkpoint=%r was ignored: with sampling_mode=%r and "
"kl_coeff=0 this run pins no reference policy.",
peft_cfg.reference_checkpoint,
sampling_mode,
)
@torch.no_grad()
def _clamp_peft_m(self):
# DoRA magnitude m is unbounded by construction. Clamping keeps RLFT
# adapter updates near the frozen prior's scale instead of letting a few
# large magnitudes dominate the token logits.
bound = getattr(self.config.model.actor.peft, "m_clamp", None)
if bound is None:
return
peft = self.model._actor.prior_with_peft
for module in peft.base_prior._transformer.modules():
if hasattr(module, "m"):
module.m.clamp_(-bound, bound)
[docs]
def actor_step(self, batch_dict: Dict) -> Tuple[Tensor, Dict]:
return self._actor_step_discrete_ppo(batch_dict)
def _actor_step_discrete_ppo(self, batch_dict: Dict) -> Tuple[Tensor, Dict]:
"""Compute the discrete-token PPO loss for the PEFT actor."""
actor = self.model._actor
prior_tokens = batch_dict["prior_tokens"].detach()
old_neglogp = batch_dict["neglogp"].detach()
advantages = batch_dict["advantages"].detach()
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
# PPO is applied in token space: rollout sampled GPC prior tokens from
# the adapter, then optimization replays those exact tokens under
# the current adapter to compute a categorical likelihood ratio.
prior_dict = actor.build_prior_input(batch_dict, tokens=prior_tokens)
logits = self.actor(prior_dict)
prior_logits = None
prior_with_peft = getattr(actor, "prior_with_peft", None)
loss_temperature = getattr(prior_with_peft, "temperature", 1.0)
loss_top_p = getattr(prior_with_peft, "top_p", 1.0)
if (
prior_with_peft is not None
and getattr(prior_with_peft, "sampling_mode", None) == "prior_constraint"
):
prior_logits = prior_with_peft.forward_prior(prior_dict)
loss_top_p = prior_with_peft.prior_top_p
ppo_loss, ppo_log_dict = compute_discrete_latent_ppo_loss(
logits=logits,
selected=prior_tokens,
old_neglogp=old_neglogp,
advantages=advantages,
e_clip=self.e_clip,
entropy_coef=self.config.entropy_coef,
temperature=loss_temperature,
top_p=loss_top_p,
prior_logits=prior_logits,
log_prefix="actor",
)
kl_coeff = actor.kl_coeff
kl_prior_loss = torch.tensor(0.0, device=logits.device)
if kl_coeff > 0:
# Compare the same transformed distribution used to sample PPO
# actions, not raw full-vocabulary logits.
if prior_logits is None:
prior_logits = actor.prior_with_peft.forward_prior(prior_dict)
kl_prior_loss = kl_divergence_sampling_distribution(
logits,
prior_logits,
p=loss_top_p,
temperature=loss_temperature,
prior_constraint=(
prior_with_peft is not None
and getattr(prior_with_peft, "sampling_mode", None)
== "prior_constraint"
),
reduction="mean",
)
loss = ppo_loss + kl_coeff * kl_prior_loss
log_dict = {
"actor/kl_prior_loss": kl_prior_loss.detach(),
"actor/kl_coeff": kl_coeff,
"actor/adv_mean": advantages.mean().detach(),
"actor/adv_std": advantages.std().detach(),
"losses/actor_loss": loss.detach(),
"stats/reward_mean": batch_dict["rewards"].mean().detach(),
}
log_dict.update(ppo_log_dict)
return loss, log_dict
[docs]
def critic_step(self, batch_dict: Dict) -> Tuple[Tensor, Dict]:
"""Critic MSE loss against computed returns."""
batch_td = TensorDict(
batch_dict,
batch_size=batch_dict["returns"].shape[0],
)
batch_td = self.critic(batch_td)
out_key = self.model._critic.out_keys[0]
values = batch_td[out_key].squeeze(-1)
returns = batch_dict["returns"].detach()
critic_loss = torch.nn.functional.mse_loss(values, returns)
log_dict = {
"losses/critic_loss": critic_loss.detach(),
"stats/value_mean": values.mean().detach(),
"stats/return_mean": returns.mean().detach(),
}
return critic_loss, log_dict
def _load_training_state(self, state_dict):
warm_start_from_sft = self.has_critic and "critic_optimizer" not in state_dict
if not warm_start_from_sft:
super()._load_training_state(state_dict)
return
log.info(
"Using SFT checkpoint as RLFT initialization; restoring actor "
"optimizer state and leaving critic optimizer fresh."
)
if "actor_optimizer" in state_dict:
self.actor_optimizer.load_state_dict(state_dict["actor_optimizer"])
self.current_epoch = 0
self.step_count = 0
self.fit_start_time = None
self.best_evaluated_score = None
