protomotions.agents.ppo.agent module

Proximal Policy Optimization (PPO) agent implementation.

This module implements the PPO algorithm for reinforcement learning. PPO is an on-policy algorithm that uses clipped surrogate objectives for stable policy updates. It collects experience through environment interaction and performs multiple epochs of minibatch updates using Generalized Advantage Estimation (GAE).

Key Classes:

• PPO: Main PPO agent class extending BaseAgent

References

Schulman et al. “Proximal Policy Optimization Algorithms” (2017)

class protomotions.agents.ppo.agent.PPO(fabric, env, config, root_dir=None)

Bases: BaseAgent

Proximal Policy Optimization (PPO) agent.

Implements the PPO algorithm for training reinforcement learning policies. PPO uses clipped surrogate objectives to enable stable policy updates while maintaining sample efficiency. This implementation supports actor-critic architecture with separate optimizers for policy and value networks.

The agent collects experience through environment interaction, computes advantages using Generalized Advantage Estimation (GAE), and performs multiple epochs of minibatch updates on the collected data.

Parameters:

• fabric (MockFabric) – Lightning Fabric instance for distributed training.

• env (BaseEnv) – Environment instance to train on.

• config (PPOAgentConfig) – PPO-specific configuration including learning rates, clip parameters, etc.

• root_dir (Path | None) – Optional root directory for saving outputs.

Variables:

• tau – GAE lambda parameter for advantage estimation.

• e_clip – PPO clipping parameter for policy updates.

• actor – Policy network.

• critic – Value network.

Example

>>> fabric = Fabric(devices=4)
>>> env = Steering(config, robot_config, simulator_config, device)
>>> agent = PPO(fabric, env, config)
>>> agent.setup()
>>> agent.train()

config: PPOAgentConfig

motion_index: MotionIndex | None = None

__init__(fabric, env, config, root_dir=None)

Initialize the base agent.

Sets up distributed training infrastructure, initializes tracking metrics, and creates the evaluator. Subclasses should call super().__init__() first.

Parameters:

• fabric (MockFabric) – Lightning Fabric for distributed training and device management.

• env (BaseEnv) – Environment instance for agent-environment interaction.

• config (PPOAgentConfig) – Configuration containing hyperparameters and training settings.

• root_dir (Path | None) – Optional directory for saving outputs (uses logger dir if None).

property has_critic: bool

PPO owns a value critic by construction.

create_model()

Create PPO actor-critic model.

Instantiates the PPO model with actor and critic networks, applies weight initialization, and returns the model.

Returns:

PPOModel instance with initialized weights.

add_agent_info_to_obs(obs)

Add agent-specific observations to the environment observations.

Must be stateless — reads agent state set by pre_collect_step() but does not modify it. Called for the actor’s obs during rollout. Also called during evaluation and setup (where pre_collect_step is not called), so it must not depend on pre_collect_step having run.

create_optimizers(model)

Create separate optimizers for actor and critic.

Uses Fabric to prepare both model/optimizer pairs for distributed training, matching the public PPO setup path.

Parameters:

model (PPOModel) – PPOModel with actor and critic networks.

property actor_module

Underlying actor module, unwrapped from DDP when needed.

property critic_module

Underlying critic module, unwrapped from DDP when needed.

load_parameters(state_dict, load_training_state=True)

Load agent parameters from state dictionary.

Always restores model weights. When load_training_state is true, also restores optimizer-era state such as counters, reward normalization, and evaluator state.

Parameters:

state_dict – Dictionary containing saved agent state from checkpoint. Expected keys: epoch, step_count, run_start_time, best_evaluated_score, running_reward_norm (if normalization enabled), model.

get_state_dict(state_dict)

Get complete state dictionary for checkpointing.

Collects all agent state including model weights, training progress, and normalization statistics into a single dictionary for saving.

Parameters:

state_dict – Existing state dict to update (typically empty dict).

Returns:

Updated state dictionary containing all agent state.

save(checkpoint_name='last.ckpt', new_high_score=False)

Save model checkpoint and environment state.

Rank 0 saves the main checkpoint (shared model weights). Each unique-task rank saves its env checkpoint. Global barrier ensures all ranks wait until saving is complete.

Parameters:

• checkpoint_name – Name of checkpoint file (e.g., “last.ckpt” or “epoch_100.ckpt”)

• new_high_score – Whether this is a new high score (will also save as score_based.ckpt)

register_algorithm_experience_buffer_keys()

Register algorithm-specific keys in the experience buffer.

Subclasses override this to add custom keys to the experience buffer (e.g., AMP adds discriminator observations, ASE adds latent codes).

record_rollout_step(

next_obs_td,

actions,

rewards,

dones,

terminated,

done_indices,

extras,

step,

)

Record PPO-specific data: next value estimates for GAE computation.

perform_optimization_step(batch_dict, batch_idx)

Perform one PPO optimization step on a minibatch.

Computes actor and critic losses, performs backpropagation, clips gradients, and updates both networks.

Parameters:

• batch_dict – Dictionary containing minibatch data (obs, actions, advantages, etc.).

• batch_idx – Index of current batch (unused but kept for compatibility).

Returns:

Dictionary of training metrics (losses, clip fraction, etc.).

Return type:

Dict

actor_step(batch_dict)

Compute actor loss and perform policy update.

Computes PPO clipped surrogate objective plus optional bounds loss and extra algorithm-specific losses.

Parameters:

batch_dict – Minibatch containing obs, actions, old neglogp, advantages.

Returns:

• actor_loss: Total actor loss for backprop

• log_dict: Dictionary of actor metrics for logging

Return type:

Tuple of (actor_loss, log_dict) where

calculate_extra_actor_loss(batch_td)

Calculate additional actor losses beyond PPO objective.

Supports L2C2 regularization: penalizes Lipschitz ratio between actor outputs on noisy vs clean observations (Kobayashi 2022).

Parameters:

batch_td – Minibatch data (post actor forward, contains mean_action).

Returns:

Tuple of (extra_loss, log_dict) with additional loss and metrics.

Return type:

Tuple[MockTensor, Dict]

critic_step(batch_dict)

optimize_model()