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

first commit

Browse Source
This commit is contained in:
Hongru
2026-04-02 02:24:35 +00:00
commit 96cd594be9
8 changed files with 579 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+46
View File
@@ -0,0 +1,46 @@
# Python cache and bytecode
__pycache__/
*.py[cod]
*$py.class
# Compiled extensions
*.so
# Build and packaging
.Python
build/
dist/
*.egg-info/
.eggs/
# Virtual environments
.venv/
venv/
env/
ENV/
# Tool caches
.pytest_cache/
.mypy_cache/
.ruff_cache/
.ipynb_checkpoints/
# IDE/editor
.vscode/
.idea/
# OS files
.DS_Store
Thumbs.db
# Logs and runtime files
*.log
# Training artifacts
*.png
*.pt
*.pth
*.ckpt
# Environment files
.env
agent/__init__.py
+2
View File
@@ -0,0 +1,2 @@
from agent.ppo import PPOAgent
from agent.trpo import TRPOAgent
agent/ppo.py
+150
View File
@@ -0,0 +1,150 @@
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()
agent/trpo.py
+239
View File
@@ -0,0 +1,239 @@
import numpy as np
import torch
import torch.nn.functional as F
from torch.distributions import Normal
from torch.nn.utils import parameters_to_vector, vector_to_parameters
from networks import PolicyNetwork, ValueNetwork
class TRPOAgent:
def __init__(self, state_dim, action_dim, action_bound, hidden_dim=128, kl_margin=0.01, gamma=0.99, tau=0.97, cg_iters=10):
self.gamma = gamma
self.tau = tau
self.kl_margin = kl_margin
self.cg_iters = cg_iters
self.actor = PolicyNetwork(state_dim, action_dim, action_bound, hidden_dim)
self.critic = ValueNetwork(state_dim, hidden_dim)
self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=1e-3)
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) 计算优势函数
"""
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
# --- 关键修复 1：将固定的 old_dist 作为参数传入 ---
def _hessian_vector_product(self, states, old_dist, vector, damping=0.1):
"""
计算海森矩阵与向量的乘积 (Hvp)
这里的 old_dist 必须是不带梯度的历史常数分布
"""
dist = self.actor.evaluate(states)
# 计算当前分布与固定的旧分布之间的 KL 散度
kl = torch.distributions.kl_divergence(old_dist, dist).mean()
# 一阶导数
grads = torch.autograd.grad(kl, self.actor.parameters(), create_graph=True)
flat_grad_kl = torch.cat([grad.view(-1) for grad in grads])
# 与给定向量点乘
kl_v = (flat_grad_kl * vector).sum()
# 二阶导数
grads = torch.autograd.grad(kl_v, self.actor.parameters())
flat_grad_grad_kl = torch.cat([grad.contiguous().view(-1) for grad in grads])
# 加入阻尼系数，保证矩阵正定，避免数值不稳定发散
return flat_grad_grad_kl + vector * damping
# --- 关键修复 2：共轭梯度法同样接收 old_dist ---
def _conjugate_gradient(self, states, old_dist, b, nsteps, residual_tol=1e-10):
"""
共轭梯度法，近似求解 Hx = b
"""
x = torch.zeros_like(b)
r = b.clone()
p = b.clone()
rdotr = torch.dot(r, r)
for _ in range(nsteps):
Hp = self._hessian_vector_product(states, old_dist, p)
alpha = rdotr / torch.dot(p, Hp)
x += alpha * p
r -= alpha * Hp
new_rdotr = torch.dot(r, r)
if new_rdotr < residual_tol:
break
p = r + new_rdotr / rdotr * p
rdotr = new_rdotr
return x
def update(self, memory):
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]
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)
# 【修复】：加大 Critic 的训练力度，从 10 提升到 40 Epochs
# 确保裁判的眼光足够准确
for _ in range(40):
critic_loss = F.mse_loss(self.critic(states).squeeze(), returns)
self.critic_optimizer.zero_grad()
critic_loss.backward()
self.critic_optimizer.step()
with torch.no_grad():
old_mean, old_std = self.actor(states)
old_dist = Normal(old_mean, old_std)
old_log_probs = old_dist.log_prob(actions).sum(dim=1)
def compute_surrogate_loss():
dist = self.actor.evaluate(states)
log_probs = dist.log_prob(actions).sum(dim=1)
ratio = torch.exp(log_probs - old_log_probs)
surrogate_loss = (ratio * advantages).mean()
return surrogate_loss, dist
surrogate_loss, dist = compute_surrogate_loss()
loss_grad = torch.autograd.grad(surrogate_loss, self.actor.parameters())
loss_grad_flat = torch.cat([grad.view(-1) for grad in loss_grad])
step_dir = self._conjugate_gradient(states, old_dist, loss_grad_flat, self.cg_iters)
shs = 0.5 * torch.dot(step_dir, self._hessian_vector_product(states, old_dist, step_dir))
if shs < 1e-8:
return
lm = torch.sqrt(shs / self.kl_margin)
fullstep = step_dir / lm
old_params = parameters_to_vector(self.actor.parameters())
# 线性搜索
success = False
step_size = 1.0
for _ in range(10):
new_params = old_params + step_size * fullstep
vector_to_parameters(new_params, self.actor.parameters())
with torch.no_grad():
new_surrogate_loss, new_dist = compute_surrogate_loss()
kl = torch.distributions.kl_divergence(old_dist, new_dist).mean()
# 【修复】：增加极小的浮点数宽容度，防止在极小提升时被误判失败而拒绝更新
if new_surrogate_loss >= surrogate_loss - 1e-8 and kl <= self.kl_margin * 1.5:
success = True
break
step_size *= 0.5
if not success:
vector_to_parameters(old_params, self.actor.parameters())
"""
利用收集到的轨迹数据更新 Actor 和 Critic
"""
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]
# 1. 拟合价值网络 (Critic)
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)
for _ in range(10):
critic_loss = F.mse_loss(self.critic(states).squeeze(), returns)
self.critic_optimizer.zero_grad()
critic_loss.backward()
self.critic_optimizer.step()
# --- 关键修复 3：在截断梯度的环境下生成严格的旧分布 ---
with torch.no_grad():
old_mean, old_std = self.actor(states)
old_dist = Normal(old_mean, old_std)
old_log_probs = old_dist.log_prob(actions).sum(dim=1)
def compute_surrogate_loss():
# 计算替代目标函数 (Surrogate Objective)
dist = self.actor.evaluate(states)
log_probs = dist.log_prob(actions).sum(dim=1)
ratio = torch.exp(log_probs - old_log_probs)
surrogate_loss = (ratio * advantages).mean()
return surrogate_loss, dist
surrogate_loss, dist = compute_surrogate_loss()
loss_grad = torch.autograd.grad(surrogate_loss, self.actor.parameters())
loss_grad_flat = torch.cat([grad.view(-1) for grad in loss_grad])
# 传入 old_dist，确保海森矩阵计算包含准确的曲率信息
step_dir = self._conjugate_gradient(states, old_dist, loss_grad_flat, self.cg_iters)
shs = 0.5 * torch.dot(step_dir, self._hessian_vector_product(states, old_dist, step_dir))
# 增加数值保护：防止 shs 出现负数或极小值导致报错
if shs < 1e-8:
return
lm = torch.sqrt(shs / self.kl_margin)
fullstep = step_dir / lm
old_params = parameters_to_vector(self.actor.parameters())
# 线性搜索 (Line Search)
success = False
step_size = 1.0
for _ in range(10):
new_params = old_params + step_size * fullstep
vector_to_parameters(new_params, self.actor.parameters())
with torch.no_grad():
new_surrogate_loss, new_dist = compute_surrogate_loss()
kl = torch.distributions.kl_divergence(old_dist, new_dist).mean()
if new_surrogate_loss > surrogate_loss and kl <= self.kl_margin:
success = True
break
step_size *= 0.5
if not success:
vector_to_parameters(old_params, self.actor.parameters())
environment.yml
+11
View File
@@ -0,0 +1,11 @@
name: rl_trpo
channels:
- conda-forge
dependencies:
- python=3.10
- pip
- numpy>=1.24
- matplotlib>=3.7
- pip:
- torch>=2.1
- gymnasium[classic-control]>=0.29
main.py
+71
View File
@@ -0,0 +1,71 @@
import gymnasium as gym
import matplotlib.pyplot as plt
import numpy as np
from agent.ppo import PPOAgent
def main():
env = gym.make('Pendulum-v1')
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.shape[0]
action_bound = float(env.action_space.high[0])
agent = PPOAgent(state_dim=state_dim, action_dim=action_dim, action_bound=action_bound)
num_episodes = 200
batch_size = 2000
episode_rewards = []
state, _ = env.reset()
memory = []
current_ep_reward = 0
episodes_completed = 0
print("开始训练 PPO 智能体...")
step_count = 0
while episodes_completed < num_episodes:
action = agent.get_action(state)
# 交互
next_state, reward, terminated, truncated, _ = env.step(action)
# 【极其关键的修复】：只有真正死亡 (terminated) 才清零未来价值
# 绝对不能把时间截断 (truncated) 算作 mask=0
mask = 0.0 if terminated else 1.0
done = terminated or truncated
memory.append([state, action, reward, next_state, mask])
state = next_state
current_ep_reward += reward
step_count += 1
if done:
episode_rewards.append(current_ep_reward)
episodes_completed += 1
state, _ = env.reset()
current_ep_reward = 0
if episodes_completed % 10 == 0:
avg_reward = np.mean(episode_rewards[-10:])
print(f"Episode: {episodes_completed}, 平均奖励 (最近10轮): {avg_reward:.2f}")
if step_count >= batch_size:
agent.update(memory)
memory.clear()
step_count = 0
env.close()
plt.figure(figsize=(10, 5))
plt.plot(episode_rewards)
plt.title('PPO Learning Curve on Pendulum-v1')
plt.xlabel('Episode')
plt.ylabel('Total Reward')
plt.grid(True)
plt.savefig('ppo_learning_curve_final.png')
plt.show()
if __name__ == '__main__':
main()
networks.py
+56
View File
@@ -0,0 +1,56 @@
import torch
import torch.nn as nn
from torch.distributions import Normal
class ValueNetwork(nn.Module):
"""
状态值函数网络 (Critic)，用于估计状态的内在价值 V(s)
增加了隐藏层容量，以确保在多轮迭代中能够更准确地拟合优势函数
"""
def __init__(self, state_dim, hidden_dim=128):
super(ValueNetwork, self).__init__()
self.net = nn.Sequential(
nn.Linear(state_dim, hidden_dim),
nn.Tanh(),
nn.Linear(hidden_dim, hidden_dim),
nn.Tanh(),
nn.Linear(hidden_dim, 1)
)
def forward(self, state):
# 返回对当前状态的价值评估
return self.net(state)
class PolicyNetwork(nn.Module):
"""
策略网络 (Actor)，输出高斯分布的均值和标准差，适用于连续动作空间
"""
def __init__(self, state_dim, action_dim, action_bound, hidden_dim=128):
super(PolicyNetwork, self).__init__()
self.action_bound = action_bound # 环境允许的最大物理动作幅度
self.net = nn.Sequential(
nn.Linear(state_dim, hidden_dim),
nn.Tanh(),
nn.Linear(hidden_dim, hidden_dim),
nn.Tanh(),
nn.Linear(hidden_dim, action_dim),
nn.Tanh() # 关键修复：强制均值输出在 [-1, 1] 之间，防止动作空间爆炸
)
# 将初始对数标准差设为 -0.5 (对应的标准差约为 0.6)
# 较小的初始方差有助于防止初期探索步子迈得太大导致系统崩溃
self.action_log_std = nn.Parameter(torch.full((1, action_dim), -0.5))
def forward(self, state):
# 计算动作均值并将其缩放到实际的物理边界内
action_mean = self.net(state) * self.action_bound
# 限制标准差的范围以提高数值稳定性
action_std = torch.exp(self.action_log_std.expand_as(action_mean))
return action_mean, action_std
def evaluate(self, state):
# 评估当前状态，返回构建好的高斯动作分布
mean, std = self.forward(state)
dist = Normal(mean, std)
return dist
requirements.txt
+4
View File
@@ -0,0 +1,4 @@
numpy>=1.24
matplotlib>=3.7
torch>=2.1
gymnasium[classic-control]>=0.29