强化学习基础：从 Q-Learning 到 PPO

# 强化学习交互循环的简化模拟
import numpy as np

class SimpleEnvironment:
    """一个极简的强化学习环境：寻找宝藏
    
    状态: 0-4 共 5 个位置
    动作: 0=向左, 1=向右
    奖励: 到达位置 4 得 +10，其余位置得 -1
    """
    def __init__(self):
        self.state = 0
        self.done = False

    def reset(self):
        self.state = 0
        self.done = False
        return self.state

    def step(self, action):
        if action == 0:  # 向左
            self.state = max(0, self.state - 1)
        else:  # 向右
            self.state = min(4, self.state + 1)

        reward = 10 if self.state == 4 else -1
        self.done = (self.state == 4)
        return self.state, reward, self.done

# 演示交互
env = SimpleEnvironment()
state = env.reset()
print(f"初始状态: {state}")

# 随机策略
total_reward = 0
for step in range(20):
    action = np.random.choice([0, 1])
    next_state, reward, done = env.step(action)
    total_reward += reward
    print(f"  步 {step+1}: 动作={'←' if action==0 else '→'}, "
          f"状态={state}→{next_state}, 奖励={reward}")
    state = next_state
    if done:
        print(f"  找到宝藏！总奖励: {total_reward}")
        break

# 折扣累积回报计算：理解 gamma 的作用

def compute_returns(rewards, gamma=0.9):
    """计算每一步的折扣累积回报
    
    G_t = r_{t+1} + gamma*r_{t+2} + gamma^2*r_{t+3} + ...
    
    gamma 决定了"远见"程度：
    - gamma=0：只看眼前（贪婪）
    - gamma=0.99：非常长远
    - gamma=1：完全平等看待未来
    """
    returns = np.zeros(len(rewards))
    G = 0
    for t in reversed(range(len(rewards))):
        G = rewards[t] + gamma * G
        returns[t] = G
    return returns

# 不同 gamma 对回报的影响
rewards = [-1, -1, -1, -1, -1, -1, -1, -1, -1, 10]

print("不同折扣因子下的回报：")
for gamma in [0.0, 0.5, 0.9, 0.99]:
    returns = compute_returns(rewards, gamma)
    print(f"  gamma={gamma:.2f}: G_0 = {returns[0]:.2f}, G_5 = {returns[5]:.2f}")

# 用字典表示一个完整的 MDP：悬崖行走（Cliff Walking）
# 4×12 的网格，从左下角走到右下角，避开悬崖

class CliffWalkingMDP:
    """悬崖行走 MDP
    
    4行12列的网格世界：
    S = 起点 (3, 0)
    G = 终点 (3, 11)
    第3行中间10个格子是悬崖
    
    到达悬崖：-100 奖励，回到起点
    到达终点：0 奖励，episode 结束
    每步移动：-1 奖励
    """
    def __init__(self, rows=4, cols=12):
        self.rows = rows
        self.cols = cols
        self.actions = [0, 1, 2, 3]  # 上, 右, 下, 左
        self.action_names = ['上', '右', '下', '左']
        self.deltas = [(-1, 0), (0, 1), (1, 0), (0, -1)]
        self.cliff_rows = {rows - 1}
        self.cliff_cols = set(range(1, cols - 1))
        self.n_states = rows * cols

    def state_to_pos(self, state):
        return divmod(state, self.cols)

    def pos_to_state(self, row, col):
        return row * self.cols + col

    def step(self, state, action):
        row, col = self.state_to_pos(state)
        dr, dc = self.deltas[action]
        new_row = max(0, min(self.rows - 1, row + dr))
        new_col = max(0, min(self.cols - 1, col + dc))

        # 检查是否落入悬崖
        if new_row == self.rows - 1 and 0 < new_col < self.cols - 1:
            return self.pos_to_state(self.rows - 1, 0), -100, True

        # 检查是否到达终点
        new_state = self.pos_to_state(new_row, new_col)
        done = (new_state == self.pos_to_state(self.rows - 1, self.cols - 1))
        reward = 0 if done else -1

        return new_state, reward, done

mdp = CliffWalkingMDP()
state = mdp.pos_to_state(3, 0)  # 起点
print(f"起点状态: {state}")
next_state, reward, done = mdp.step(state, 1)  # 向右
print(f"向右走: 状态 {state} → {next_state}, 奖励={reward}")

# Bellman 方程的数值求解：值迭代

import numpy as np

def value_iteration_simple(n_states, transitions, rewards, gamma=0.9,
                           theta=1e-6, max_iter=1000):
    """值迭代求解 Bellman 最优方程
    
    transitions[s][a] = [(prob, next_state, reward), ...]
    """
    V = np.zeros(n_states)
    n_actions = len(transitions[0])

    for iteration in range(max_iter):
        delta = 0
        V_new = np.zeros(n_states)

        for s in range(n_states):
            q_values = np.zeros(n_actions)
            for a in range(n_actions):
                for prob, sp, r in transitions[s][a]:
                    q_values[a] += prob * (r + gamma * V[sp])
            V_new[s] = np.max(q_values)
            delta = max(delta, abs(V_new[s] - V[s]))

        V = V_new
        if delta < theta:
            print(f"收敛！迭代 {iteration+1} 次")
            break

    # 从最优值函数提取最优策略
    policy = np.zeros(n_states, dtype=int)
    for s in range(n_states):
        q_values = np.zeros(n_actions)
        for a in range(n_actions):
            for prob, sp, r in transitions[s][a]:
                q_values[a] += prob * (r + gamma * V[sp])
        policy[s] = np.argmax(q_values)

    return V, policy

# 简单 3 状态示例
transitions = [
    # 状态 0
    [
        [(1.0, 1, 5)],   # 动作 0: 一定到状态 1，奖励 5
        [(0.5, 0, 0), (0.5, 2, 10)]  # 动作 1: 50% 留在 0, 50% 到 2
    ],
    # 状态 1
    [
        [(1.0, 2, 3)],
        [(1.0, 0, 1)]
    ],
    # 状态 2 (终止状态)
    [
        [(1.0, 2, 0)],
        [(1.0, 2, 0)]
    ]
]

V, policy = value_iteration_simple(3, transitions, None, gamma=0.9)
print(f"最优值函数: V = {V}")
print(f"最优策略: pi = {policy}")

# Q-Learning 从零实现
import numpy as np
import gymnasium as gym

def q_learning(env, episodes=1000, alpha=0.1, gamma=0.99,
               epsilon=1.0, epsilon_min=0.01, epsilon_decay=0.995):
    """Q-Learning 算法实现"""
    n_states = env.observation_space.n
    n_actions = env.action_space.n
    Q = np.zeros((n_states, n_actions))
    episode_rewards = []

    for ep in range(episodes):
        state, _ = env.reset()
        total_reward = 0
        terminated = False
        truncated = False

        while not (terminated or truncated):
            # Epsilon-greedy 选择动作
            if np.random.random() < epsilon:
                action = env.action_space.sample()  # 探索
            else:
                action = np.argmax(Q[state])  # 利用

            next_state, reward, terminated, truncated, _ = env.step(action)

            # Q-Learning 核心更新
            if terminated:
                td_target = reward  # 终止状态没有未来值
            else:
                td_target = reward + gamma * np.max(Q[next_state])

            td_error = td_target - Q[state, action]
            Q[state, action] += alpha * td_error

            state = next_state
            total_reward += reward

        episode_rewards.append(total_reward)
        epsilon = max(epsilon_min, epsilon * epsilon_decay)

        if (ep + 1) % 100 == 0:
            avg_reward = np.mean(episode_rewards[-100:])
            print(f"Episode {ep+1:5d} | Avg Reward: {avg_reward:.2f} | "
                  f"Epsilon: {epsilon:.3f}")

    return Q, episode_rewards

# 在 FrozenLake 上运行
env = gym.make("FrozenLake-v1", is_slippery=True)
Q, rewards = q_learning(env, episodes=2000, alpha=0.1)
print(f"\nQ 表形状: {Q.shape}")
print(f"最终 100 集平均奖励: {np.mean(rewards[-100:]):.3f}")

# 探索-利用权衡的可视化

import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

def plot_epsilon_decay():
    """展示不同 epsilon 衰减策略的效果"""
    episodes = np.arange(1, 501)

    # 三种衰减策略
    eps_fast = np.maximum(0.01, 1.0 * (0.990 ** episodes))
    eps_medium = np.maximum(0.01, 1.0 * (0.995 ** episodes))
    eps_slow = np.maximum(0.01, 1.0 * (0.998 ** episodes))

    fig, ax = plt.subplots(figsize=(10, 6))
    ax.plot(episodes, eps_fast, label='快速衰减 (0.990)', linewidth=2)
    ax.plot(episodes, eps_medium, label='中速衰减 (0.995)', linewidth=2)
    ax.plot(episodes, eps_slow, label='慢速衰减 (0.998)', linewidth=2)
    ax.set_xlabel('Episode')
    ax.set_ylabel('Epsilon (探索概率)')
    ax.set_title('Epsilon-Greedy 探索率衰减曲线')
    ax.legend()
    ax.grid(True, alpha=0.3)
    ax.set_ylim(0, 1.05)
    plt.savefig('epsilon_decay.png', dpi=100)
    print("已保存 epsilon_decay.png")

plot_epsilon_decay()

# DQN 核心组件实现

import torch
import torch.nn as nn
import numpy as np
from collections import deque

class DQNNetwork(nn.Module):
    """简单的 DQN 网络（MLP 版本，适合离散状态空间）"""
    def __init__(self, state_dim, action_dim, hidden_dim=128):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim)
        )

    def forward(self, x):
        return self.network(x)

    def get_action(self, state, epsilon=0.1):
        """Epsilon-greedy 动作选择"""
        if np.random.random() < epsilon:
            return np.random.randint(self.network[-1].out_features)
        with torch.no_grad():
            state = torch.FloatTensor(state).unsqueeze(0)
            q_values = self.network(state)
            return q_values.argmax().item()


class ReplayBuffer:
    """经验回放缓冲区"""
    def __init__(self, capacity=10000):
        self.buffer = deque(maxlen=capacity)

    def push(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size=64):
        indices = np.random.choice(len(self.buffer), batch_size, replace=False)
        batch = [self.buffer[i] for i in indices]
        states, actions, rewards, next_states, dones = zip(*batch)
        return (torch.FloatTensor(states),
                torch.LongTensor(actions),
                torch.FloatTensor(rewards),
                torch.FloatTensor(next_states),
                torch.FloatTensor(dones))

    def __len__(self):
        return len(self.buffer)

# DQN 训练循环

def train_dqn(env, episodes=500, batch_size=64, gamma=0.99,
              lr=1e-3, target_update=10, epsilon_start=1.0,
              epsilon_end=0.01, epsilon_decay=0.995):
    """DQN 完整训练流程"""
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n

    # 主网络和目标网络
    q_network = DQNNetwork(state_dim, action_dim)
    target_network = DQNNetwork(state_dim, action_dim)
    target_network.load_state_dict(q_network.state_dict())

    optimizer = torch.optim.Adam(q_network.parameters(), lr=lr)
    replay_buffer = ReplayBuffer(capacity=10000)
    epsilon = epsilon_start

    for ep in range(episodes):
        state, _ = env.reset()
        total_reward = 0

        while True:
            # 选择动作
            action = q_network.get_action(state, epsilon)

            # 执行动作
            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated

            # 存储经验
            replay_buffer.push(state, action, reward, next_state, float(done))
            state = next_state
            total_reward += reward

            # 从缓冲区采样并训练
            if len(replay_buffer) >= batch_size:
                states, actions, rewards, next_states, dones =                     replay_buffer.sample(batch_size)

                # 当前 Q 值
                current_q = q_network(states).gather(
                    1, actions.unsqueeze(1)).squeeze(1)

                # 目标 Q 值（用目标网络）
                with torch.no_grad():
                    max_next_q = target_network(next_states).max(1)[0]
                    target_q = rewards + gamma * max_next_q * (1 - dones)

                # 更新
                loss = nn.MSELoss()(current_q, target_q)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

            if done:
                break

        epsilon = max(epsilon_end, epsilon * epsilon_decay)

        # 更新目标网络
        if (ep + 1) % target_update == 0:
            target_network.load_state_dict(q_network.state_dict())

        if (ep + 1) % 50 == 0:
            print(f"Episode {ep+1} | Reward: {total_reward:.1f} | "
                  f"Epsilon: {epsilon:.3f}")

    return q_network

# REINFORCE 算法实现

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class PolicyNetwork(nn.Module):
    """策略网络：输出动作的概率分布"""
    def __init__(self, state_dim, action_dim, hidden_dim=64):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim)
        )

    def forward(self, x):
        logits = self.network(x)
        return F.softmax(logits, dim=-1)

    def select_action(self, state):
        """根据策略采样动作，同时返回 log_prob 用于训练"""
        state = torch.FloatTensor(state)
        probs = self.forward(state)
        dist = torch.distributions.Categorical(probs)
        action = dist.sample()
        return action.item(), dist.log_prob(action)


def reinforce_train(env, episodes=500, lr=1e-2, gamma=0.99):
    """REINFORCE 训练循环"""
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n
    policy = PolicyNetwork(state_dim, action_dim)
    optimizer = torch.optim.Adam(policy.parameters(), lr=lr)

    for ep in range(episodes):
        state, _ = env.reset()
        log_probs, rewards = [], []
        total_reward = 0

        while True:
            action, log_prob = policy.select_action(state)
            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated

            log_probs.append(log_prob)
            rewards.append(reward)
            total_reward += reward
            state = next_state

            if done:
                break

        # 计算折扣回报
        returns = []
        G = 0
        for r in reversed(rewards):
            G = r + gamma * G
            returns.insert(0, G)
        returns = torch.tensor(returns)

        # 标准化回报（降低方差）
        returns = (returns - returns.mean()) / (returns.std() + 1e-8)

        # 策略梯度更新
        policy_loss = torch.stack([
            -lp * ret for lp, ret in zip(log_probs, returns)
        ]).sum()

        optimizer.zero_grad()
        policy_loss.backward()
        optimizer.step()

        if (ep + 1) % 50 == 0:
            print(f"Episode {ep+1} | Reward: {total_reward:.1f}")

    return policy

# PPO 核心组件

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class ActorCritic(nn.Module):
    """PPO 的 Actor-Critic 网络"""
    def __init__(self, state_dim, action_dim, hidden_dim=128):
        super().__init__()
        self.shared = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
        )
        # Actor：策略输出
        self.actor = nn.Linear(hidden_dim, action_dim)
        # Critic：状态值估计
        self.critic = nn.Linear(hidden_dim, 1)

    def forward(self, state):
        features = self.shared(state)
        action_probs = F.softmax(self.actor(features), dim=-1)
        state_value = self.critic(features)
        return action_probs, state_value

    def evaluate(self, state, action):
        """评估动作的对数概率和状态值"""
        action_probs, state_value = self.forward(state)
        dist = torch.distributions.Categorical(action_probs)
        log_prob = dist.log_prob(action)
        entropy = dist.entropy()
        return log_prob, state_value, entropy


def compute_gae(rewards, values, dones, gamma=0.99, lam=0.95):
    """广义优势估计（GAE）"""
    advantages = []
    gae = 0
    for t in reversed(range(len(rewards))):
        # TD 误差: delta = r + gamma*V(s')*not_done - V(s)
        next_value = 0.0 if t == len(rewards) - 1 else values[t + 1]
        next_value *= (1 - dones[t])
        delta = rewards[t] + gamma * next_value - values[t]
        gae = delta + gamma * lam * (1 - dones[t]) * gae
        advantages.insert(0, gae)

    advantages = torch.tensor(advantages)
    returns = advantages + torch.tensor(values)
    return advantages, returns

# PPO 更新步骤

def ppo_update(policy, states, actions, old_log_probs, advantages, returns,
               clip_epsilon=0.2, vf_coeff=0.5, ent_coeff=0.01,
               epochs=4, batch_size=64):
    """PPO 多轮更新：在收集的数据上做多个 epoch 的 SGD"""
    optimizer = torch.optim.Adam(policy.parameters(), lr=3e-4)
    n = len(states)

    for epoch in range(epochs):
        # 随机打乱
        indices = torch.randperm(n)
        for start in range(0, n, batch_size):
            end = start + batch_size
            batch_idx = indices[start:end]

            b_states = states[batch_idx]
            b_actions = actions[batch_idx]
            b_old_log_probs = old_log_probs[batch_idx]
            b_advantages = advantages[batch_idx]
            b_returns = returns[batch_idx]

            # 评估当前策略
            log_probs, state_values, entropy = policy.evaluate(
                b_states, b_actions)

            # 概率比
            ratio = torch.exp(log_probs - b_old_log_probs)

            # Clip 目标
            surr1 = ratio * b_advantages
            surr2 = torch.clamp(ratio, 1 - clip_epsilon, 1 + clip_epsilon) * b_advantages
            actor_loss = -torch.min(surr1, surr2).mean()

            # Critic 损失
            critic_loss = F.mse_loss(state_values.squeeze(), b_returns)

            # 熵正则化（鼓励探索）
            entropy_loss = -entropy.mean() * ent_coeff

            # 总损失
            loss = actor_loss + vf_coeff * critic_loss + entropy_loss

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

import gymnasium as gym
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from collections import namedtuple

# 定义经验存储结构
Transition = namedtuple('Transition',
    ['state', 'action', 'log_prob', 'reward', 'value', 'done'])

class ActorCritic(nn.Module):
    def __init__(self, state_dim=4, action_dim=2, hidden=64):
        super().__init__()
        self.backbone = nn.Sequential(
            nn.Linear(state_dim, hidden),
            nn.Tanh(),
            nn.Linear(hidden, hidden),
            nn.Tanh(),
        )
        self.actor = nn.Linear(hidden, action_dim)
        self.critic = nn.Linear(hidden, 1)

    def forward(self, x):
        h = self.backbone(x)
        return F.softmax(self.actor(h), dim=-1), self.critic(h)

    def act(self, state):
        state = torch.FloatTensor(state)
        probs, value = self.forward(state)
        dist = torch.distributions.Categorical(probs)
        action = dist.sample()
        return action.item(), dist.log_prob(action), value.item()


def ppo_cartpole(episodes=500, gamma=0.99, lam=0.95,
                 clip_eps=0.2, lr=3e-4, epochs=4, batch=32):
    """用 PPO 训练 CartPole"""
    env = gym.make("CartPole-v1")
    policy = ActorCritic()
    optimizer = torch.optim.Adam(policy.parameters(), lr=lr)
    episode_rewards = []

    for ep in range(episodes):
        state, _ = env.reset()
        transitions = []
        total_reward = 0

        # 收集一集数据
        while True:
            action, log_prob, value = policy.act(state)
            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            transitions.append(Transition(
                state, action, log_prob, reward, value, float(done)
            ))
            total_reward += reward
            state = next_state
            if done:
                break

        episode_rewards.append(total_reward)

        # 计算 GAE
        rewards = [t.reward for t in transitions]
        values = [t.value for t in transitions]
        dones = [t.done for t in transitions]

        advantages, returns = compute_gae(rewards, values, dones, gamma, lam)

        # 转换为 Tensor
        states = torch.FloatTensor([t.state for t in transitions])
        actions = torch.LongTensor([t.action for t in transitions])
        old_log_probs = torch.stack([t.log_prob for t in transitions])

        # PPO 多轮更新
        for _ in range(epochs):
            log_probs, state_values, _ = policy.evaluate(states, actions)
            ratio = torch.exp(log_probs - old_log_probs)
            surr1 = ratio * advantages
            surr2 = torch.clamp(ratio, 1-clip_eps, 1+clip_eps) * advantages
            actor_loss = -torch.min(surr1, surr2).mean()
            critic_loss = F.mse_loss(state_values.squeeze(), returns)
            loss = actor_loss + 0.5 * critic_loss

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

        if (ep + 1) % 50 == 0:
            avg = np.mean(episode_rewards[-50:])
            print(f"Episode {ep+1} | Avg Reward: {avg:.1f} | "
                  f"Last: {total_reward}")

    return policy, episode_rewards

# 辅助函数（需要在 ActorCritic 中添加 evaluate 方法）
print("完整代码包含 GAE 计算和 evaluate 方法，见上方 PPO 章节")
print("在本地运行需要安装: pip install gymnasium torch numpy")