Proximal Policy Optimization (PPO)¶

PPO is a reliable policy-gradient method that balances implementation simplicity with stable learning. Our implementation trains the actor and critic on batches of collected rollouts, using a clipped surrogate objective, policy entropy, and GAE-style advantage estimation.

Components¶

Actor (Gaussian policy): parameters \(\mu, \sigma\) define \(\mathcal{N}(\mu, \sigma^2)\)
Critic: scalar value estimate \(V(s)\)
Experience collection: rollout of length rollout_len storing \(s,a,\log\pi(a|s), r, d, V(s)\)
Training: mini-batches across num_epochs with clipped probability ratios

Theory¶

Probability ratio:

\[ r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{\text{old}}}(a_t|s_t)} = \exp\big(\log \pi_\theta - \log \pi_{\theta_{\text{old}}}\big) \]

Clipped surrogate (Actor):

\[ \mathcal{L}_\text{actor} = -\,\mathbb{E}\Big[\min\big( r_t\,A_t,\ \mathrm{clip}(r_t,\ 1-\varepsilon,\ 1+\varepsilon)\,A_t \big) \Big] \]

Critic loss (Value):

\[ \mathcal{L}_\text{critic} = \mathbb{E}\big[ (R_t - V_\phi(s_t))^2 \big] \]

Entropy regularization:

\[ \mathcal{L}_\text{entropy} = -\beta\,\mathbb{E}\big[\mathcal{H}[\pi_\theta(\cdot|s_t)]\big] \]

Total loss: \(\mathcal{L} = \mathcal{L}_\text{actor} + \mathcal{L}_\text{critic} + \mathcal{L}_\text{entropy}\)
Advantage (GAE-like): preprocess1 returns \(\text{return} = V + \sum\gamma\lambda\,\delta\), with \(A = \text{return} - V\)

\[ \delta_t = r_t + \gamma V(s_{t+1}) - V(s_t),\quad \hat{A}_t \approx \sum_{l=0}^{\infty} (\gamma\lambda)^l\, \delta_{t+l} \]

Implementation details¶

Policy: Actor.forward(..., continous_actions=True) outputs mu = tanh(Wx) and log_std = tanh(Wx) stretched to [log_std_min, log_std_max]; \(\sigma = e^{\log \sigma}\). Actions are sampled from Normal(mu, sigma).
Probability ratios: computed as torch.exp(new_probs - old_probs) for stability versus dividing densities.
Entropy: actor_loss uses -new_distr.entropy().mean() and later adds + entropy_coef * entropy, effectively subtracting entropy with a coefficient to encourage stochastic policies.
GAE & bootstrap: preprocess1 appends next_value to values, then iterates backward computing \(\delta\) and accumulating \(g\) with \(\lambda=0.8\); final returns = V + g, advantages = returns - V.
Mini-batches: ppo_iter samples mini_batch_size indices repeatedly each epoch.
Auxiliary head: actor also predicts rewards (self.r) for an optional auxillary_task (reward MSE); not included in the default loss.

Training pseudocode¶

for episode in range(max_episodes):
  rollout = collect(rollout_len)
  next_value = V(s_T)
  returns, advantages = GAE(rollout.rewards, rollout.values, dones, gamma, lambda)
  for epoch in range(num_epochs):
    for batch in mini_batches(rollout, returns, advantages):
      ratios = exp(new_logp - old_logp)
      a_loss = -mean(min(ratios*A, clip(ratios)*A)) + entropy_coef * (-entropy)
      c_loss = mse(returns - V(s))
      update(actor, critic)
  log TensorBoard metrics

Hyperparameters and mapping¶

clip_pram = ε — clipping threshold for probability ratios
num_epochs, batch_size — epochs and mini-batch size per update
rollout_len — rollout length prior to updates
entropy_coef — weight of the entropy term (watch sign in code)
actor_lr, critic_lr — Adam learning rates
gamma, lambda(=0.8) — discount and GAE parameter inside preprocess1

Quick start¶

import gymnasium as gym
from tensoraerospace.agent.ppo.model import PPO

# Create environment (F-16 example)
env = gym.make('LinearLongitudinalF16-v0', number_time_steps=2000)

# Initialize PPO
agent = PPO(
    env=env,
    gamma=0.99,
    max_episodes=50,
    rollout_len=2048,
    clip_pram=0.2,
    num_epochs=64,
    batch_size=64,
    entropy_coef=0.005,
    actor_lr=1e-3,
    critic_lr=5e-3,
)

# Train
agent.train()

# Save and load
agent.save('./runs')

# Load a trained agent
agent = PPO.load('./runs')

Tip

For continuous actions use the Gaussian policy; clamp log_std (as in code) and normalize features.

Practical tips¶

Increase rollout_len for stabler advantage estimates
Balance clip_pram (typically 0.1–0.3) and entropy_coef for exploration
Multiple epochs (num_epochs) with smaller batch_size help convergence—watch for overfitting

Auxiliary Tasks¶

The PPO implementation includes an optional auxiliary task mechanism for reward prediction. This auxiliary head helps the agent learn better state representations by predicting expected rewards alongside the main policy optimization.

How it works¶

The Actor network includes an additional output layer self.r that predicts the reward
The auxiliary loss is computed as MSE between predicted and actual rewards
This loss can be added to the main PPO loss via the auxillary_task method in the Agent class

Usage¶

# Auxiliary task is computed separately from main training
aux_loss = agent.auxillary_task(states, rewards)

The auxiliary task encourages the network to encode reward-relevant features in its hidden representations, potentially improving sample efficiency and generalization.

Unified training interface¶

PPO follows the shared unified train() API from BaseRLModel:

stats = agent.train(
    num_episodes=200,   # optional: overrides self.max_episodes
    max_steps=1024,     # optional: overrides self.rollout_len
)

Calling agent.train() with no arguments is still supported and uses the hyperparameters set at construction time. Note that PPO's learn(states, actions, adv, old_probs, returns, rewards, old_values) method is an internal per-batch gradient update helper and is kept untouched by the unified interface.

API reference¶

PPO(env, gamma=0.99, max_episodes=30, rollout_len=2048, clip_pram=0.2, num_epochs=64, batch_size=64, entropy_coef=0.005, actor_lr=0.001, critic_lr=0.005, gae_lambda=0.95, max_grad_norm=0.5, target_kl=None, normalize_obs=True, normalize_reward=False, actor_hidden_dim=256, critic_hidden_dim=256, actor_log_std_min=-20.0, actor_log_std_max=0.0, eval_freq=10, seed=336699, device=None, log_dir=None, save_best_model=True, best_model_dir=None, save_best_async=True, wandb_project=None, wandb_entity=None, wandb_run_name=None, wandb_tags=None, wandb_config=None) ¶

Bases: BaseRLModel

Proximal Policy Optimization (PPO) reinforcement learning agent.

PPO is a policy gradient method that uses a clipped objective function to ensure stable and monotonic policy improvements. This implementation includes: - Actor-Critic architecture with separate networks - Generalized Advantage Estimation (GAE) - Observation and reward normalization - Value function clipping - Gradient clipping for stability - KL divergence early stopping - TensorBoard logging

The agent is designed for continuous control tasks in aerospace applications.

Attributes:

Name	Type	Description
`actor`		Policy network that outputs action distributions.
`critic`		Value network that estimates state values.
`gamma`		Discount factor for future rewards.
`clip_pram`		PPO clipping parameter epsilon.
`gae_lambda`		GAE lambda for advantage estimation.
`max_grad_norm`		Maximum gradient norm for clipping.
`target_kl`		Target KL divergence for early stopping.
`normalize_obs`		Whether to normalize observations.
`normalize_reward`		Whether to normalize rewards.
`obs_rms`		Running statistics for observation normalization.
`ret_rms`		Running statistics for return normalization.
`writer`		TensorBoard summary writer.

Initialize agent with given environment and discount coefficient.

Parameters:

Name	Type	Description	Default
`env`	`Any`	Environment object with which agent will interact.	required
`gamma`	`float`	Discount coefficient. Defaults to 0.99.	`0.99`
`max_episodes`	`int`	Maximum number of training episodes.	`30`
`rollout_len`	`int`	Number of steps per rollout.	`2048`
`clip_pram`	`float`	PPO clipping parameter epsilon.	`0.2`
`num_epochs`	`int`	Number of optimization epochs per rollout.	`64`
`batch_size`	`int`	Mini-batch size for SGD.	`64`
`entropy_coef`	`float`	Entropy bonus coefficient.	`0.005`
`actor_lr`	`float`	Learning rate for actor network.	`0.001`
`critic_lr`	`float`	Learning rate for critic network.	`0.005`
`gae_lambda`	`float`	GAE lambda parameter for advantage estimation.	`0.95`
`max_grad_norm`	`float`	Maximum gradient norm for clipping.	`0.5`
`target_kl`	`Optional[float]`	Target KL divergence for early stopping.	`None`
`normalize_obs`	`bool`	Whether to normalize observations.	`True`
`normalize_reward`	`bool`	Whether to normalize rewards.	`False`
`actor_hidden_dim`	`int`	Hidden layer size for actor network.	`256`
`critic_hidden_dim`	`int`	Hidden layer size for critic network.	`256`
`eval_freq`	`int`	Frequency (in episodes) for evaluation.	`10`
`seed`	`int`	Random seed.	`336699`
`device`	`Union[str, device, None]`	Torch device to run the training on. If None, the agent will auto-select CUDA (if available), else MPS (if available), else CPU.	`None`
`save_best_model`	`bool`	Whether to save the best checkpoint during training.	`True`
`best_model_dir`	`Union[str, Path, None]`	Directory for the best checkpoint (config.json + weights). If None, defaults to "{cwd}/best_model_PPO/".	`None`
`save_best_async`	`bool`	If True, save best checkpoint in a background thread (recommended to avoid slowing training).	`True`

`act(state, deterministic=False)` ¶

Select action for given state.

Parameters:

Name	Type	Description	Default
`state`	`ndarray \| Tensor`	Current environment state.	required
`deterministic`	`bool`	If True, use mean action (no sampling).	`False`

Returns:

Name	Type	Description
`tuple`	`Tuple[Tensor, ndarray, Tensor]`	Tuple containing action, mean action and log probability.

`actor_loss(probs, entropy, actions, adv, old_probs)` ¶

Calculate actor losses.

Parameters:

Name	Type	Description	Default
`probs`	`Tensor`	Action probabilities of new policy.	required
`entropy`	`Tensor`	Action entropy.	required
`actions`	`Tensor`	Actions taken.	required
`adv`	`Tensor`	Advantages.	required
`old_probs`	`Tensor`	Action probabilities of old policy.	required

Returns:

Name	Type	Description
`Tensor`	`Tensor`	Actor loss function value.

`eval()` ¶

Switch actor and critic networks to evaluation mode.

Returns:

Type	Description
`PPO`	self for method chaining.

`close()` ¶

Flush and stop background saver (safe to call multiple times).

`learn(states, actions, adv, old_probs, discnt_rewards, rewards, old_values)` ¶

Agent training procedure.

Parameters:

Name	Type	Description	Default
`states`	`Tensor`	States experienced by agent.	required
`actions`	`Tensor`	Actions taken by agent.	required
`adv`	`Tensor`	Advantages.	required
`old_probs`	`Tensor`	Log probabilities of previous actions.	required
`discnt_rewards`	`Tensor`	Discounted rewards.	required
`rewards`	`Tensor`	Actual received rewards.	required
`old_values`	`Tensor`	Previous value function estimates.	required

Returns:

Name	Type	Description
`dict`	`Dict[str, float]`	Dictionary with training metrics.

`test_reward()` ¶

Test model by executing one episode with deterministic actions.

Returns:

Name	Type	Description
`float`	`float`	Total reward per episode.

`preprocess1(states, actions, rewards, dones, values, probs, gamma)` ¶

Preprocess transitions for buffer.

Parameters:

Name	Type	Description	Default
`states`	`list[Tensor]`	List of states.	required
`actions`	`list[Tensor]`	List of actions.	required
`rewards`	`list[Tensor]`	List of rewards.	required
`dones`	`list[Tensor]`	List of boolean values indicating episode termination.	required
`values`	`list[Tensor]`	State values.	required
`probs`	`list[Tensor]`	Log probabilities of actions.	required
`gamma`	`float`	Discount coefficient.	required

Returns:

Name	Type	Description
`tuple`	`Tuple[Tensor, Tensor, list[Tensor], Tensor, Tensor, Tensor]`	Tuple containing processed states, actions, rewards, advantages and probabilities.

`train(num_episodes=None, *, max_steps=None, save_best=False, save_path=None, verbose=True, **kwargs)` ¶

Train the PPO agent through interaction with the environment.

This method implements the complete PPO training loop

Collect rollout data by interacting with environment
Compute advantages using GAE
Update policy and value function using mini-batch SGD
Log metrics to TensorBoard
Periodically evaluate policy performance

Parameters:

Name	Type	Description	Default
`num_episodes`	`Optional[int]`	Number of training episodes. When `None` (the default), PPO falls back to `self.max_episodes` which was set at construction time. This preserves the original no-argument call style.	`None`
`max_steps`	`Optional[int]`	Optional override for `self.rollout_len` so callers can shorten each rollout via the unified API.	`None`
`save_best`	`bool`	Reserved for unified interface; PPO already handles best-model saving via its internal background saver, so this flag is currently a no-op.	`False`
`save_path`	`Optional[str]`	Reserved for unified interface (see `save_best`).	`None`
`verbose`	`bool`	Reserved for symmetry with other agents.	`True`
`**kwargs`	`Any`	Additional algorithm-specific keyword arguments (currently ignored by PPO).	`{}`

Returns:

Name	Type	Description
`dict`	`dict`	Training metrics dictionary with episode rewards
	`dict`	collected so far, the final running average and any
	`dict`	early-stopping flag PPO may have triggered.

Training can be stopped early using KL divergence thresholds (target_kl) or by setting self.target = True.

Note

All metrics are logged to TensorBoard including actor/critic losses, rewards, entropy, KL divergence, clip fraction, and explained variance.

`get_param_env()` ¶

Get environment and agent parameters for serialization.

This method extracts all necessary information to reconstruct the agent and its environment, including hyperparameters, network architectures, and environment specifications.

Returns:

Type	Description
`Dict[str, Dict[str, Any]]`	Dictionary with two main keys: - 'env': Dictionary containing environment name and parameters. - 'policy': Dictionary containing agent name and hyperparameters.

Note

For TensorAeroSpace environments, full environment parameters are serialized. For other environments, only the class name is stored.

`save(path=None)` ¶

Save the PPO model to disk.

This method saves all components needed to restore the agent

Configuration file (config.json) with hyperparameters
Actor network weights (actor.pth)
Critic network weights (critic.pth)
Actor optimizer state (actor_opt.pth) for resuming training
Critic optimizer state (critic_opt.pth) for resuming training
Training state (train_state.json): best_reward, timestamps, etc.
Observation normalization statistics (obs_rms.npz, if enabled)
Return normalization statistics (ret_rms.npz, if enabled)

The model is saved in a timestamped directory with format: {path}/{Month}{Day}_{Hour}-{Minute}-{Second}_PPO/

Parameters:

Name	Type	Description	Default
`path`	`Union[str, Path, None]`	Directory where the model will be saved. If None, uses current working directory. Defaults to None.	`None`

Returns:

Type	Description
`Path`	Path to the created timestamped directory.

Example

agent.save('/path/to/models')

Saves to: /path/to/models/Oct05_14-30-45_PPO/¶

`__load(path)` `classmethod` ¶

Load a PPO model from disk (internal method).

This private method handles the complete restoration of a saved PPO agent, including network weights, configuration, and normalization statistics.

Parameters:

Name	Type	Description	Default
`path`	`Union[str, Path]`	Directory containing the saved model files (config.json, actor.pth, critic.pth, and optional normalization files).	required

Returns:

Type	Description
`PPO`	Restored PPO agent instance with loaded weights and configuration.

Raises:

Type	Description
`TheEnvironmentDoesNotMatch`	If the agent type in the saved config does not match the current class.
`FileNotFoundError`	If required model files are not found.

Note

This is a private method. Use from_pretrained() for loading models.

`from_pretrained(repo_name, access_token=None, version=None)` `classmethod` ¶

Load a pretrained PPO model from local path or Hugging Face Hub.

This method provides a unified interface for loading models from either

Local filesystem paths
Hugging Face Hub repositories

The method automatically detects whether repo_name is a local path or a remote repository and handles downloading/loading appropriately.

Parameters:

Name	Type	Description	Default
`repo_name`	`str`	Either a local directory path containing saved model files, or a Hugging Face Hub repository name (e.g., 'username/model-name').	required
`access_token`	`Optional[str]`	Hugging Face API token for accessing private repositories. Only required for private models. Defaults to None.	`None`
`version`	`Optional[str]`	Specific version/tag of the model to load from Hub. Defaults to None (loads latest version).	`None`

Returns:

Type	Description
`PPO`	Loaded PPO agent instance ready for inference or further training.

Examples:

Load from local path:

>>> agent = PPO.from_pretrained('./saved_models/my_agent')

Load from Hugging Face Hub:

>>> agent = PPO.from_pretrained('username/ppo-pendulum-v1')

Load specific version with auth:

>>> agent = PPO.from_pretrained(
...     'username/private-model',
...     access_token='hf_xxx',
...     version='v1.0.0'
... )

`Actor(input_dim, out_dim, hidden_dim=256, *, log_std_min=-20.0, log_std_max=0.0)` ¶

Bases: Module

Policy network for PPO algorithm with continuous action spaces.

The actor implements a Gaussian policy that outputs mean and standard deviation for continuous action distributions.

Architecture

Two hidden layers with ReLU activation
Separate output heads for mean (mu) and log std (delta)
Tanh activation on outputs to bound actions

Attributes:

Name	Type	Description
`d1`		First hidden layer.
`d2`		Second hidden layer.
`mu`		Mean output layer for action distribution.
`delta`		Log standard deviation output layer.
`log_std_min`		Minimum allowed log std value.
`log_std_max`		Maximum allowed log std value.

Initialize actor network.

Parameters:

Name	Type	Description	Default
`input_dim`	`int`	Dimension of input observations.	required
`out_dim`	`int`	Dimension of action space.	required
`hidden_dim`	`int`	Number of units in hidden layers. Defaults to 256.	`256`

`forward(input_data)` ¶

Perform forward pass to compute action distribution and sample action.

Parameters:

Name	Type	Description	Default
`input_data`	`Tensor`	Input observation tensor of shape (batch_size, input_dim).	required

Returns:

Type	Description
`Tuple[Tensor, Normal]`	Tuple containing: - Sampled action tensor of shape (batch_size, action_dim). - Normal distribution object representing the action distribution.

`Critic(input_dim, hidden_dim=256)` ¶

Bases: Module

Value function network for PPO algorithm.

The critic estimates the expected return (value) from a given state, which is used to compute advantages for policy updates.

Architecture

Two hidden layers with ReLU activation
Final linear layer outputs scalar value estimate

Attributes:

Name	Type	Description
`d1`		First hidden layer.
`d2`		Second hidden layer.
`v`		Value output layer.

Initialize critic network.

Parameters:

Name	Type	Description	Default
`input_dim`	`int`	Dimension of input observations.	required
`hidden_dim`	`int`	Number of units in hidden layers. Defaults to 256.	`256`

`forward(input_data)` ¶

Perform forward pass to compute state value.

Parameters:

Name	Type	Description	Default
`input_data`	`Tensor`	Input observation tensor of shape (batch_size, input_dim).	required

Returns:

Type	Description
`Tensor`	Value estimates of shape (batch_size, 1).

`ppo_iter(epoch, mini_batch_size, states, actions, log_probs, returns, advantages, rewards, values)` ¶

Create mini-batch iterator for PPO training with shuffled indices.

This function generates mini-batches by randomly shuffling the data indices for each epoch, which helps prevent overfitting and improves generalization.

Parameters:

Name	Type	Description	Default
`epoch`	`int`	Number of epochs to iterate over the data.	required
`mini_batch_size`	`int`	Size of each mini-batch.	required
`states`	`Tensor`	State tensor of shape (batch_size, state_dim).	required
`actions`	`Tensor`	Action tensor of shape (batch_size, action_dim).	required
`log_probs`	`Tensor`	Log probability tensor of shape (batch_size, 1).	required
`returns`	`Tensor`	Return tensor of shape (batch_size, 1).	required
`advantages`	`Tensor`	Advantage tensor of shape (batch_size, 1).	required
`rewards`	`Tensor`	Reward tensor of shape (batch_size, 1).	required
`values`	`Tensor`	Old value estimates of shape (batch_size, 1).	required

Yields:

Type	Description
`Tensor`	Tuple containing mini-batches of (states, actions, log_probs, returns,
`Tensor`	advantages, rewards, values) for each iteration.

References¶

Proximal Policy Optimization Algorithms

Tested on¶

Unity environment

Proximal Policy Optimization (PPO)¶

Components¶

Theory¶

Implementation details¶

Training pseudocode¶

Hyperparameters and mapping¶

Quick start¶

Practical tips¶

Auxiliary Tasks¶

How it works¶

Usage¶

Unified training interface¶

API reference¶

act(state, deterministic=False) ¶

actor_loss(probs, entropy, actions, adv, old_probs) ¶

eval() ¶

close() ¶

learn(states, actions, adv, old_probs, discnt_rewards, rewards, old_values) ¶

test_reward() ¶

preprocess1(states, actions, rewards, dones, values, probs, gamma) ¶

train(num_episodes=None, *, max_steps=None, save_best=False, save_path=None, verbose=True, **kwargs) ¶

get_param_env() ¶

save(path=None) ¶

Saves to: /path/to/models/Oct05_14-30-45_PPO/¶

__load(path) classmethod ¶

from_pretrained(repo_name, access_token=None, version=None) classmethod ¶

Actor(input_dim, out_dim, hidden_dim=256, *, log_std_min=-20.0, log_std_max=0.0) ¶

forward(input_data) ¶

Critic(input_dim, hidden_dim=256) ¶

forward(input_data) ¶

ppo_iter(epoch, mini_batch_size, states, actions, log_probs, returns, advantages, rewards, values) ¶

References¶

Tested on¶

`act(state, deterministic=False)` ¶

`actor_loss(probs, entropy, actions, adv, old_probs)` ¶

`eval()` ¶

`close()` ¶

`learn(states, actions, adv, old_probs, discnt_rewards, rewards, old_values)` ¶

`test_reward()` ¶

`preprocess1(states, actions, rewards, dones, values, probs, gamma)` ¶

`train(num_episodes=None, *, max_steps=None, save_best=False, save_path=None, verbose=True, **kwargs)` ¶

`get_param_env()` ¶

`save(path=None)` ¶

`__load(path)` `classmethod` ¶

`from_pretrained(repo_name, access_token=None, version=None)` `classmethod` ¶

`Actor(input_dim, out_dim, hidden_dim=256, *, log_std_min=-20.0, log_std_max=0.0)` ¶

`forward(input_data)` ¶

`Critic(input_dim, hidden_dim=256)` ¶

`forward(input_data)` ¶

`ppo_iter(epoch, mini_batch_size, states, actions, log_probs, returns, advantages, rewards, values)` ¶