protomotions.agents.peft.model module#

Top-level discrete-prior PEFT model container.

The pretrained prior is passed in by the agent. This class only wires the PEFT actor to the optional PPO critic. Rollout is the module forward path.

class protomotions.agents.peft.model.DiscretePriorPEFTModel(*args, **kwargs)[source]#

Bases: BaseModel

Discrete-prior PEFT actor plus optional PPO critic.

forward is reserved for rollout TensorDicts. SFT teacher forcing lives on DiscretePriorPEFTSFTModel so RLFT and SFT flows stay separate.

__init__(
config,
pretrained_prior_model,
mimic_target_poses_dim=0,
)[source]#
optional_checkpoint_modules()[source]#

SFT checkpoints can initialize RLFT with a fresh value critic.

materialize(tensordict)[source]#

Initialize shapes with teacher forcing, never sampling or pinning a reference.

forward(tensordict)[source]#

Run rollout inference and critic evaluation.

forward_rollout(tensordict)[source]#