# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""Configuration classes for PPO agent.
This module defines all configuration dataclasses for the Proximal Policy Optimization (PPO)
algorithm, including actor-critic architecture parameters, optimization settings, and
training hyperparameters.
Key Classes:
- PPOAgentConfig: Main PPO agent configuration
- PPOModelConfig: PPO model (actor-critic) configuration
- PPOActorConfig: Policy network configuration
- AdvantageNormalizationConfig: Advantage normalization settings
"""
from typing import Dict, List, Optional
from dataclasses import dataclass, field
from enum import Enum
import math
from protomotions.agents.common.config import (
ModuleContainerConfig,
)
from protomotions.agents.base_agent.config import (
OptimizerConfig,
BaseAgentConfig,
BaseModelConfig,
)
[docs]
class StdMode(str, Enum):
FIXED = "fixed"
GLOBAL = "global"
PER_MOTION = "per_motion"
[docs]
@dataclass
class PPOActorConfig:
"""Configuration for PPO Actor network."""
mu_key: str = field(metadata={"help": "The key of the output of the mu model."})
# The recipe lists the actor's observation keys once, here. A plain mean
# network (mu_model) that does not declare its own in_keys inherits them in
# __post_init__; a composite mean network (e.g. an autoencoder) that reads its
# obs through its own sub-models declares them there and is left untouched.
in_keys: List[str] = field(
default_factory=list, metadata={"help": "Input observation keys."}
)
out_keys: List[str] = field(
default_factory=lambda: ["action", "mean_action", "neglogp"],
metadata={"help": "Output keys: action, mean_action, neglogp."},
)
_target_: str = "protomotions.agents.ppo.model.PPOActor"
mu_model: ModuleContainerConfig = field(
default_factory=ModuleContainerConfig,
metadata={"help": "Neural network model for action mean."},
)
num_out: int = field(
default=None, metadata={"help": "Number of actions. Set from robot config."}
)
actor_logstd: float = field(
default=-2.9, metadata={"help": "Initial log std for action distribution."}
)
std_mode: StdMode = StdMode.FIXED
min_logstd: float = math.log(0.001)
max_logstd: float = math.log(0.5)
# Filled by training setup before constructing the actor. Not an observation.
num_motions: Optional[int] = None
def __post_init__(self):
# A plain mean network inherits the actor's obs keys; a composite one that
# already declares them through its sub-models keeps its own.
if not self.mu_model.in_keys:
self.mu_model.in_keys = list(self.in_keys)
def _union_keys(configs, attr: str) -> List[str]:
"""Ordered, de-duplicated union of ``attr`` over component configs."""
return list(
dict.fromkeys(
key for config in configs for key in (getattr(config, attr, None) or [])
)
)
[docs]
@dataclass
class PPOModelConfig(BaseModelConfig):
"""Configuration for PPO Model (Actor-Critic).
``in_keys`` and ``out_keys`` are NOT passed: the recipe lists its observation
keys once on the actor's mean network and once on the critic, and the model
derives the unions from its component configs. They are ``init=False``, so
passing a list is a construction error (and the saved form skips them and
recomputes on load).
"""
_target_: str = "protomotions.agents.ppo.model.PPOModel"
in_keys: List[str] = field(init=False, default_factory=list)
out_keys: List[str] = field(init=False, default_factory=list)
actor: PPOActorConfig = field(
default_factory=PPOActorConfig,
metadata={"help": "Actor (policy) network configuration."},
)
critic: ModuleContainerConfig = field(
default_factory=ModuleContainerConfig,
metadata={"help": "Critic (value) network configuration."},
)
actor_optimizer: OptimizerConfig = field(
default_factory=lambda: OptimizerConfig(lr=2e-5),
metadata={"help": "Optimizer settings for actor network."},
)
critic_optimizer: OptimizerConfig = field(
default_factory=lambda: OptimizerConfig(lr=1e-4),
metadata={"help": "Optimizer settings for critic network."},
)
def _key_source_configs(self) -> List:
"""Component configs whose in_keys/out_keys union into the model's."""
return [self.actor, self.critic]
def __post_init__(self):
sources = self._key_source_configs()
self.in_keys = _union_keys(sources, "in_keys")
self.out_keys = _union_keys(sources, "out_keys")
[docs]
@dataclass
class AdvantageNormalizationConfig:
"""Configuration for advantage normalization."""
enabled: bool = field(
default=True, metadata={"help": "Whether to normalize advantages."}
)
shift_mean: bool = field(
default=True, metadata={"help": "Subtract mean from advantages."}
)
# EMA parameters
use_ema: bool = field(
default=True, metadata={"help": "Use EMA for normalization statistics."}
)
ema_alpha: float = field(
default=0.05, metadata={"help": "EMA weight for new data."}
)
min_std: float = field(
default=0.02, metadata={"help": "Minimum std to prevent extreme normalization."}
)
clamp_range: float = field(
default=4.0,
metadata={"help": "Clamp normalized advantages to [-range, range]."},
)
[docs]
@dataclass
class AdaptiveLRConfig:
"""Configuration for adaptive learning rate based on KL divergence."""
enabled: bool = field(
default=False,
metadata={"help": "Enable adaptive learning rate based on KL divergence."},
)
desired_kl: float = field(
default=0.01,
metadata={"help": "Target KL divergence for adaptive learning rate."},
)
min_lr: float = field(
default=1e-5,
metadata={"help": "Minimum learning rate for both actor and critic."},
)
max_lr: float = field(
default=1e-2,
metadata={"help": "Maximum learning rate for both actor and critic."},
)
[docs]
@dataclass
class L2C2Config:
"""L2C2 (Lipschitz-ratio) actor regularization (Kobayashi 2022).
Penalizes the ratio ||mu(noisy) - mu(clean)||^2 / ||noisy - clean||^2
so the actor's Lipschitz constant stays bounded.
"""
enabled: bool = field(
default=False, metadata={"help": "Enable L2C2 regularization."}
)
lambda_l2c2: float = field(default=0.1, metadata={"help": "L2C2 loss coefficient."})
obs_pairs: Dict[str, str] = field(
default_factory=dict,
metadata={"help": "Map from noisy actor obs key to clean counterpart key."},
)
[docs]
@dataclass
class PPOAgentConfig(BaseAgentConfig):
"""Main configuration class for PPO Agent."""
_target_: str = "protomotions.agents.ppo.agent.PPO"
# Model configuration
model: PPOModelConfig = field(
default_factory=PPOModelConfig, metadata={"help": "Model configuration."}
)
# PPO hyperparameters
tau: float = field(
default=0.95, metadata={"help": "GAE lambda for advantage estimation."}
)
e_clip: float = field(
default=0.2, metadata={"help": "PPO clipping parameter epsilon."}
)
clip_critic_loss: bool = field(
default=True, metadata={"help": "Clip critic loss similar to actor."}
)
# Actor update control
actor_clip_frac_threshold: Optional[float] = field(
default=0.6,
metadata={"help": "Skip actor update if clip_frac > threshold (e.g., 0.5)."},
)
# Entropy regularization for learned exploration.
entropy_coef: float = field(
default=0.005,
metadata={"help": "Entropy bonus coefficient for learnable std exploration."},
)
# L2C2 regularization
l2c2: L2C2Config = field(
default_factory=L2C2Config, metadata={"help": "L2C2 settings."}
)
# Adaptive learning rate
adaptive_lr: AdaptiveLRConfig = field(
default_factory=AdaptiveLRConfig,
metadata={"help": "Adaptive learning rate settings."},
)
# Value normalization
advantage_normalization: AdvantageNormalizationConfig = field(
default_factory=AdvantageNormalizationConfig,
metadata={"help": "Advantage normalization settings."},
)
