Hongru/RL_TRPO
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

修复并行采样并完善训练文档

Browse Source
This commit is contained in:
Hongru
2026-04-02 09:48:59 +00:00
parent 771eba8607
commit 428e6f7f81
6 changed files with 943 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agent/ppo.py
+72 -27
View File
@@ -2,7 +2,7 @@ import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Normal
from networks import PolicyNetwork, ValueNetwork
@@ -32,6 +32,7 @@ class PPOAgent:
k_epochs=10,
minibatch_size=64,
max_grad_norm=0.5,
device="cpu",
):
self.gamma = gamma
self.tau = tau
@@ -39,37 +40,58 @@ class PPOAgent:
self.k_epochs = k_epochs # 每次更新对同一批数据的训练轮数
self.minibatch_size = minibatch_size
self.max_grad_norm = max_grad_norm # 梯度裁剪阈值
self.device = torch.device(device)
self.actor = PolicyNetwork(state_dim, action_dim, action_bound, hidden_dim)
self.critic = ValueNetwork(state_dim, hidden_dim)
self.actor = PolicyNetwork(state_dim, action_dim, action_bound, hidden_dim).to(self.device)
self.critic = ValueNetwork(state_dim, hidden_dim).to(self.device)
self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)
self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=critic_lr)
def _to_tensor(self, array_like):
return torch.as_tensor(array_like, dtype=torch.float32, device=self.device)
def get_action(self, state):
state_tensor = torch.FloatTensor(state).unsqueeze(0)
state_tensor = self._to_tensor(state)
single_state = state_tensor.ndim == 1
if single_state:
state_tensor = state_tensor.unsqueeze(0)
with torch.no_grad():
dist = self.actor.evaluate(state_tensor)
action = dist.sample()
return action.squeeze(0).numpy()
action_np = action.detach().cpu().numpy()
return action_np[0] if single_state else action_np
def get_value(self, state):
with torch.no_grad():
return self.critic(torch.FloatTensor(state).unsqueeze(0)).squeeze().item()
state_tensor = self._to_tensor(state)
single_state = state_tensor.ndim == 1
if single_state:
state_tensor = state_tensor.unsqueeze(0)
def _compute_advantages(self, rewards, values, masks):
with torch.no_grad():
values = self.critic(state_tensor).squeeze(-1)
values_np = values.detach().cpu().numpy()
return float(values_np[0]) if single_state else values_np
def _compute_advantages(self, rewards, values, next_values, masks):
"""
GAE (Generalized Advantage Estimation)
Eq.(11): Â_t = δ_t + (γλ)δ_{t+1} + ... + (γλ)^{T-t+1} δ_{T-1}
Eq.(12): δ_t = r_t + γV(s_{t+1}) - V(s_t)
"""
returns = []
gae = 0
for i in reversed(range(len(rewards))):
delta = rewards[i] + self.gamma * values[i + 1] * masks[i] - values[i]
advantages = torch.zeros_like(rewards, device=self.device)
gae = torch.zeros(rewards.size(1), dtype=torch.float32, device=self.device)
for i in reversed(range(rewards.size(0))):
delta = rewards[i] + self.gamma * next_values[i] * masks[i] - values[i]
gae = delta + self.gamma * self.tau * masks[i] * gae
returns.insert(0, gae + values[i])
return returns
advantages[i] = gae
returns = advantages + values
return returns, advantages
def update(self, memory):
"""
@@ -80,25 +102,32 @@ class PPOAgent:
L = L^CLIP - c1 * L^VF (Eq.9，c2=0 不加熵)
2. 梯度裁剪 (max_grad_norm)
"""
if not memory:
return {}
# ---------- 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]
states = self._to_tensor(np.asarray(memory["states"]))
actions = self._to_tensor(np.asarray(memory["actions"]))
rewards = self._to_tensor(np.asarray(memory["rewards"]))
next_states = self._to_tensor(np.asarray(memory["next_states"]))
masks = self._to_tensor(np.asarray(memory["masks"]))
# ---------- 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)
rollout_steps, num_envs = rewards.shape
flat_states = states.reshape(rollout_steps * num_envs, -1)
flat_next_states = next_states.reshape(rollout_steps * num_envs, -1)
values = self.critic(flat_states).squeeze(-1).reshape(rollout_steps, num_envs)
next_values = self.critic(flat_next_states).squeeze(-1).reshape(rollout_steps, num_envs)
returns = self._compute_advantages(rewards, values, masks)
returns = torch.FloatTensor(returns)
values = torch.FloatTensor(values[:-1])
advantages = returns - values
returns, advantages = self._compute_advantages(rewards, values, next_values, masks)
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
states = flat_states
actions = actions.reshape(rollout_steps * num_envs, -1)
returns = returns.reshape(rollout_steps * num_envs)
advantages = advantages.reshape(rollout_steps * num_envs)
# ---------- 3. 记录旧策略 π_θold ----------
with torch.no_grad():
old_dist = self.actor.evaluate(states)
@@ -107,13 +136,17 @@ class PPOAgent:
# ---------- 4. K epochs × minibatch 联合训练 ----------
dataset_size = states.size(0)
indices = np.arange(dataset_size)
policy_loss_value = 0.0
critic_loss_value = 0.0
entropy_value = 0.0
minibatch_updates = 0
for _ in range(self.k_epochs):
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_idx = torch.as_tensor(indices[start:end], device=self.device, dtype=torch.long)
mb_states = states[mb_idx]
mb_actions = actions[mb_idx]
@@ -144,3 +177,15 @@ class PPOAgent:
self.critic_optimizer.zero_grad()
critic_loss.backward()
self.critic_optimizer.step()
policy_loss_value += float(policy_loss.detach().item())
critic_loss_value += float(critic_loss.detach().item())
entropy_value += float(dist.entropy().sum(dim=1).mean().detach().item())
minibatch_updates += 1
divisor = max(1, minibatch_updates)
return {
"policy_loss": policy_loss_value / divisor,
"critic_loss": critic_loss_value / divisor,
"entropy": entropy_value / divisor,
}