近端策略優化算法PPO的核心概念和PyTorch實現詳解

netouch 2025-08-29 發布于北京

展開全文

近端策略優化(Proximal Policy Optimization, PPO)作為強化學習領域的重要算法，在眾多實際應用中展現出卓越的性能。本文將詳細介紹PPO算法的核心原理，并提供完整的PyTorch實現方案。

PPO算法在強化學習任務中具有顯著優勢：即使未經過精細的超參數調優，也能在Atari游戲環境等復雜場景中取得優異表現。該算法不僅在傳統強化學習任務中表現出色，還被廣泛應用于大語言模型的對齊優化過程。因此掌握PPO算法對于深入理解現代強化學習技術具有重要意義。

本文將通過Lunar Lander環境演示PPO算法的完整實現過程。文章重點闡述算法的核心概念和實現細節，通過適當的修改，本實現方案可擴展至其他強化學習環境。本文專注于高層次的算法理解，為讀者提供系統性的技術資源。

打開今日頭條查看圖片詳情

PPO算法核心組件

PPO算法由四個核心組件構成：環境交互模塊、智能體決策系統、優勢函數計算以及策略更新裁剪機制。每個組件在算法整體架構中發揮著關鍵作用。

環境交互模塊

環境是智能體進行學習和決策的載體。這里我們選用Lunar Lander作為測試環境，這是一個二維物理模擬場景，要求著陸器在月球表面的指定區域安全著陸。環境模塊負責提供狀態觀測信息，接收智能體的動作指令，并根據任務完成情況反饋相應的獎勵信號。

有效的環境設計和獎勵函數是成功訓練的基礎。智能體需要從環境中獲取充分的狀態信息以做出合理決策，同時需要通過明確的獎勵信號了解其行為的優劣程度。獎勵信號的質量直接影響智能體的學習效率和最終性能。

智能體決策系統

PPO采用演員-評論家(Actor-Critic)架構設計智能體決策系統。演員網絡負責根據當前狀態選擇最優動作，而評論家網絡則評估當前狀態的價值期望。

演員網絡的作用類似于決策執行者，根據觀測到的環境狀態輸出動作概率分布。評論家網絡則充當價值評估者，預測在當前狀態下能夠獲得的累積獎勵期望值。當評論家的價值估計出現偏差時，通常表明智能體的策略仍有改進空間。

優勢函數

優勢函數用于量化特定動作相對于評論家期望值的優劣程度。正優勢值表示該動作的表現優于期望，應當增強此類行為的選擇概率；負優勢值則表示表現不佳，需要降低此類行為的選擇概率。

相比直接使用原始獎勵值，優勢函數能夠提供更穩定的訓練信號。智能體僅在實際表現與預期之間存在顯著差異時才進行大幅度的策略調整，這種機制有效避免了訓練過程中的不必要波動。本實現采用廣義優勢估計(Generalized Advantage Estimation, GAE)方法計算優勢值。

PPO策略更新裁剪機制

PPO算法的核心創新在于引入策略更新的裁剪機制，這是其相對于傳統策略梯度方法的關鍵改進。在強化學習訓練過程中，過大的策略更新可能導致訓練失穩，使智能體突然丟失已學習的有效策略。

這種現象可以類比為在狹窄山脊上行走的登山者：如果步伐過大或方向偏離，很容易失足跌落深谷，重新攀登將耗費大量時間和精力。PPO通過實施裁剪約束，確保每次策略更新都在安全范圍內進行，保持學習過程的穩定性和連續性。

依賴庫安裝與環境配置

本實現需要安裝gymnasium庫及其相關依賴來運行Lunar Lander環境，PyTorch用于神經網絡的構建和訓練，以及tensordict庫來管理訓練數據。tensordict是一個先進的數據管理工具，允許將PyTorch張量作為字典元素進行操作，支持通過鍵值索引和檢索數據項。這種設計使得數據管理更加靈活高效，同時保持張量和tensordict在GPU上的計算能力，與PyTorch工作流程無縫集成。

!pip install swig gymnasium torch tensordict pyvirtualdisplay
!pip install 'gymnasium[box2d]'

為確保實驗結果的可重現性，我們設置統一的隨機種子：

import random
import torch
import numpy as np
seed = 777
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.backends.cudnn.deterministic = True

需要注意的是，由于向量化環境初始化過程中存在的隨機性因素，完全的結果重現仍然面臨技術挑戰。雖然代碼能夠正常運行且智能體訓練過程穩定，但具體的數值結果在不同運行之間可能存在差異。

Lunar Lander環境配置

Lunar Lander是gymnasium庫中的經典強化學習環境，任務目標是控制著陸器在二維空間中導航，最終在月球表面的指定著陸區域安全降落。智能體根據著陸過程中的姿態穩定性、著陸柔和度以及任務完成速度獲得相應獎勵。著陸器具備四種控制動作，環境在每個時間步提供八維狀態向量描述著陸器的當前狀態。詳細的環境說明可參考官方文檔。

為提高訓練數據收集效率，我們創建10個并行的Lunar Lander環境實例。同時配置獨立的評估環境，用于記錄智能體在各個訓練階段的表現視頻，便于直觀評估訓練效果。

import gymnasium as gym
from gymnasium.wrappers import RecordVideo

# 創建將運行10個模擬的向量化環境
env_name = 'LunarLander-v3'
num_envs = 10
envs = gym.make_vec(env_name, num_envs=num_envs, vectorization_mode='sync')

# 創建我們的評估錄制環境，用于當我們
# 想要測試我們的智能體在訓練的各個階段表現如何時使用
env = gym.make(env_name, render_mode='rgb_array')
trigger = lambda t: True

recording_output_directory = './checkpoint_videos'
recording_env = RecordVideo(env,
recording_output_directory, episode_trigger=trigger)

環境交互輔助類設計

為提高代碼的模塊化程度和可維護性，我們設計了專門的環境交互輔助類，負責處理訓練數據收集和評估過程。該類封裝了與環境的所有交互操作，包括訓練環境、評估環境、智能體以及計算設備的管理。

主要功能包括：數據rollout收集用于獲取訓練樣本，以及評估rollout執行用于性能測試和視頻記錄。評估過程中的視頻錄制功能由RecordVideo包裝器自動完成。

from tensordict import TensorDict
import torch

# 關于這個類的快速說明
# 它假設訓練環境在終端狀態時自動重置，例如向量化環境
# 并且它假設評估環境不會自動重置
# 我還將num_steps_per_rollout固定為收集器初始化時的值，但這可以在
# get rollout函數中參數化
class PPORolloutCollector:
def __init__(self, agent, envs, num_steps_per_rollout, device, eval_env=None):
self.agent = agent
self.envs = envs
self.num_envs = envs.num_envs

self.num_steps_per_rollout = num_steps_per_rollout
self.device = device
self.eval_env = eval_env

self.obs_shape =
envs.single_observation_space.shape
self.action_shape =
envs.single_action_space.shape
self.initial_buffer_shape = (num_steps_per_rollout, envs.num_envs)

# 繼續重置環境并存儲觀察
# 并將next_done設置為false(我們假設環境不能以終端狀態開始)
obs, _ = envs.reset()
self.next_obs = torch.Tensor(obs).to(device)
self.next_done = torch.zeros(num_envs).to(device)

# 創建一個空緩沖區，將保存觀察、來自智能體的動作、該動作的對數概率
# 當前觀察的評論家估計、從環境中獲得的采取動作的實際獎勵，以及
# 動作是否導致終端狀態
# 我們故意選擇不為每個觀察記錄'下一狀態'，這基本上會使
# 緩沖區的大小翻倍，由于我們按順序收集觀察并且在這里不打亂
# 任何下游操作都可以只使用數組中的下一個值作為下一狀態
def _create_buffer(self):
return TensorDict({
'obs': torch.zeros(self.initial_buffer_shape + self.obs_shape).to(self.device),
'actions': torch.zeros(self.initial_buffer_shape + self.action_shape).to(self.device),
'log_probs': torch.zeros(self.initial_buffer_shape).to(self.device),
'rewards': torch.zeros(self.initial_buffer_shape).to(self.device),
'dones': torch.zeros(self.initial_buffer_shape).to(self.device),
'critic_values': torch.zeros(self.initial_buffer_shape).to(self.device),
})

# 將收集num_steps_per_rollout個觀察對我們的訓練環境的函數
def get_next_rollout(self):
buffer = self._create_buffer()

# 獲取最后記錄的觀察以及該觀察是否為終端
next_obs = self.next_obs
next_done = self.next_done

# 收集rollout
for t in range(
self.num_steps_per_rollout):
# 記錄當前觀察和終端狀態
buffer['obs'][t] = next_obs
buffer['dones'][t] = next_done

# 查詢智能體下一個動作、該動作的對數概率和評論家估計
with torch.no_grad():
action, log_prob, entropy =
self.agent.get_actor_values(next_obs)
critic_value =
self.agent.get_critic_value(next_obs)

# 記錄值
buffer['actions'][t] = action
buffer['log_probs'][t] = log_prob
buffer['critic_values'][t] = critic_value.flatten()

# 執行動作
next_obs, reward, terminations, truncations, infos = envs.step(action.cpu().numpy())

# 形狀化并存儲獎勵
reward = torch.tensor(reward).to(self.device).view(-1)
buffer['rewards'][t] = reward

# 一些環境會終止(意味著智能體處于最終狀態)，
# 其他會截斷(例如達到時間限制但不在終端狀態)
# 這些是重要的區別，但對我們的目的來說意味著同樣的事情，模擬結束了
# 所以如果任一為真且模擬重置，我們將next done設置為true
next_done = np.logical_or(terminations, truncations)
# 為下一輪存儲下一個obs和next done
next_obs, next_done = torch.Tensor(next_obs).to(self.device), torch.Tensor(next_done).to(self.device)

# 在緩沖區中存儲下一個obs和next done
# 這是為了在稍后計算優勢時處理邊緣情況
buffer['next_obs'] = next_obs
buffer['next_done'] = next_done
# 我們還需要這個下一狀態的評論家估計，我們將使用它來引導最終狀態的獎勵
# 當我們計算gae時
with torch.no_grad():
buffer['next_value'] =
self.agent.get_critic_value(next_obs).reshape(1, -1)
self.next_obs = next_obs
self.next_done = next_done
return buffer

# 用于在我們的環境上評估智能體，將運行整個模擬直到終止
# 與上面非常相似，唯一的區別是我們將手動檢查環境是否終止
# 然后結束循環
# 將返回一個普通的python字典，包含獎勵、熵、獎勵平均值、我們智能體的平均熵，以及每次運行的總獎勵
def run_eval_rollout(self, num_episodes: int = 5):
assert self.eval_env is not None, 'No eval_env provided.'

rewards_per_timestep = []
entropies_per_timestep = []
final_rewards = []
total_entropies = []

for _ in range(num_episodes):
obs, _ = self.eval_env.reset()
obs = torch.tensor(obs, device=self.device).unsqueeze(0)
done = False

episode_rewards = []
episode_entropies = []

while not done:
with torch.no_grad():
action, _, entropy =
self.agent.get_actor_values(obs)
action = action.squeeze()
obs_np, reward, term, trunc, _ = self.eval_env.step(action.cpu().numpy())
done = term or trunc

obs = torch.tensor(obs_np, device=self.device).unsqueeze(0)
episode_rewards.append(float(reward))
episode_entropies.append(entropy.item())

rewards_per_timestep.append(episode_rewards)

entropies_per_timestep.append(episode_entropies)
final_rewards.append(sum(episode_rewards))
total_entropies.append(sum(episode_entropies) / len(episode_entropies))

return {
'rewards_per_timestep': rewards_per_timestep,
'average_reward_per_run': sum(final_rewards) / len(final_rewards),
'average_entropy_per_run': sum(total_entropies) / len(total_entropies),
'entropies_per_timestep': entropies_per_timestep,
'final_rewards': final_rewards,
}

智能體網絡架構設計

本實現采用演員-評論家架構構建智能體決策系統。該架構的一個重要特點是演員網絡和評論家網絡可以完全獨立，也可以共享部分網絡層。在復雜環境中，通常建議讓兩個網絡共享前期特征提取層，使得雙方都能從對方的損失更新中受益。然而，在本示例中，為了清晰展示兩個網絡的獨立性，我們采用完全分離的網絡結構。

借助PyTorch模塊的參數管理機制，我們無需將兩個網絡實現為獨立的類，可以在單一Agent類中統一管理。

import torch.nn as nn
from
torch.distributions.categorical import Categorical

# 一個初始化輔助函數。在強化學習問題中，正交初始化
# 網絡的權重可能會有幫助，這意味著每層的輸出盡可能
# 不相關。這可以通過幫助梯度更好地流動和避免可能
# 從樸素初始化中產生的相關特征問題來改善訓練穩定性和效率。
# 可以把這想象成確保網絡不會以交叉的線路開始。
# 這個函數取自CleanRL的PPO實現。
def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
torch.nn.init.orthogonal_(layer.weight, std)
torch.nn.init.constant_(layer.bias, bias_const)
return layer

class Agent(nn.Module):
def __init__(self, envs):
super().__init__()
# 我們的輸入數組有多大？
# 注意如果我們做的是圖像觀察之類的
# 我們需要重新設計這個，但lunar lander只提供
# 代表世界狀態的數字數組
input_shape = np.array(
envs.single_observation_space.shape).prod()

# 我們的智能體將有多少個動作？
action_shape =
envs.single_action_space.n

# 創建一個3層評論家，它將預測
# 我們的智能體在給定觀察下預期收到的總獎勵
self.critic = nn.Sequential(
layer_init(nn.Linear(input_shape, 64)),
nn.Tanh(),
layer_init(nn.Linear(64, 64)),
nn.Tanh(),
layer_init(nn.Linear(64,1), std=1.0)
)

# 創建一個3層演員，它將輸出一個概率分布
# 表示在給定觀察下最好采取哪個動作的概率
self.actor = nn.Sequential(
layer_init(nn.Linear(input_shape, 64)),
nn.Tanh(),
layer_init(nn.Linear(64, 64)),
nn.Tanh(),
layer_init(nn.Linear(64, action_shape), std=1.0)
)

def save(self, path):
torch.save(self.state_dict(), path)

def load(self, path):
self.load_state_dict(torch.load(path))

def get_critic_value(self, x):
return self.critic(x)

def get_actor_values(self, x, action=None):
logits = self.actor(x)
probs = Categorical(logits=logits)
if action is None:
action = probs.sample()
return action, probs.log_prob(action), probs.entropy()

網絡初始化采用正交初始化策略，這是強化學習領域的常見實踐。正交初始化確保各網絡層輸出之間的去相關性，有助于改善梯度流動并避免特征相關性問題，從而提升訓練的穩定性和效率。這種初始化方法可以形象地理解為確保網絡不會以'交叉連線'的狀態開始訓練。

創建智能體實例并配置計算設備：

# 如果可用，使用gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
agent = Agent(envs).to(device)
print('using device: ', device)

配置優化器，選擇Adam優化算法并設置學習率為2.5e-4，這是PPO訓練的常用起始參數：

import torch.optim as optim
learning_rate = 2.5e-4
optimizer = optim.Adam(agent.parameters(), lr=learning_rate, eps=1e-5)

優勢函數理論與實現

優勢函數用于量化特定動作相對于評論家價值估計的表現差異。當智能體在游戲中的典型得分為15分時，這個分數就成為性能基準。如果智能體采用新策略獲得17分，則優勢值為2，表明新策略的表現超出預期。

采用優勢函數替代原始獎勵信號能夠穩定訓練過程，防止智能體滿足于次優策略。當我們使用優勢校正值訓練評論家網絡時：

打開今日頭條查看圖片詳情

相比直接預測原始回報，這種方法通常能實現更快的收斂速度和更穩定的評論家更新過程。

廣義優勢估計(GAE)算法

廣義優勢估計是計算優勢函數的經典方法。該方法不直接比較完整軌跡回報與價值估計(這種方式存在較大噪聲)，而是采用時間差分殘差的平滑求和。這種設計在偏差和方差之間取得平衡，即使在評論家網絡估計不夠精確的情況下也能保持實用性。

時間差分殘差的定義為：

打開今日頭條查看圖片詳情

基于TD殘差，GAE的遞歸定義為：

打開今日頭條查看圖片詳情

在實際實現中，我們沿著軌跡反向計算優勢值，在每個回合結束時將累積值重置為零。這種方法能夠產生穩定且低方差的優勢估計，可直接用于PPO算法的策略更新。

# 我們的compute_gae函數將接收緩沖區和2個超參數，gamma和gae_lambda，它們控制未來結果對優勢和回報的影響程度
# 我們的函數將向緩沖區添加2個新值：
# 1. 'advantages'，即每個時間步的gae計算優勢，將用于幫助調整我們智能體的演員部分
# 2. 'returns'，即基于優勢調整的評論家值，可以將其視為使用優勢將評論家的值推向實際獎勵
# returns將用于訓練評論家做出更好的價值預測
def compute_gae(buffer, gamma=0.99, gae_lambda=0.95):
next_done = buffer['next_done']
critic_values = buffer['critic_values']
rewards = buffer['rewards']
dones = buffer['dones']
num_steps_per_rollout = len(rewards)
last_step_critic_values = buffer['next_value']

advantages = torch.zeros_like(rewards).to(device)
last_gae_value = 0
with torch.no_grad():

# 由于gae是前瞻的并使用未來值，我們可以聰明一點并向后工作
# 我們將首先計算最后的優勢，然后當我們計算倒數第二個時，我們將有遞歸的
# 未來值保存，以此類推
for t in reversed(range(num_steps_per_rollout)):
# 檢查我們是否在最后一個時間步，如果是，則使用last_step_critic_values，這是對
# 來自我們環境的最后狀態的評論家預測，我們還沒有實際對其采取行動的狀態。否則我們向前
# 在緩沖區中查看以獲得下一個值
if t == num_steps_per_rollout-1:
next_nonterminal = 1.0 - next_done
next_values = last_step_critic_values
else:
# 否則我們使用緩沖區中為下一個時間步存儲的值
next_nonterminal = 1.0 - dones[t+1]
next_values = critic_values[t+1]

# 計算未來獎勵，由gamma折扣
# 如果狀態是終端狀態，我們也將未來獎勵歸零

discounted_estimated_future_reward = gamma * next_values * next_nonterminal

# 接下來我們計算此時間步的時間差分(TD)誤差
td_error = rewards[t] +
discounted_estimated_future_reward - critic_values[t]

# 現在我們遞歸計算優勢
current_advantage = td_error + gamma * gae_lambda * next_nonterminal * last_gae_value

# 記錄這個優勢
advantages[t] = current_advantage
# 由于我們向后計算優勢，我們可以存儲當前優勢并將其用作
# 下一步的未來優勢！
last_gae_value = current_advantage
# 存儲我們計算的優勢
buffer['advantages'] = advantages
# 并計算和存儲我們的回報
returns = advantages + critic_values
buffer['returns'] = returns
return buffer

PPO策略更新機制

在完成優勢計算后，我們實施PPO的核心策略更新過程。PPO通過引入策略更新裁剪機制，有效避免了傳統策略梯度方法中可能出現的破壞性大幅更新。

首先定義新策略πθ與舊策略πθ_old之間的概率比率：

打開今日頭條查看圖片詳情

該比率反映了新策略相對于舊策略選擇相同動作的概率變化程度。

裁剪代理目標函數

PPO的核心目標函數用于指導演員網絡的優化過程：

打開今日頭條查看圖片詳情

該函數的設計目標是控制策略更新幅度，確保更新方向基于優勢信號的指導，同時避免過度偏離當前策略。

價值函數損失

評論家網絡采用均方誤差損失函數，目標是最小化價值預測與GAE目標值之間的差異：

打開今日頭條查看圖片詳情

熵正則化項

熵值反映模型決策的確定性程度。高熵表示模型在多個動作間的概率分布較為均勻，低熵則表示模型對特定動作具有強烈偏好。為促進探索行為，我們在損失函數中加入熵正則化項。當模型的策略熵過低時，該項會施加懲罰，防止策略過早收斂到局部最優解。

打開今日頭條查看圖片詳情

綜合損失函數

PPO的最終損失函數是上述三個組件的加權組合：

打開今日頭條查看圖片詳情

其中，系數c1控制評論家網絡的學習速率，c2控制探索行為的強度。

def run_ppo_update_step(update_epochs_per_rollout, agent, optimizer, buffer, batch_size, minibatch_size,
clip_coef=0.2, ent_coef=0.01, vf_coef=0.5,

should_normalize_advantages=True, max_grad_norm=0.5):
# 為訓練準備緩沖區
b_obs = buffer['obs'].reshape((-1,) +
envs.single_observation_space.shape)
b_logprobs = buffer['log_probs'].reshape(-1)
b_actions = buffer['actions'].reshape((-1,) +
envs.single_action_space.shape)
b_advantages = buffer['advantages'].reshape(-1)
b_returns = buffer['returns'].reshape(-1)
b_inds = np.arange(batch_size)

for epoch in range(update_epochs_per_rollout):
np.random.shuffle(b_inds)

for start in range(0, batch_size, minibatch_size):
end = start + minibatch_size

mb_inds = b_inds[start:end]

_, new_log_prob, entropy = agent.get_actor_values(b_obs[mb_inds], action = b_actions[mb_inds])
new_critic_value = agent.get_critic_value(b_obs[mb_inds])

log_ratio = new_log_prob - b_logprobs[mb_inds]
ratio = log_ratio.exp()

mb_advantages = b_advantages[mb_inds]
# 將優勢標準化作為小批量的一部分
# 在小批量級別而不是整個緩沖區進行標準化通常有助于
# 穩定訓練并導致更好的結果！
if
should_normalize_advantages:
mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8)

# 策略損失
pg_loss1 = -mb_advantages * ratio
pg_loss2 = -mb_advantages * torch.clamp(ratio, 1 - clip_coef, 1 + clip_coef)
pg_loss = torch.max(pg_loss1, pg_loss2).mean()

# 價值損失
new_critic_value = new_critic_value.view(-1)
v_loss = 0.5 * ((new_critic_value - b_returns[mb_inds]) ** 2).mean()

entropy_loss = entropy.mean()
loss = pg_loss - ent_coef * entropy_loss + v_loss * vf_coef

optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(agent.parameters(), max_grad_norm)
optimizer.step()

完整的PPO訓練流程

下面將整合所有組件，構建完整的PPO訓練循環。該循環包含數據收集、優勢計算、策略更新以及定期評估等關鍵步驟。

import time
import os

# 定義一些超參數
num_training_iterations = 1000 # 我們應該在停止之前做多少次訓練迭代？
# 實際上，1000可能過度了，特別是對于這個環境，
# 但我們無論如何都會保持在1k！

num_steps_per_rollout = 1000 # 每次rollout我們應該做多少個時間步？
minibatch_size = 256 # 在更新期間我們的小批量應該有多大？
# 對于PPO，每次迭代我們可以有多次更新，盡管如果我們有太多，我們更可能達到裁剪函數的最大值并浪費計算
# 記住，每次迭代只采取小步驟，所以每個rollout 4個epochs(在整個rollout上訓練4次)應該沒問題
update_epochs_per_rollout = 4

iterations_per_eval = 100 # 我們應該多久做一次評估？在我們的情況下，我們將訓練100次迭代，然后執行評估步驟
runs_per_eval = 5 # 評估時我們應該運行模擬多少次？單次運行可能很幸運，所以我們想對幾次進行平均

# 定義保存我們智能體檢查點的位置并確保目錄存在

output_model_checkpoints_base = './checkpoints'
os.makedirs(
output_model_checkpoints_base, exist_ok=True)

# 初始化rollout收集器
rollout_collector = PPORolloutCollector(
agent=agent,
envs=envs,
num_steps_per_rollout=num_steps_per_rollout,
device=device,
eval_env=recording_env
)

# 跟蹤運行需要多長時間！
elapsed_timesteps = 0
start_time = time.time()

# 創建一些列表來記錄每個評估步驟的平均獎勵和平均熵
# 我們稍后會在notebook中繪制這些！
average_eval_reward = []
average_eval_entropy = []

for iteration in range(num_training_iterations):
agent.train()
# 我們將為每次迭代打印一個臨時行，以在訓練期間跟蹤它的進度，并驗證它仍然活著
print(f'\rtraining iteration: {iteration}/{num_training_iterations} - runtime: {time.time() - start_time:.2f}', end='', flush=True)

# 為這個訓練迭代收集rollout
buffer =
rollout_collector.get_next_rollout()

# 為我們rollout中的觀察計算GAE優勢
buffer = compute_gae(buffer)

# 使用rollout數據運行ppo更新步驟
run_ppo_update_step(
update_epochs_per_rollout,
agent,
optimizer,
buffer,
batch_size=num_envs * num_steps_per_rollout,
minibatch_size=minibatch_size
)

# 跟蹤到目前為止我們已經模擬了多少個時間步
elapsed_timesteps += num_envs * num_steps_per_rollout

# 定期對我們的智能體進行檢查點和評估運行，
# 我們也會為最終迭代手動觸發這個
if iteration % iterations_per_eval == 0 or iteration == num_training_iterations - 1:
with torch.no_grad():
agent.eval()

# 為我們的評估rollout收集評估統計數據
eval_stats =
rollout_collector.run_eval_rollout(num_episodes=runs_per_eval)
avg_reward = eval_stats['average_reward_per_run']
avg_entropy = eval_stats['average_entropy_per_run']

average_eval_reward.append(avg_reward)

average_eval_entropy.append(avg_entropy)

# 保存我們的智能體，以防我們需要在將來的運行中恢復
ckpt_path = os.path.join(
output_model_checkpoints_base, f'checkpoint_{iteration}.pth')
agent.save(ckpt_path)
agent.train()

# 打印此評估運行的統計數據
elapsed = time.time() - start_time
print(f'\riteration {iteration}/{num_training_iterations} | timesteps: {elapsed_timesteps} | '
f'avg_eval_reward: {avg_reward:.2f} | avg_eval_entropy: {avg_entropy:.2f} | runtime: {elapsed:.2f}s')

# 將我們的智能體保存到最終模型
final_ckpt_path = os.path.join(
output_model_checkpoints_base, '!final_model.pth')
agent.save(final_ckpt_path)

訓練過程包括環境數據收集、GAE優勢計算、PPO策略更新以及定期性能評估。評估階段會生成視頻文件用于可視化分析，同時記錄平均獎勵和熵值變化以監控訓練進展。

針對本示例環境，1000次訓練迭代可能超出實際需求，但這個設置能夠確保充分的策略收斂。在實際應用中，可根據具體任務需求調整訓練輪數。

實驗結果如下所示：

打開今日頭條查看圖片詳情

訓練結果分析與可視化

通過繪制訓練過程中的關鍵指標，我們可以直觀地評估智能體的學習進展。主要關注兩個指標：各代次的平均獎勵值以及對應的策略熵變化。

import matplotlib.pyplot as plt
import numpy as np

generations = [i for i in range(len(average_eval_entropy))]

# 創建一個圖形和兩個排列在1行2列中的子圖
fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, figsize=(10, 4))

# 在第一個子圖(ax1)上繪制平均獎勵
ax1.plot(generations, average_eval_reward, label='Avg. Eval Rewards', color='blue')
ax1.set_title('Avg. Eval Rewards')
ax1.set_xlabel('Generations')
ax1.set_ylabel('Avg. Reward')
ax1.grid(True)

# 在第二個子圖(ax2)上繪制平均熵
ax2.plot(generations, average_eval_entropy, label='Avg. Eval Entropy', color='red')
ax2.set_title('Avg. Eval Entropy')
ax2.set_xlabel('Generations')
ax2.set_ylabel('Avg. Entropy')
ax2.grid(True)

# 調整布局以防止標題/標簽重疊
plt.tight_layout()

# 顯示圖表
plt.show()

打開今日頭條查看圖片詳情

從訓練曲線可以觀察到預期的學習模式：智能體獲得的累積獎勵隨訓練進展而穩步提升，同時策略熵呈現下降趨勢。需要注意的是，由于PPO損失函數中包含熵正則化項，策略熵不會降至過低水平，這種設計保證了智能體在優化過程中保持適度的探索能力。熵值的下降反映了策略的逐步成熟和決策確定性的提高。

智能體行為可視化

除了數值指標分析，直觀觀察智能體的實際執行行為同樣重要。通過配置的評估環境，系統會自動將智能體的執行過程錄制為MP4視頻文件。下面展示如何檢索和展示這些訓練記錄。

首先定位最新生成的評估視頻文件：

import os
import re

# 這將提供給定目錄中我們錄制環境產生的最新文件
def
get_latest_episode_video_file(directory):
# 正則表達式匹配文件格式并捕獲片段編號
pattern = re.compile(r'rl-video-episode-(\d+)\.mp4')
latest_file = None
highest_episode = -1

# 搜索目錄中的文件
for filename in os.listdir(directory):
match = pattern.match(filename)
if match:
episode_number = int(match.group(1)) # 提取片段編號
# 檢查這個片段編號是否是找到的最高的
if episode_number > highest_episode:
highest_episode = episode_number
latest_file = os.path.join(directory, filename) # 存儲完整路徑

return latest_file

# 獲取最新錄制的文件路徑
latest_eval_recording =
get_latest_episode_video_file(
recording_output_directory)

實現視頻嵌入功能，將錄制的行為視頻直接集成到Jupyter環境中：

import io
import base64

from IPython import display
from IPython.display import HTML
# 這個函數將接收視頻文件的位置，然后
# 使用虛擬顯示將視頻嵌入到notebook中
def embed_video(video_file):
# 打開并從視頻中讀取原始數據
video_data = io.open(video_file, 'r+b').read()
# 現在我們必須將數據編碼為base64才能與
# 虛擬顯示一起工作
encoded_data = base64.b64encode(video_data)
# 現在我們使用display.display函數獲取一些html
# 和編碼數據，并將html嵌入到notebook中！
display.display(HTML(data='''<video alt='test' autoplay
loop controls style='height: 400px;'>
<source src='data:video/mp4;base64,{0}' type='video/mp4' />
</video>'''.format(encoded_data.decode('ascii'))))
embed_video(latest_eval_recording)