protomotions.agents.amp.model module#

AMP model components including discriminator network.

This module implements the AMP-specific neural networks, particularly the discriminator that distinguishes between agent and reference motion data.

Key Classes:

  • Discriminator: Binary classifier for agent vs. reference motions

  • AMPModel: PPO model extended with discriminator

class protomotions.agents.amp.model.Discriminator(*args, **kwargs)[source]#

Bases: ProtoMotionsTensorDictModule

Discriminator network for AMP style rewards.

Binary classifier that distinguishes between agent-generated and reference motion data. Uses ModuleContainer structure - just chains models together.

Parameters:

config (DiscriminatorConfig) – DiscriminatorConfig (extends ModuleContainerConfig).

Variables:

  • models – ModuleContainer list of modules.

  • in_keys – Input keys from config.

  • out_keys – Output keys from config.

__init__(config)[source]#
config: DiscriminatorConfig#
forward(tensordict, log_internals=False)[source]#

Forward pass through discriminator.

Parameters:

tensordict (MockTensorDict) – TensorDict containing observations.

Returns:

TensorDict with discriminator output added.

Return type:

MockTensorDict

compute_disc_reward(disc_logits, eps=0.0001)[source]#

Compute style reward from discriminator logits.

Converts discriminator logits to reward using negative log probability. Higher reward means motion is more similar to reference data.

Parameters:

  • disc_logits (MockTensor) – Discriminator logits.

  • eps (float) – Small constant for numerical stability.

Returns:

Style rewards for each sample (higher = more reference-like).

Return type:

MockTensor

all_discriminator_weights()[source]#

Get all discriminator weight matrices (works with LazyLinear).

Returns:

List of weight parameters from all linear layers in discriminator.

logit_weights()[source]#

Get the final layer weights (logit layer).

Returns:

List containing the output layer weight parameter.

Return type:

List[<Mock object at 0x7a9284e585d0>[]]

class protomotions.agents.amp.model.AMPModelComponentsMixin[source]#

Bases: object

Adds AMP discriminator modules to a host model.

optional_checkpoint_modules()[source]#
class protomotions.agents.amp.model.AMPModel(*args, **kwargs)[source]#

Bases: AMPModelComponentsMixin, PPOModel

AMP model with actor, task critic, disc critic, and discriminator networks.

Extends PPOModel by adding a discriminator network that provides style rewards and a separate critic for estimating discriminator reward values.

Parameters:

config (AMPModelConfig) – AMPModelConfig specifying all networks.

Variables:

  • _actor – Policy network.

  • _critic – Task value network.

  • _disc_critic – Discriminator reward value network.

  • _discriminator – Style discriminator network.

config: AMPModelConfig#
__init__(config)[source]#
forward(tensordict, log_internals=False)[source]#

Forward pass through PPO, and discriminator.

Parameters:

tensordict (MockTensorDict) – TensorDict containing observations.

Returns:

TensorDict with all model outputs added.

Return type:

MockTensorDict