PPO：近端策略优化

# TRPO 的共轭梯度求解（简化版）
def conjugate_gradient(fvp_fn, b, nsteps=10, residual_tol=1e-10):
    """Solve Ax = b using conjugate gradient"""
    x = torch.zeros_like(b)
    r = b.clone()
    p = b.clone()
    rdotr = torch.dot(r, r)
    for _ in range(nsteps):
        Av = fvp_fn(p)
        alpha = rdotr / torch.dot(p, Av)
        x += alpha * p
        r -= alpha * Av
        new_rdotr = torch.dot(r, r)
        if new_rdotr < residual_tol:
            break
        p = r + (new_rdotr / rdotr) * p
        rdotr = new_rdotr
    return x

# PPO 的简化实现对比
def ppo_update(policy, old_policy, states, actions, advantages,
               clip_eps=0.2, lr=3e-4):
    """PPO 用一阶 SGD 替代 TRPO 的共轭梯度"""
    optimizer = torch.optim.Adam(policy.parameters(), lr=lr)
    for epoch in range(4):  # 多次 epoch 提升样本效率
        dist = policy(states)
        old_dist = old_policy(states)

        log_prob = dist.log_prob(actions)
        old_log_prob = old_dist.log_prob(actions).detach()
        ratio = torch.exp(log_prob - old_log_prob)

        surr1 = ratio * advantages
        surr2 = torch.clamp(ratio, 1 - clip_eps, 1 + clip_eps) * advantages
        loss = -torch.min(surr1, surr2).mean()

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

import torch

def compute_ppo_loss(log_probs, old_log_probs, advantages, clip_eps=0.2):
    """PPO-Clip 核心：裁剪目标函数"""
    ratio = torch.exp(log_probs - old_log_probs)

    # 未裁剪的 surrogate
    surr1 = ratio * advantages

    # 裁剪后的 surrogate
    clipped_ratio = torch.clamp(ratio, 1.0 - clip_eps, 1.0 + clip_eps)
    surr2 = clipped_ratio * advantages

    # 取最小值，确保策略更新不超过信任区域
    loss = -torch.min(surr1, surr2)

    # 近似 KL 散度（用于监控，不参与梯度）
    approx_kl = (old_log_probs - log_probs).mean().item()

    return loss.mean(), approx_kl

def ppo_epoch(policy, old_policy, states, actions, advantages,
             values, old_values, returns, clip_eps=0.2,
             value_coef=0.5, ent_coef=0.01):
    """完整的 PPO epoch 更新"""
    dist = policy.get_dist(states)
    log_probs = dist.log_prob(actions)
    entropy = dist.entropy().mean()

    # Value loss with clipping
    pred_values = policy.get_value(states)
    value_clipped = old_values + (pred_values - old_values).clamp(
        -clip_eps, clip_eps
    )
    value_loss1 = (pred_values - returns).pow(2)
    value_loss2 = (value_clipped - returns).pow(2)
    value_loss = 0.5 * torch.max(value_loss1, value_loss2).mean()

    # Total loss
    policy_loss, _ = compute_ppo_loss(
        log_probs, old_log_probs, advantages, clip_eps
    )
    total_loss = policy_loss + value_coef * value_loss - ent_coef * entropy

    return total_loss

class PPOPenalty:
    """PPO-Penalty 实现：自适应 KL 惩罚"""
    def __init__(self, kl_target=0.01, beta_init=1.0,
                 beta_adjust_rate=2.0, clip_high=10.0):
        self.kl_target = kl_target
        self.beta = beta_init
        self.adjust_rate = beta_adjust_rate
        self.clip_high = clip_high

    def compute_loss(self, log_probs, old_log_probs, advantages):
        ratio = torch.exp(log_probs - old_log_probs)
        kl_approx = (old_log_probs - log_probs).mean()

        # 自适应调整 beta
        if kl_approx > self.kl_target * 1.5:
            self.beta = min(self.beta * self.adjust_rate, self.clip_high)
        elif kl_approx < self.kl_target / 1.5:
            self.beta = max(self.beta / self.adjust_rate, 1e-4)

        # 带 KL 惩罚的目标
        surr = ratio * advantages
        penalty = self.beta * kl_approx
        loss = -(surr.mean() - penalty)
        return loss, kl_approx.item()

def compare_ppo_variants(env_name="Hopper-v4", seed=42):
    """对比 PPO-Clip 和 PPO-Penalty 的训练效果"""
    results = {"clip": [], "penalty": []}

    for variant in ["clip", "penalty"]:
        env = gym.make(env_name)
        agent = PPOAgent(env, variant=variant, seed=seed)

        for episode in range(500):
            state, _ = env.reset(seed=seed)
            episode_reward = 0
            done = False
            while not done:
                action = agent.select_action(state)
                next_state, reward, done, _, _ = env.step(action)
                agent.store(state, action, reward, done)
                state = next_state
                episode_reward += reward

            agent.update()
            results[variant].append(episode_reward)
            print(f"{variant}: ep={episode}, reward={episode_reward:.1f}")

    return results

def compute_gae(rewards, values, dones, gamma=0.99, lam=0.95):
    """GAE 广义优势估计"""
    advantages = []
    gae = 0.0
    n = len(rewards)

    for t in reversed(range(n)):
        if t == n - 1:
            next_value = 0.0
        else:
            next_value = values[t + 1]

        # TD error: delta_t = r_t + gamma * V(s_{t+1}) - V(s_t)
        delta = rewards[t] + gamma * next_value * (1 - dones[t]) - values[t]

        # GAE: A_t = delta_t + gamma * lambda * A_{t+1}
        gae = delta + gamma * lam * (1 - dones[t]) * gae
        advantages.insert(0, gae)

    return torch.tensor(advantages, dtype=torch.float32)

def compute_returns_gae(rewards, values, dones,
                     gamma=0.99, lam=0.95):
    """完整的 GAE + Returns 计算"""
    advantages = compute_gae(rewards, values, dones, gamma, lam)

    # Returns = advantages + values（用于 value loss）
    returns = advantages + torch.tensor(values, dtype=torch.float32)

    # 标准化优势（重要技巧：提升训练稳定性）
    advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

    return advantages, returns

def visualize_lambda_effect():
    """不同 lambda 值的偏差-方差权衡"""
    lambdas = [0.0, 0.5, 0.95, 1.0]
    for lam in lambdas:
        advantages = compute_gae(rewards, values, dones, lam=lam)
        print(f"lambda={lam:.2f}: mean={advantages.mean():.3f}, "
              f"std={advantages.std():.3f}, "
              f"range=[{advantages.min():.3f}, {advantages.max():.3f}]")

import gymnasium as gym
from gymnasium.vector import AsyncVectorEnv

def create_parallel_envs(env_id, num_envs=8, seed=42):
    """创建并行向量环境"""
    def make_env(env_id, idx):
        def thunk():
            env = gym.make(env_id)
            env = gym.wrappers.RecordEpisodeStatistics(env)
            if idx == 0:
                env = gym.wrappers.RenderMode(env, "human")
            return env
        return thunk

    envs = AsyncVectorEnv(
        [make_env(env_id, i) for i in range(num_envs)],
        seed=seed,
    )
    return envs

def parallel_rollout(agent, envs, num_steps=128):
    """在多个环境中并行收集经验"""
    obs, _ = envs.reset()
    states, actions, rewards, dones, log_probs, values = [], [], [], [], [], []

    for _ in range(num_steps):
        action, log_prob, value = agent.get_action_and_value(obs)
        next_obs, reward, terminated, truncated, info = envs.step(action)
        done = terminated | truncated

        states.append(obs)
        actions.append(action)
        rewards.append(reward)
        dones.append(done)
        log_probs.append(log_prob)
        values.append(value)
        obs = next_obs

    return states, actions, rewards, dones, log_probs, values, obs

def analyze_parallel_scaling():
    """分析并行环境数量对训练的影响"""
    configs = [
        {"num_envs": 4, "rollout": 512, "total": 2048},
        {"num_envs": 8, "rollout": 256, "total": 2048},
        {"num_envs": 16, "rollout": 128, "total": 2048},
        {"num_envs": 32, "rollout": 64, "total": 2048},
        {"num_envs": 64, "rollout": 128, "total": 8192},
    ]

    print(f"{'Env数':>6} | {'Rollout':>8} | {'总样本':>8} | {'评估'}")
    print("-" * 42)
    for cfg in configs:
        n = cfg["num_envs"]
        r = cfg["rollout"]
        t = cfg["total"]
        comment = "推荐" if t == 2048 and 8 <= n <= 32 else "极端"
        print(f"{n:>6} | {r:>8} | {t:>8} | {comment}")

# 评估经验法则：
# - num_envs=8~32, rollout=128~256 是最佳平衡点
# - 总样本 < 1024 时梯度方差过大
# - 总样本 > 16384 时单次更新可能超出信任区域

PPO_CONFIGS = {
    "default": {
        "learning_rate": 3e-4,
        "n_steps": 2048,
        "batch_size": 64,
        "n_epochs": 10,
        "gamma": 0.99,
        "gae_lambda": 0.95,
        "clip_range": 0.2,
        "ent_coef": 0.0,
        "vf_coef": 0.5,
        "max_grad_norm": 0.5,
        "use_sde": False,
    },
    "continuous_control": {
        "learning_rate": 3e-4,
        "n_steps": 2048,
        "batch_size": 64,
        "n_epochs": 10,
        "gamma": 0.99,
        "gae_lambda": 0.95,
        "clip_range": 0.2,
        "ent_coef": 0.0,
        "vf_coef": 0.5,
        "max_grad_norm": 0.5,
        "use_sde": True,  # State-Dependent Exploration
    },
    "atari_discrete": {
        "learning_rate": 2.5e-4,
        "n_steps": 128,
        "batch_size": 256,
        "n_epochs": 4,
        "gamma": 0.99,
        "gae_lambda": 0.95,
        "clip_range": 0.1,
        "ent_coef": 0.01,
        "vf_coef": 0.5,
        "max_grad_norm": 0.5,
    },
}

class PPOWithEarlyStopping:
    """带 Early Stopping 的 PPO 实现"""
    def __init__(self, clip_eps=0.2, target_kl=0.03):
        self.clip_eps = clip_eps
        self.target_kl = target_kl

    def update(self, states, actions, advantages, returns,
               old_log_probs, old_values, policy, value_net,
               optimizer, n_epochs=10):
        """PPO 更新，当 KL 超过阈值时提前停止"""
        batch_size = states.shape[0]
        indices = np.arange(batch_size)

        for epoch in range(n_epochs):
            np.random.shuffle(indices)
            for start in range(0, batch_size, self.batch_size):
                end = start + self.batch_size
                batch_idx = indices[start:end]

                # 计算 KL 散度
                with torch.no_grad():
                    new_dist = policy(states[batch_idx])
                    old_dist = old_policy(states[batch_idx])
                    kl = torch.distributions.kl_divergence(
                        old_dist, new_dist
                    ).mean().item()

                if kl > self.target_kl:
                    print(f"Early stop at epoch {epoch}, "
                          f"KL={kl:.4f} > {self.target_kl}")
                    return  # 提前终止

                # 正常 PPO 更新
                loss = compute_ppo_loss(...)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

import torch
import torch.nn as nn
import torch.optim as optim
import gymnasium as gym
import numpy as np
from torch.distributions import Categorical

class PPONetwork(nn.Module):
    """Actor-Critic 网络"""
    def __init__(self, state_dim, action_dim, hidden_dim=64):
        super().__init__()
        self.actor = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, action_dim),
        )
        self.critic = 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):
        logits = self.actor(state)
        value = self.critic(state).squeeze(-1)
        return logits, value

def train_ppo_lunarlander(total_timesteps=500000, seed=42):
    """完整的 PPO LunarLander 训练"""
    env = gym.make("LunarLander-v2", render_mode=None)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 初始化
    net = PPONetwork(state_dim=8, action_dim=4).to(device)
    optimizer = optim.Adam(net.parameters(), lr=2.5e-4, eps=1e-5)

    # 训练循环
    num_iterations = total_timesteps // 2048
    for iteration in range(num_iterations):
        # 线性衰减学习率
        lr = 2.5e-4 * (1 - iteration / num_iterations)
        for param_group in optimizer.param_groups:
            param_group["lr"] = lr

        # 收集经验 (rollout)
        states, actions, rewards, dones, log_probs, values = [], [], [], [], [], []
        state, _ = env.reset(seed=seed + iteration)
        for _ in range(2048):
            state_t = torch.FloatTensor(state).unsqueeze(0).to(device)
            with torch.no_grad():
                logits, value = net(state_t)
                dist = Categorical(logits=logits)
                action = dist.sample()
                log_prob = dist.log_prob(action)

            states.append(state)
            actions.append(action.item())
            log_probs.append(log_prob.item())
            values.append(value.item())

            next_state, reward, done, truncated, _ = env.step(action.item())
            rewards.append(reward)
            dones.append(done or truncated)
            state = next_state if not (done or truncated) else env.reset()[0]

        # GAE + PPO 更新
        advantages = compute_gae(rewards, values, dones)
        returns = torch.tensor(advantages) + torch.tensor(values)
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

        for epoch in range(10):
            # 转换为 tensor
            states_t = torch.FloatTensor(states).to(device)
            actions_t = torch.LongTensor(actions).to(device)
            advantages_t = advantages.to(device)
            old_log_probs_t = torch.FloatTensor(log_probs).to(device)

            logits, _ = net(states_t)
            dist = Categorical(logits=logits)
            new_log_probs = dist.log_prob(actions_t)

            loss, kl = compute_ppo_loss(
                new_log_probs, old_log_probs_t, advantages_t
            )

            optimizer.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(net.parameters(), 0.5)
            optimizer.step()

            if kl > 0.03:
                break

        if iteration % 10 == 0:
            print(f"Iter {iteration}: loss={loss:.4f}, kl={kl:.4f}")

    return net