protomotions.agents.peft.prior_amp_model module#

Discrete-prior PEFT RLFT model with AMP discriminator outputs.

class protomotions.agents.peft.prior_amp_model.DiscretePriorPEFTRLFTAMPModel(*args, **kwargs)[source]#

Bases: AMPModelComponentsMixin, DiscretePriorPEFTModel

Discrete-prior PEFT actor/critic with AMP discriminator outputs.

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

AMP discriminator state is optional for PEFT full-checkpoint loads.

experience_buffer_keys()[source]#

Keys produced during rollout that the experience buffer must store.

forward_rollout(tensordict)[source]#