02 - SARSA算法：On-Policy TD控制¶

import numpy as np
from collections import defaultdict

def sarsa(env, num_episodes=1000, alpha=0.1, gamma=0.9, epsilon=0.1):
    """
    SARSA算法实现

    参数:
        env: 环境
        num_episodes: 训练的episode数量
        alpha: 学习率
        gamma: 折扣因子
        epsilon: ε-贪婪探索率

    返回:
        Q: 动作值函数
        policy: 学到的策略
        rewards_history: 每episode的总奖励
    """
    n_actions = env.get_action_space()
    Q = defaultdict(lambda: np.zeros(n_actions))  # defaultdict访问不存在的键时返回默认值
    rewards_history = []

    def epsilon_greedy(state, Q, epsilon):
        """ε-贪婪策略"""
        if np.random.random() < epsilon:
            return np.random.randint(n_actions)
        else:
            return np.argmax(Q[state])

    for episode in range(num_episodes):
        state, _ = env.reset()
        action = epsilon_greedy(state, Q, epsilon)

        total_reward = 0
        done = False

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

            # 选择下一个动作（ε-贪婪）
            next_action = epsilon_greedy(next_state, Q, epsilon)

            # SARSA更新
            if not done:
                td_target = reward + gamma * Q[next_state][next_action]
            else:
                td_target = reward

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

            # 转移到下一个状态
            state = next_state
            action = next_action

        rewards_history.append(total_reward)

        # 衰减探索率
        epsilon = max(0.01, epsilon * 0.995)

        if (episode + 1) % 100 == 0:
            avg_reward = np.mean(rewards_history[-100:])
            print(f"Episode {episode + 1}/{num_episodes}, "
                  f"Avg Reward: {avg_reward:.2f}, Epsilon: {epsilon:.3f}")

    # 提取最终策略
    policy = {}
    for state in Q.keys():
        policy[state] = np.argmax(Q[state])

    return Q, policy, rewards_history

def sarsa_lambda(env, num_episodes=1000, alpha=0.1, gamma=0.9,
                 epsilon=0.1, lambda_=0.9):
    """
    SARSA(λ)算法

    参数:
        lambda_: 迹衰减参数
    """
    n_actions = env.get_action_space()
    Q = defaultdict(lambda: np.zeros(n_actions))

    def epsilon_greedy(state, Q, epsilon):
        if np.random.random() < epsilon:
            return np.random.randint(n_actions)
        return np.argmax(Q[state])

    for episode in range(num_episodes):
        # 初始化资格迹
        E = defaultdict(lambda: np.zeros(n_actions))

        state, _ = env.reset()
        action = epsilon_greedy(state, Q, epsilon)

        done = False
        while not done:
            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            next_action = epsilon_greedy(next_state, Q, epsilon)

            # 计算TD误差
            if not done:
                td_target = reward + gamma * Q[next_state][next_action]
            else:
                td_target = reward
            td_error = td_target - Q[state][action]

            # 更新资格迹
            E[state][action] += 1

            # 更新所有Q值
            for s in E.keys():
                for a in range(n_actions):
                    Q[s][a] += alpha * td_error * E[s][a]
                    E[s][a] *= gamma * lambda_

            state = next_state
            action = next_action

    return Q

class WindyGridWorld:
    """有风网格世界"""

    def __init__(self, size=7):
        self.size = size
        self.start = (3, 0)
        self.goal = (3, 7)
        # 每列的风力（向上）
        self.wind = [0, 0, 0, 1, 1, 1, 2, 2, 1, 0]

    def step(self, action):
        x, y = self.agent_pos

        # 动作效果
        if action == 0:  # 上
            x -= 1
        elif action == 1:  # 下
            x += 1
        elif action == 2:  # 左
            y -= 1
        elif action == 3:  # 右
            y += 1

        # 加上风力
        x -= self.wind[y]

        # 边界检查
        x = np.clip(x, 0, self.size - 1)
        y = np.clip(y, 0, 9)

        self.agent_pos = (x, y)
        reward = -1
        done = (self.agent_pos == self.goal)

        return self._get_state(), reward, done, False, {}

# 用SARSA解决
env = WindyGridWorld()
Q, policy, rewards = sarsa(env, num_episodes=5000, alpha=0.5, epsilon=0.1)

def compare_algorithms(env, num_runs=10, num_episodes=1000):
    """对比SARSA和Q-Learning"""

    sarsa_rewards = []
    qlearning_rewards = []

    for run in range(num_runs):
        # SARSA
        _, _, sarsa_hist = sarsa(env, num_episodes)
        sarsa_rewards.append(sarsa_hist)

        # Q-Learning
        _, _, q_hist = q_learning(env, num_episodes)
        qlearning_rewards.append(q_hist)

    # 绘制对比图
    plt.figure(figsize=(10, 6))

    sarsa_mean = np.mean(sarsa_rewards, axis=0)
    q_mean = np.mean(qlearning_rewards, axis=0)

    plt.plot(sarsa_mean, label='SARSA', alpha=0.7)
    plt.plot(q_mean, label='Q-Learning', alpha=0.7)

    plt.xlabel('Episode')
    plt.ylabel('Total Reward')
    plt.title('SARSA vs Q-Learning')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.show()