151 lines
5.7 KiB
Python
151 lines
5.7 KiB
Python
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
import torch.nn.functional as F
|
|||
|
|
from torch.distributions import Normal
|
|||
|
|
from networks import PolicyNetwork, ValueNetwork
|
|||
|
|
|
|||
|
|
|
|||
|
|
class PPOAgent:
|
|||
|
|
"""
|
|||
|
|
PPO-Clip (Proximal Policy Optimization, Clipped version)
|
|||
|
|
|
|||
|
|
Reference: Schulman et al., "Proximal Policy Optimization Algorithms", 2017.
|
|||
|
|
|
|||
|
|
与 TRPO 的核心区别：
|
|||
|
|
- 不再使用共轭梯度 + 线搜索求解约束优化
|
|||
|
|
- 用 clip(ratio, 1-ε, 1+ε) 限制策略更新幅度
|
|||
|
|
- 支持多 epoch 的小批量更新（每次从经验池中采样）
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
state_dim,
|
|||
|
|
action_dim,
|
|||
|
|
action_bound,
|
|||
|
|
hidden_dim=128,
|
|||
|
|
gamma=0.99,
|
|||
|
|
tau=0.97,
|
|||
|
|
lr=3e-4,
|
|||
|
|
clip_eps=0.2,
|
|||
|
|
k_epochs=10,
|
|||
|
|
minibatch_size=64,
|
|||
|
|
critic_epochs=10,
|
|||
|
|
entropy_coef=0.0,
|
|||
|
|
):
|
|||
|
|
self.gamma = gamma
|
|||
|
|
self.tau = tau
|
|||
|
|
self.clip_eps = clip_eps # PPO-Clip 的裁剪范围 ε
|
|||
|
|
self.k_epochs = k_epochs # 每次更新对同一批数据的训练轮数
|
|||
|
|
self.minibatch_size = minibatch_size
|
|||
|
|
self.entropy_coef = entropy_coef # 熵正则化系数（可选，鼓励探索）
|
|||
|
|
|
|||
|
|
self.actor = PolicyNetwork(state_dim, action_dim, action_bound, hidden_dim)
|
|||
|
|
self.critic = ValueNetwork(state_dim, hidden_dim)
|
|||
|
|
|
|||
|
|
self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=lr)
|
|||
|
|
self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=lr)
|
|||
|
|
|
|||
|
|
def get_action(self, state):
|
|||
|
|
state_tensor = torch.FloatTensor(state).unsqueeze(0)
|
|||
|
|
with torch.no_grad():
|
|||
|
|
dist = self.actor.evaluate(state_tensor)
|
|||
|
|
action = dist.sample()
|
|||
|
|
return action.squeeze(0).numpy()
|
|||
|
|
|
|||
|
|
def _compute_advantages(self, rewards, values, masks):
|
|||
|
|
"""
|
|||
|
|
GAE (Generalized Advantage Estimation)
|
|||
|
|
"""
|
|||
|
|
returns = []
|
|||
|
|
gae = 0
|
|||
|
|
for i in reversed(range(len(rewards))):
|
|||
|
|
delta = rewards[i] + self.gamma * values[i + 1] * masks[i] - values[i]
|
|||
|
|
gae = delta + self.gamma * self.tau * masks[i] * gae
|
|||
|
|
returns.insert(0, gae + values[i])
|
|||
|
|
return returns
|
|||
|
|
|
|||
|
|
def update(self, memory):
|
|||
|
|
"""
|
|||
|
|
PPO-Clip 更新
|
|||
|
|
|
|||
|
|
Algorithm 1 (Schulman et al. 2017):
|
|||
|
|
for iteration=1, 2, ... do
|
|||
|
|
for actor=1, 2, ..., N do
|
|||
|
|
Run policy π_θold in environment for T timesteps
|
|||
|
|
Compute advantage estimates Aˆ1, ..., AˆT
|
|||
|
|
end for
|
|||
|
|
Optimize surrogate L wrt θ, with K epochs and minibatch size M ≤ NT
|
|||
|
|
θold ← θ
|
|||
|
|
end for
|
|||
|
|
"""
|
|||
|
|
# ---------- 1. 准备数据 ----------
|
|||
|
|
states = torch.FloatTensor(np.array([m[0] for m in memory]))
|
|||
|
|
actions = torch.FloatTensor(np.array([m[1] for m in memory]))
|
|||
|
|
rewards = [m[2] for m in memory]
|
|||
|
|
next_states = torch.FloatTensor(np.array([m[3] for m in memory]))
|
|||
|
|
masks = [m[4] for m in memory]
|
|||
|
|
|
|||
|
|
# ---------- 2. GAE 优势估计 ----------
|
|||
|
|
with torch.no_grad():
|
|||
|
|
values = self.critic(states).squeeze().numpy().tolist()
|
|||
|
|
next_value = self.critic(next_states[-1].unsqueeze(0)).squeeze().item()
|
|||
|
|
values.append(next_value)
|
|||
|
|
|
|||
|
|
returns = self._compute_advantages(rewards, values, masks)
|
|||
|
|
returns = torch.FloatTensor(returns)
|
|||
|
|
values = torch.FloatTensor(values[:-1])
|
|||
|
|
advantages = returns - values
|
|||
|
|
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
|
|||
|
|
|
|||
|
|
# ---------- 3. 训练 Critic (Value Network) ----------
|
|||
|
|
for _ in range(self.critic_epochs):
|
|||
|
|
critic_loss = F.mse_loss(self.critic(states).squeeze(), returns)
|
|||
|
|
self.critic_optimizer.zero_grad()
|
|||
|
|
critic_loss.backward()
|
|||
|
|
self.critic_optimizer.step()
|
|||
|
|
|
|||
|
|
# ---------- 4. 训练 Actor (PPO-Clip 损失) ----------
|
|||
|
|
# 在 no_grad 下记录旧策略的对数概率（对应 Algorithm 1 中的 π_θold）
|
|||
|
|
with torch.no_grad():
|
|||
|
|
old_dist = self.actor.evaluate(states)
|
|||
|
|
old_log_probs = old_dist.log_prob(actions).sum(dim=1)
|
|||
|
|
|
|||
|
|
dataset_size = states.size(0)
|
|||
|
|
indices = np.arange(dataset_size)
|
|||
|
|
|
|||
|
|
for _ in range(self.k_epochs):
|
|||
|
|
# 每轮随机打乱，分成多个 minibatch
|
|||
|
|
np.random.shuffle(indices)
|
|||
|
|
|
|||
|
|
for start in range(0, dataset_size, self.minibatch_size):
|
|||
|
|
end = start + self.minibatch_size
|
|||
|
|
mb_idx = indices[start:end]
|
|||
|
|
|
|||
|
|
mb_states = states[mb_idx]
|
|||
|
|
mb_actions = actions[mb_idx]
|
|||
|
|
mb_advantages = advantages[mb_idx]
|
|||
|
|
mb_old_log_probs = old_log_probs[mb_idx]
|
|||
|
|
|
|||
|
|
# 当前策略的对数概率
|
|||
|
|
dist = self.actor.evaluate(mb_states)
|
|||
|
|
log_probs = dist.log_prob(mb_actions).sum(dim=1)
|
|||
|
|
|
|||
|
|
# 概率比 r_t(θ) = π_θ(a|s) / π_θold(a|s)
|
|||
|
|
ratio = torch.exp(log_probs - mb_old_log_probs)
|
|||
|
|
|
|||
|
|
# ---------- PPO-Clip 核心 ----------
|
|||
|
|
# L^CLIP(θ) = E[min(r_t(θ) * A_t, clip(r_t(θ), 1-ε, 1+ε) * A_t)]
|
|||
|
|
surr1 = ratio * mb_advantages
|
|||
|
|
surr2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * mb_advantages
|
|||
|
|
policy_loss = -torch.min(surr1, surr2).mean()
|
|||
|
|
# ---------------------------------
|
|||
|
|
|
|||
|
|
# 可选：熵正则化（鼓励探索）
|
|||
|
|
entropy_loss = -self.entropy_coef * dist.entropy().mean() if self.entropy_coef > 0 else 0
|
|||
|
|
|
|||
|
|
total_loss = policy_loss + entropy_loss
|
|||
|
|
|
|||
|
|
self.actor_optimizer.zero_grad()
|
|||
|
|
total_loss.backward()
|
|||
|
|
self.actor_optimizer.step()
|