# 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 protomotions.agents.common.config import (
ModuleContainerConfig,
)
from protomotions.agents.base_agent.config import (
OptimizerConfig,
BaseAgentConfig,
BaseModelConfig,
)
[docs]
@dataclass
class PPOActorConfig:
"""Configuration for PPO Actor network."""
mu_key: str = field(metadata={"help": "The key of the output of the mu model."})
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."}
)
learnable_std: bool = field(
default=False,
metadata={"help": "Make action log std learnable (requires_grad=True)."},
)
[docs]
@dataclass
class PPOModelConfig(BaseModelConfig):
"""Configuration for PPO Model (Actor-Critic)."""
_target_: str = "protomotions.agents.ppo.model.PPOModel"
out_keys: List[str] = field(
default_factory=lambda: ["action", "mean_action", "neglogp", "value"],
metadata={"help": "Output keys including actions and value estimate."},
)
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."},
)
[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 (used when actor has learnable_std)
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."},
)
