From 96cd594be9055dc2fe24c8021d3ae6970b27cabd Mon Sep 17 00:00:00 2001 From: lhr Date: Thu, 2 Apr 2026 02:24:35 +0000 Subject: [PATCH] first commit --- .gitignore | 46 +++++++++ agent/__init__.py | 2 + agent/ppo.py | 150 +++++++++++++++++++++++++++++ agent/trpo.py | 239 ++++++++++++++++++++++++++++++++++++++++++++++ environment.yml | 11 +++ main.py | 71 ++++++++++++++ networks.py | 56 +++++++++++ requirements.txt | 4 + 8 files changed, 579 insertions(+) create mode 100644 .gitignore create mode 100644 agent/__init__.py create mode 100644 agent/ppo.py create mode 100644 agent/trpo.py create mode 100644 environment.yml create mode 100644 main.py create mode 100644 networks.py create mode 100644 requirements.txt diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..bdd270b --- /dev/null +++ b/.gitignore @@ -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 diff --git a/agent/__init__.py b/agent/__init__.py new file mode 100644 index 0000000..4e0b03d --- /dev/null +++ b/agent/__init__.py @@ -0,0 +1,2 @@ +from agent.ppo import PPOAgent +from agent.trpo import TRPOAgent diff --git a/agent/ppo.py b/agent/ppo.py new file mode 100644 index 0000000..86e5930 --- /dev/null +++ b/agent/ppo.py @@ -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() diff --git a/agent/trpo.py b/agent/trpo.py new file mode 100644 index 0000000..607a63e --- /dev/null +++ b/agent/trpo.py @@ -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()) diff --git a/environment.yml b/environment.yml new file mode 100644 index 0000000..316162f --- /dev/null +++ b/environment.yml @@ -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 diff --git a/main.py b/main.py new file mode 100644 index 0000000..177c93f --- /dev/null +++ b/main.py @@ -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() diff --git a/networks.py b/networks.py new file mode 100644 index 0000000..b46b105 --- /dev/null +++ b/networks.py @@ -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 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..a3ed8f6 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +numpy>=1.24 +matplotlib>=3.7 +torch>=2.1 +gymnasium[classic-control]>=0.29