# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""Discrete-prior PEFT actor built around a pretrained discrete latent prior.
This module owns observation routing, FSQ encode/decode helpers, and actor
rollout outputs. Token-level adapter generation lives in ``prior_with_peft``;
the pretrained prior model itself lives in the Supervised latent-prior package.
"""
from __future__ import annotations
import logging
import torch
from torch import nn
from tensordict import TensorDict
from protomotions.agents.common.discrete_latent import (
DiscreteLatentDecoder,
DiscreteLatentTargetEncoder,
FSQTokenization,
)
from protomotions.agents.common.latent import (
LATENT_KEY,
LATENT_LOGPROB_KEY,
)
from protomotions.agents.common.pretrained import freeze_module
from protomotions.agents.peft.utils.adapter_state import (
REFERENCE_STATE_PREFIXES,
build_adapter_state_dict,
build_reference_state_dict,
reference_state_from_policy_state,
refuse_reference_without_conditioning,
validate_adapter_checkpoint,
load_adapter_state_dict as load_adapter_state,
load_reference_state_dict as load_reference_state,
)
from protomotions.agents.peft.utils.frozen_prior_contract import (
require_frozen_prior_attr,
resolve_frozen_prior_input_keys,
)
from protomotions.agents.peft.prior_with_peft import (
DiscretePriorWithPEFT,
frozen_copy,
)
from protomotions.utils.hydra_replacement import get_class
log = logging.getLogger(__name__)
[docs]
class DiscretePriorPEFTActor(nn.Module):
"""Actor that wraps DiscretePriorWithPEFT and the frozen decoder for action generation.
The actor owns the public observation contract. Task observations flow
through ``config.peft.model`` into ``config.peft.condition_key``; frozen
prior context keys are discovered from the pretrained prior and appended to
``in_keys`` internally and passed to the frozen prior as they are; the PEFT
model may read them but must not write them. The target encoder is optional
and is used only for SFT batches that include ``mimic_target_poses``.
"""
[docs]
def __init__(
self,
config,
pretrained_prior_model,
mimic_target_poses_dim: int = 0,
):
super().__init__()
peft_cfg = config.peft
self.out_keys = list(config.out_keys)
latent_decoder = require_frozen_prior_attr(
pretrained_prior_model,
"latent_decoder",
DiscreteLatentDecoder,
)
latent_tokenization = require_frozen_prior_attr(
pretrained_prior_model,
"latent_tokenization",
FSQTokenization,
)
self.frozen_prior_input_keys = resolve_frozen_prior_input_keys(
pretrained_prior_model
)
prior_transformer = pretrained_prior_model.prior
self.condition_key = peft_cfg.condition_key
actor_model_config = peft_cfg.model
if actor_model_config is None:
raise AssertionError(
"DiscretePriorPEFTActor requires config.peft.model. "
"DiscretePriorPEFTActorConfig should resolve the default model."
)
ActorPEFTModelClass = get_class(actor_model_config._target_)
self.actor_peft_model = ActorPEFTModelClass(config=actor_model_config)
# The reference policy's frozen copy of it, built when a reference is
# pinned or loaded.
self.reference_peft_model = None
if self.condition_key not in self.actor_peft_model.out_keys:
raise AssertionError(
"DiscretePriorPEFTActor PEFT model must produce condition_key "
f"{self.condition_key!r} for adapter "
f"conditioning. out_keys={self.actor_peft_model.out_keys}"
)
# The frozen prior, and so the RLFT reference, read their context
# straight from the observations they were trained on.
rewritten_context = [
key
for key in self.actor_peft_model.out_keys
if key in self.frozen_prior_input_keys
]
if rewritten_context:
raise AssertionError(
"DiscretePriorPEFTActor PEFT model must write only conditioning "
f"features, not the frozen prior's context keys {rewritten_context}. "
"The frozen prior reads those straight from the observations."
)
missing_actor_inputs = [
key for key in self.actor_peft_model.in_keys if key not in config.in_keys
]
if missing_actor_inputs:
raise AssertionError(
"DiscretePriorPEFTActor in_keys must include actor model in_keys. "
f"Missing: {missing_actor_inputs}"
)
self.in_keys = list(
dict.fromkeys(
[
*config.in_keys,
*self.actor_peft_model.in_keys,
*self.frozen_prior_input_keys,
]
)
)
log.info(
"[PEFT] Loaded frozen prior expects actor input keys=%s; "
"actor.peft.model.in_keys=%s; actor.peft.model.out_keys=%s",
tuple(self.frozen_prior_input_keys),
tuple(self.actor_peft_model.in_keys),
tuple(self.actor_peft_model.out_keys),
)
# Frozen decoder from the current discrete latent prior model.
self.latent_decoder = latent_decoder
self.latent_tokenization = latent_tokenization
self.decoder_latent_key = latent_decoder.latent_key
freeze_module(self.latent_decoder)
# Frozen encoder from prior (only needed for SFT with mimic_target_poses).
if mimic_target_poses_dim > 0:
target_latent_encoder = require_frozen_prior_attr(
pretrained_prior_model,
"target_latent_encoder",
DiscreteLatentTargetEncoder,
)
self.target_latent_encoder = target_latent_encoder
freeze_module(self.target_latent_encoder)
else:
self.target_latent_encoder = None
# Frozen prior transformer, wrapped with PEFT adapters below.
freeze_module(prior_transformer)
# FSQ scalar-code counts stay separate from autoregressive prior tokens.
self.num_fsq_levels = latent_decoder.num_fsq_levels
self.num_fsq_scalars = latent_decoder.num_fsq_scalars
self.fsq_scalars_per_prior_token = (
latent_tokenization.fsq_scalars_per_prior_token
)
self.num_prior_tokens = latent_tokenization.num_prior_tokens
self.prior_token_vocab_size = latent_tokenization.prior_token_vocab_size
self.L = latent_decoder.quantizer.L
self.half_L = latent_decoder.quantizer.half_L
self.half_width = latent_decoder.quantizer.half_width
# PEFT-wrapped prior
self.prior_with_peft = DiscretePriorWithPEFT(
prior=prior_transformer,
rank=peft_cfg.rank,
alpha=peft_cfg.alpha,
peft_type=peft_cfg.peft_type,
temperature=peft_cfg.temperature,
top_p=peft_cfg.top_p,
sampling_mode=peft_cfg.sampling_mode,
prior_top_p=peft_cfg.prior_top_p,
condition_key=self.condition_key,
film_input_norm=peft_cfg.film_input_norm,
film_input_norm_clamp=peft_cfg.film_input_norm_clamp,
)
self.kl_coeff = peft_cfg.kl_coeff
[docs]
def optional_full_checkpoint_state_prefixes(self) -> tuple[str, ...]:
"""Frozen target encoder state is present only for SFT/checkpointing flows.
The reference policy's networks are present only once a reference is
pinned or loaded. Their prefixes come from the student-to-reference
table in ``adapter_state``, which also drives checkpoint validation.
"""
return ("target_latent_encoder.", *REFERENCE_STATE_PREFIXES)
[docs]
def adapter_state_dict(self) -> dict[str, torch.Tensor]:
"""Return only PEFT adapter weights, excluding the frozen prior."""
return build_adapter_state_dict(self)
[docs]
def load_adapter_state_dict(self, state_dict: dict, strict: bool = True):
"""Load adapter-checkpoint state, ignoring frozen-prior and critic entries.
This loads the active adapter and task state. When the checkpoint also
stores the RLFT reference policy, it is restored too, and the returned
``reference_loaded`` entry is True; resume and inference use it, while
a new training run pins its own. Otherwise any earlier reference is
dropped, so the new weights never run against a reference that came
with other weights.
"""
refuse_reference_without_conditioning(state_dict)
validate_adapter_checkpoint(self, state_dict, strict=strict)
result = load_adapter_state(self, state_dict, strict=strict)
result["reference_loaded"] = load_reference_state(self, state_dict)
if not result["reference_loaded"]:
self.reset_reference()
return result
[docs]
def reference_state_dict(self) -> dict[str, torch.Tensor]:
"""Return the reference policy's state for an adapter checkpoint.
That is its adapter tensors, its FiLM input normalizer and its copy of
the conditioning network. Empty until a reference policy has been
pinned. The reference's frozen prior weights are left out because they
equal the configured prior.
"""
if not self.reference_ready:
return {}
return build_reference_state_dict(self)
# ---- reference policy ----
# Change the reference through these methods, not through the matching
# methods of prior_with_peft: only these also handle the reference's copy
# of the conditioning network.
@property
def reference_ready(self) -> bool:
"""True once a reference policy has been pinned or loaded."""
return self.prior_with_peft.reference_ready
[docs]
def ensure_reference_modules(self) -> None:
"""Build the reference's frozen networks if missing, without pinning them."""
self.prior_with_peft.ensure_reference_modules()
if self.reference_peft_model is None:
self.reference_peft_model = frozen_copy(self.actor_peft_model)
[docs]
def capture_reference(self) -> None:
"""Pin the reference to the current student policy, replacing any earlier one."""
self.ensure_reference_modules()
self.prior_with_peft.capture_reference()
self.reference_peft_model.load_state_dict(self.actor_peft_model.state_dict())
[docs]
def capture_reference_from(self, model_state: dict) -> None:
"""Pin the reference to the trained (student) policy saved in a checkpoint.
``model_state`` comes from a full or an adapter checkpoint, and any
reference stored in it is ignored. The policy's adapters, FiLM input
normalizer and conditioning network are attached to this actor's
configured frozen prior, after their names and shapes are checked.
"""
reference_state = reference_state_from_policy_state(model_state)
if not reference_state:
raise RuntimeError(
"This checkpoint holds no PEFT adapter or conditioning-network "
"state to use as the reference policy."
)
load_reference_state(self, reference_state)
[docs]
def mark_reference_loaded(self) -> None:
"""Mark the reference state just loaded into the networks as ready to use."""
self.prior_with_peft.mark_reference_loaded()
[docs]
def reset_reference(self) -> None:
"""Drop the reference policy until one is loaded or pinned again.
Every load that restores no reference calls this, so new weights never
keep an earlier load's reference. SFT, which never uses a reference,
calls it after every load, and RLFT training start calls it before a
new run pins its own and whenever the run reads none.
"""
self.prior_with_peft.reset_reference()
self.reference_peft_model = None
[docs]
def init_peft(self, warmup_obs: dict | None = None):
if warmup_obs is not None:
warmup_obs = self.build_prior_input(warmup_obs)
self.prior_with_peft.init_peft(warmup_obs=warmup_obs)
[docs]
def train(self, mode: bool = True):
super().train(mode)
if self.target_latent_encoder is not None:
self.target_latent_encoder.eval()
self.target_latent_encoder.encoder.eval()
self.latent_decoder.eval()
self.latent_decoder.decoder.eval()
self.latent_decoder.quantizer.eval()
self.prior_with_peft.base_prior.eval()
self.prior_with_peft.train(mode)
if self.reference_peft_model is not None:
self.reference_peft_model.eval()
return self
# ---- FSQ utilities ----
[docs]
def quantize(self, z):
return self.latent_decoder.quantizer.quantize(z)
[docs]
def fsq_codes_to_fsq_indices(self, codes):
return self.latent_decoder.quantizer.codes_to_indices(codes)
[docs]
def fsq_indices_to_codes(self, indices):
return self.latent_decoder.quantizer.indices_to_codes(indices)
[docs]
def fsq_indices_to_prior_tokens(self, fsq_indices):
return self.latent_tokenization.fsq_indices_to_prior_tokens(fsq_indices)
[docs]
def prior_tokens_to_fsq_indices(self, prior_tokens):
return self.latent_tokenization.prior_tokens_to_fsq_indices(prior_tokens)
[docs]
def one_hot_prior_tokens(self, prior_tokens):
return self.latent_tokenization.one_hot_prior_tokens(prior_tokens)
[docs]
def perturb_tokens(
self,
tokens: torch.Tensor,
*,
rate: float,
mode: str,
) -> torch.Tensor:
"""Apply SFT token-noise augmentation to GPC prior tokens.
Teacher forcing on perturbed target tokens reduces exact-sequence
memorization and makes the adapter less brittle when RLFT samples drift
off the expert trajectory.
"""
if rate <= 0:
return tokens
b, n = tokens.shape
device = tokens.device
mask = torch.rand(b, n, device=device) < rate
if not mask.any():
return tokens
result = tokens.clone()
if mode == "replace" or mode == "mixed":
rand_tokens = torch.randint(
0,
self.prior_token_vocab_size,
(b, n),
device=device,
dtype=tokens.dtype,
)
if mode == "replace":
result = torch.where(mask, rand_tokens, result)
else:
use_replace = torch.rand(b, n, device=device) < 0.2
result = torch.where(mask & use_replace, rand_tokens, result)
neighbor_mask = mask & ~use_replace
if neighbor_mask.any():
result = self._neighbor_perturb(result, neighbor_mask, tokens)
elif mode == "neighbor":
result = self._neighbor_perturb(result, mask, tokens)
else:
raise ValueError(
f"Unsupported token perturbation mode {mode!r}; expected "
"'replace', 'mixed', or 'neighbor'."
)
return result
def _neighbor_perturb(
self,
result: torch.Tensor,
mask: torch.Tensor,
tokens: torch.Tensor,
) -> torch.Tensor:
b, _ = tokens.shape
# Neighbor perturbation is easiest in flat FSQ-scalar space: each scalar
# index can move by -1, 0, or +1 before packing back into prior tokens.
fsq_indices = self.prior_tokens_to_fsq_indices(tokens)
fsq_indices = fsq_indices.view(
b,
self.num_prior_tokens,
self.fsq_scalars_per_prior_token,
)
offset = torch.randint(-1, 2, fsq_indices.shape, device=tokens.device)
perturbed = (fsq_indices + offset).clamp(0, self.num_fsq_levels - 1)
# Convert back to the categorical token vocabulary expected by the prior.
perturbed = perturbed.view(b, -1)
neighbor_prior_tokens = self.fsq_indices_to_prior_tokens(perturbed)
return torch.where(mask, neighbor_prior_tokens, result)
# ---- encode / decode ----
def _encode(self, tensordict):
"""Encode observations to FSQ codes. Requires mimic_target_poses in tensordict."""
if self.target_latent_encoder is None:
raise RuntimeError(
"DiscretePriorPEFTActor._encode requires mimic_target_poses_dim > 0 "
"so the target_latent_encoder path is available."
)
encoder = self.target_latent_encoder.encoder
td = encoder(tensordict)
key = encoder.out_keys[0] if hasattr(encoder, "out_keys") else "latent"
return self.quantize(td[key])
def _decode(self, tensordict, fsq_codes):
"""Decode FSQ codes to actions. Only needs max_coords_obs + latent."""
tensordict[self.decoder_latent_key] = fsq_codes
decoder = self.latent_decoder.decoder
td = decoder(tensordict)
key = decoder.out_keys[0] if hasattr(decoder, "out_keys") else "mu"
return td[key]
[docs]
def predict_target_prior_tokens(self, tensordict):
"""Encode target poses to GPC prior tokens for SFT."""
with torch.no_grad():
codes = self._encode(tensordict)
fsq_indices = self.fsq_codes_to_fsq_indices(codes)
return self.fsq_indices_to_prior_tokens(fsq_indices)
# ---- prior input construction ----
def _to_tensordict(self, d):
if isinstance(d, TensorDict):
return d.clone()
first_tensor = next(value for value in d.values() if torch.is_tensor(value))
return TensorDict(
dict(d),
batch_size=first_tensor.shape[0],
device=first_tensor.device,
)
def _run_actor_peft_model(self, d, network: nn.Module | None = None):
"""Run a conditioning network and validate its contract.
``network`` defaults to the student's ``actor_peft_model``.
"""
network = self.actor_peft_model if network is None else network
tensordict = self._to_tensordict(d)
missing_keys = [key for key in network.in_keys if key not in tensordict]
if missing_keys:
raise ValueError(
"DiscretePriorPEFTActor actor PEFT model in_keys must be present in "
f"the input TensorDict. Missing keys: {missing_keys}"
)
tensordict = network(tensordict)
if self.condition_key not in tensordict:
raise RuntimeError(
"DiscretePriorPEFTActor actor PEFT model did not produce required "
f"condition_key {self.condition_key!r}."
)
return tensordict
# ---- forward / rollout ----
[docs]
def forward(self, input_dict: dict):
"""Teacher-forced forward -> logits (B, num_prior_tokens, prior_token_vocab_size)."""
return self.prior_with_peft(input_dict)
[docs]
def get_action_and_logp(self, tensordict):
"""Rollout step: generate action + per-token log-probs for PPO."""
prior_dict = self.build_prior_input(tensordict)
prior_tokens, logprob = self.prior_with_peft.generate(
prior_dict,
return_logits=False,
return_logprob=True,
)
neglogp = -logprob
fsq_indices = self.prior_tokens_to_fsq_indices(prior_tokens)
fsq_codes = self.fsq_indices_to_codes(fsq_indices)
action = self._decode(tensordict, fsq_codes)
tensordict["action"] = action
tensordict["mean_action"] = action
tensordict["neglogp"] = neglogp
tensordict["prior_tokens"] = prior_tokens
tensordict[LATENT_KEY] = prior_tokens
tensordict[LATENT_LOGPROB_KEY] = logprob
return tensordict
