Proximal Policy Optimization (PPO)¶

PPO — надёжный policy‑gradient метод, сочетающий простоту реализации и стабильность обучения. В нашей реализации актор и критик обучаются на батчах собранных роллаутов, используется клиппированный суррогат, энтропия политики и оценка преимуществ с обобщённой ошибкой (GAE‑подобная).

Компоненты¶

Актор (гауссовская политика): параметры \(\mu, \sigma\) → распределение \(\mathcal{N}(\mu, \sigma^2)\)
Критик: скалярная оценка \(V(s)\)
Сбор опыта: роллаут длины rollout_len с записью \(s,a,\log\pi(a|s), r, d, V(s)\)
Обучение: мини‑батчи по несколько эпох num_epochs с клиппингом вероятностных отношений

Теория¶

Отношение вероятностей:

\[ 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) \]

Клиппированный суррогат (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] \]

Потеря критика (Value):

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

Энтропийная регуляризация (стохастичность политики):

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

Полная цель: \(\mathcal{L} = \mathcal{L}_\text{actor} + \mathcal{L}_\text{critic} + \mathcal{L}_\text{entropy}\)
Преимущество (GAE‑подобное): в preprocess1 возвращается \(\text{return} = V + \sum\gamma\lambda\,\delta\), а \(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} \]

Детали реализации¶

Политика: Actor.forward(..., continous_actions=True) выводит mu = tanh(Wx) и log_std = tanh(Wx) с последующим линеарным растяжением в диапазон [log_std_min, log_std_max]; \(\sigma = e^{\log \sigma}\). Действие семплируется из Normal(mu, sigma).
Отношения вероятностей: берутся через разность лог‑плотностей new_probs - old_probs, затем экспонента (torch.exp) — это численно устойчивее, чем делить плотности напрямую.
Энтропия: в коде в actor_loss подаётся отрицательная энтропия -new_distr.entropy().mean(), а затем добавляется как + entropy_coef * entropy. Эффект равнозначен вычитанию энтропии с коэффициентом (стимулируется стохастичность политики).
GAE и бустрап: в preprocess1 добавляется next_value в values, затем по реверсу считается \(\delta\) и аккумулируется \(g\) с \(\lambda=0.8\); в итоге returns = V + g, advantages = returns - V.
Мини‑батчи: итератор ppo_iter случайно выбирает индексы размера mini_batch_size многократно в течение epoch.
Доп. голова r: актор возвращает ещё и предсказание наград (линия self.r), для вспомогательной задачи auxillary_task (MSE по наградам); по умолчанию в loss не добавляется.

Псевдокод обучения¶

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

Гиперпараметры и соответствие коду¶

clip_pram = ε — порог клиппинга вероятностных отношений
num_epochs, batch_size — количество проходов и размер мини‑батча для обновлений
rollout_len — длина роллаута перед обновлениями
entropy_coef — вес энтропийного члена (учитывая знак в реализации)
actor_lr, critic_lr — скорости обучения оптимизаторов Adam
gamma, lambda(=0.8) — скидка и параметр GAE внутри preprocess1

Быстрый старт¶

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

# Создаём среду (пример — F16)
env = gym.make('LinearLongitudinalF16-v0', number_time_steps=2000)

# Инициализация 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,
)

# Обучение
agent.train()

# Сохранение
agent.save('./runs')

Tip

Для непрерывных действий используем гауссовскую политику; полезно ограничивать log_std (как в коде) и нормировать признаки.

Практические советы¶

Увеличивайте rollout_len для более стабильной оценки преимуществ
Балансируйте clip_pram (обычно 0.1–0.3) и entropy_coef для исследовательности
Несколько эпох (num_epochs) и мелкие batch_size улучшают сходимость, но следите за переобучением

Вспомогательные задачи (Auxiliary Tasks)¶

Реализация PPO включает опциональный механизм вспомогательных задач (auxiliary tasks) для предсказания награды. Этот дополнительный выход помогает агенту формировать более качественные представления состояний, предсказывая ожидаемые награды параллельно с основной оптимизацией политики.

Как это работает¶

Сеть Actor включает дополнительный выходной слой self.r, который предсказывает награду
Вспомогательная ошибка вычисляется как MSE между предсказанной и фактической наградой
Эта ошибка может быть добавлена к основной ошибке PPO через метод auxillary_task в классе Agent

Использование¶

# Вспомогательная задача вычисляется отдельно от основного обучения
aux_loss = agent.auxillary_task(states, rewards)

Вспомогательная задача стимулирует сеть кодировать признаки, релевантные для награды, в скрытых представлениях, что потенциально улучшает эффективность использования выборки и обобщающую способность.

Унифицированный интерфейс обучения¶

PPO следует общему унифицированному API train() из BaseRLModel:

stats = agent.train(
    num_episodes=200,   # необязательно: переопределяет self.max_episodes
    max_steps=1024,     # необязательно: переопределяет self.rollout_len
)

Вызов agent.train() без аргументов также поддерживается — в этом случае используются гиперпараметры, заданные при создании. Обратите внимание, что метод PPO learn(states, actions, adv, old_probs, returns, rewards, old_values) является внутренним помощником, выполняющим один шаг градиентного обновления по батчу, и не затрагивается унифицированным интерфейсом.

Документация API¶

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.

Источники¶

Proximal Policy Optimization Algorithms

Где тестировалось¶

Unity‑среда

Proximal Policy Optimization (PPO)¶

Компоненты¶

Теория¶

Детали реализации¶

Псевдокод обучения¶

Гиперпараметры и соответствие коду¶

Быстрый старт¶

Практические советы¶

Вспомогательные задачи (Auxiliary Tasks)¶

Как это работает¶

Использование¶

Унифицированный интерфейс обучения¶

Документация API¶

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) ¶

Источники¶

Где тестировалось¶

`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)` ¶