人工智能之核心基础 机器学习 第十九章 强化学习入门

graph LR
A[Agent] -->|选择动作 a| B(Environment)
B -->|返回 新状态 s' 和 奖励 r| A
A -->|更新策略| A

import numpy as np
import random
from IPython.display import clear_output
import time

class GridWorld:
    def __init__(self):
        self.size = 4
        self.start = (0, 0)
        self.goal = (3, 3)
        self.state = self.start
        self.actions = ['up', 'down', 'left', 'right']
        self.action_map = {
            'up': (-1, 0),
            'down': (1, 0),
            'left': (0, -1),
            'right': (0, 1)
        }
    
    def reset(self):
        self.state = self.start
        return self.state
    
    def step(self, action):
        # 计算新位置
        dr, dc = self.action_map[action]
        new_r = self.state[0] + dr
        new_c = self.state[1] + dc
        
        # 边界检查
        if new_r < 0 or new_r >= self.size or new_c < 0 or new_c >= self.size:
            # 撞墙，留在原地
            new_state = self.state
            reward = -1
        else:
            new_state = (new_r, new_c)
            if new_state == self.goal:
                reward = 10  # 到达终点
            else:
                reward = -1  # 普通步
        
        self.state = new_state
        done = (new_state == self.goal)
        return new_state, reward, done
    
    def render(self):
        grid = np.full((self.size, self.size), '.')
        grid[self.goal] = 'G'
        grid[self.state] = 'A'
        print(grid)

def q_learning(env, episodes=500, alpha=0.1, gamma=0.99, epsilon=0.1):
    # 初始化Q表：Q[state][action] = value
    Q = {}
    for r in range(env.size):
        for c in range(env.size):
            Q[(r, c)] = {a: 0.0 for a in env.actions}
    
    rewards = []
    
    for episode in range(episodes):
        state = env.reset()
        total_reward = 0
        done = False
        
        while not done:
            # ε-贪婪策略：以ε概率随机探索，否则选择最优动作
            if random.random() < epsilon:
                action = random.choice(env.actions)
            else:
                action = max(Q[state], key=Q[state].get)
            
            next_state, reward, done = env.step(action)
            total_reward += reward
            
            # Q-Learning 更新
            best_next_action = max(Q[next_state], key=Q[next_state].get)
            td_target = reward + gamma * Q[next_state][best_next_action]
            td_error = td_target - Q[state][action]
            Q[state][action] += alpha * td_error
            
            state = next_state
        
        rewards.append(total_reward)
        
        # 每100轮显示一次
        if (episode + 1) % 100 == 0:
            avg_reward = np.mean(rewards[-100:])
            print(f"Episode {episode+1}, Average Reward: {avg_reward:.2f}")
    
    return Q, rewards

def sarsa(env, episodes=500, alpha=0.1, gamma=0.99, epsilon=0.1):
    Q = {}
    for r in range(env.size):
        for c in range(env.size):
            Q[(r, c)] = {a: 0.0 for a in env.actions}
    
    rewards = []
    
    for episode in range(episodes):
        state = env.reset()
        # 选择初始动作
        if random.random() < epsilon:
            action = random.choice(env.actions)
        else:
            action = max(Q[state], key=Q[state].get)
        
        total_reward = 0
        done = False
        
        while not done:
            next_state, reward, done = env.step(action)
            total_reward += reward
            
            # 选择下一个动作（用于SARSA更新）
            if random.random() < epsilon:
                next_action = random.choice(env.actions)
            else:
                next_action = max(Q[next_state], key=Q[next_state].get)
            
            # SARSA 更新
            td_target = reward + gamma * Q[next_state][next_action]
            td_error = td_target - Q[state][action]
            Q[state][action] += alpha * td_error
            
            state = next_state
            action = next_action
        
        rewards.append(total_reward)
        
        if (episode + 1) % 100 == 0:
            avg_reward = np.mean(rewards[-100:])
            print(f"SARSA Episode {episode+1}, Average Reward: {avg_reward:.2f}")
    
    return Q, rewards

# 创建环境
env = GridWorld()

# 训练 Q-Learning
print("Training Q-Learning...")
Q_q, rewards_q = q_learning(env, episodes=500)

# 训练 SARSA
print("\nTraining SARSA...")
Q_s, rewards_s = sarsa(env, episodes=500)

# 可视化学习曲线
import matplotlib.pyplot as plt

plt.plot(np.convolve(rewards_q, np.ones(50)/50, mode='valid'), label='Q-Learning')
plt.plot(np.convolve(rewards_s, np.ones(50)/50, mode='valid'), label='SARSA')
plt.xlabel('Episode')
plt.ylabel('Average Reward (50-episode window)')
plt.legend()
plt.title('Q-Learning vs SARSA Learning Curves')
plt.show()

# 测试学到的策略
def test_policy(Q, env, max_steps=20):
    state = env.reset()
    env.render()
    time.sleep(1)
    
    for step in range(max_steps):
        action = max(Q[state], key=Q[state].get)
        state, reward, done = env.step(action)
        clear_output(wait=True)
        env.render()
        time.sleep(0.5)
        if done:
            print("Reached goal!")
            break
    else:
        print("Max steps reached.")

print("\nTesting Q-Learning policy:")
test_policy(Q_q, env)