如何基于FrozenLake冰湖游戏的Q-Learning代码修改实现SARSA算法

你的原代码是基于模型的Q值迭代实现，要改成SARSA仅需要针对两种算法的核心差异做3处调整：

改动点1：给动作选择添加epsilon贪婪探索 #

SARSA是同策略算法，要求训练时的采样策略和更新价值用的策略保持一致，所以不能完全用随机采样+纯贪婪选动作的组合，需要给动作选择逻辑添加可控的探索概率。

改动点2：调整价值更新逻辑 #

原Q-Learning更新时对下一状态取最大动作价值，SARSA需要按当前动作选择策略，计算下一状态所有可能动作的期望价值，替换原来的max取值逻辑。

改动点3：调整采样逻辑（可选，提升收敛速度） #

原来的play_n_random_steps完全随机采样，可以改成用epsilon贪婪策略采样，和更新策略保持一致，符合同策略要求，收敛更快。

#!/usr/bin/env python3
import gym
import collections
import random
from torch.utils.tensorboard import SummaryWriter

ENV_NAME = "FrozenLake-v1" # 原v0版本已废弃，换成v1版本兼容新gym
GAMMA = 0.9
TEST_EPISODES = 150
EPSILON = 0.1 # 探索概率

class Agent:
    def __init__(self):
        self.env = gym.make(ENV_NAME)
        self.state = self.env.reset()
        self.rewards = collections.defaultdict(float)
        self.transits = collections.defaultdict(collections.Counter)
        self.values = collections.defaultdict(float)

    def play_n_steps(self, count):
        # 改动点3：把完全随机采样改成epsilon贪婪采样，符合同策略要求
        for _ in range(count):
            action = self.select_action(self.state, epsilon=EPSILON)
            new_state, reward, is_done, _ = self.env.step(action)
            self.rewards[(self.state, action, new_state)] = reward
            self.transits[(self.state, action)][new_state] += 1
            self.state = self.env.reset() if is_done else new_state

    def select_action(self, state, epsilon=0.0):
        # 改动点1：添加epsilon贪婪逻辑，epsilon=0时就是纯贪婪用于测试
        if random.random() < epsilon:
            return self.env.action_space.sample()
        best_action, best_value = None, None
        for action in range(self.env.action_space.n):
            action_value = self.values[(state, action)]
            if best_value is None or best_value < action_value:
                best_value = action_value
                best_action = action
        return best_action

    def play_episode(self, env):
        total_reward = 0.0
        state = env.reset()
        while True:
            action = self.select_action(state) # 测试用纯贪婪
            new_state, reward, is_done, _ = env.step(action)
            total_reward += reward
            if is_done:
                break
            state = new_state
        return total_reward

    def sarsa_iteration(self):
        # 改动点2：替换Q值迭代逻辑为SARSA迭代逻辑
        for state in range(self.env.observation_space.n):
            for action in range(self.env.action_space.n):
                action_value = 0.0
                target_counts = self.transits[(state, action)]
                total = sum(target_counts.values())
                for tgt_state, count in target_counts.items():
                    reward = self.rewards[(state, action, tgt_state)]
                    # 这里不再取最大Q值，而是按当前策略计算下一状态动作价值的期望
                    num_actions = self.env.action_space.n
                    best_action = self.select_action(tgt_state)
                    expected_q = 0.0
                    for a in range(num_actions):
                        if a == best_action:
                            prob = 1 - EPSILON + EPSILON / num_actions
                        else:
                            prob = EPSILON / num_actions
                        expected_q += prob * self.values[(tgt_state, a)]
                    action_value += (count / total) * (reward + GAMMA * expected_q)
                self.values[(state, action)] = action_value


if __name__ == "__main__":
    test_env = gym.make(ENV_NAME)
    agent = Agent()
    writer = SummaryWriter(comment="-sarsa-iteration")

    iter_no = 0
    best_reward = 0.0
    while True:
        iter_no += 1
        agent.play_n_steps(100)
        agent.sarsa_iteration()

        reward = 0.0
        for _ in range(TEST_EPISODES):
            reward += agent.play_episode(test_env)
        reward /= TEST_EPISODES
        writer.add_scalar("reward", reward, iter_no)
        if reward > best_reward:
            print("最佳奖励更新 %.3f -> %.3f" % (best_reward, reward))
            best_reward = reward
        if reward > 0.80:
            print("在 %d 轮迭代后解决问题！" % iter_no)
            break
    writer.close()

内容的提问来源于stack exchange，提问作者ronanwa