Q-Learning算法实战指南：从原理到代码落地

pip install gymnasium numpy matplotlib  # Gymnasium是新版Gym，兼容旧代码

import gymnasium as gym
import numpy as np
import random
from matplotlib import pyplot as plt

# 1. 初始化环境
env = gym.make('CartPole-v1', render_mode="rgb_array")  # 用rgb_array模式，后续可渲染视频

# 2. 状态离散化函数（把连续状态切成10个区间）
def discretize_state(state):
    # 每个维度的范围（来自CartPole的官方文档）
    state_bins = [
        np.linspace(-2.4, 2.4, 10),  # 小车位置
        np.linspace(-3.0, 3.0, 10),  # 小车速度
        np.linspace(-0.209, 0.209, 10),  # 杆子角度（弧度）
        np.linspace(-3.0, 3.0, 10)   # 杆子角速度
    ]
    # 把每个维度的连续值转换成区间索引（比如位置-2.4→0，-1.2→1，…，2.4→9）
    discretized = tuple(np.digitize(s, bin) for s, bin in zip(state, state_bins))
    return discretized

# 3. 初始化Q表（用字典存储，键是离散状态，值是动作Q值的列表）
q_table = {}
# 超参数设置（我踩过无数坑后的“入门友好值”）
gamma = 0.95  # 未来奖励权重
alpha = 0.1   # 学习率
epsilon = 0.1 # 探索率（10%概率随机选动作，90%选最优）
episodes = 1000  # 训练轮数

# 4. 训练循环
reward_history = []  # 记录每轮的总奖励，方便画图
for episode in range(episodes):
    state, _ = env.reset()  # 重置环境，得到初始状态
    discretized_state = discretize_state(state)
    total_reward = 0
    done = False

    while not done:
        # 4.1 ε-贪心策略：兼顾探索（试新动作）和利用（用已知最优）
        if random.uniform(0, 1) < epsilon:
            action = env.action_space.sample()  # 随机选动作（探索）
        else:
            # 如果状态不在Q表中，初始化动作Q值为0
            if discretized_state not in q_table:
                q_table[discretized_state] = [0.0, 0.0]
            action = np.argmax(q_table[discretized_state])  # 选Q值最大的动作（利用）

        # 4.2 执行动作，得到下一个状态、奖励、是否结束
        next_state, reward, done, _, _ = env.step(action)
        discretized_next_state = discretize_state(next_state)
        total_reward += reward

        # 4.3 更新Q表
        # 如果当前状态不在Q表，初始化
        if discretized_state not in q_table:
            q_table[discretized_state] = [0.0, 0.0]
        # 如果下一个状态不在Q表，初始化（避免KeyError）
        if discretized_next_state not in q_table:
            q_table[discretized_next_state] = [0.0, 0.0]

        # 计算TD误差，更新Q值
        current_q = q_table[discretized_state][action]
        max_next_q = np.max(q_table[discretized_next_state])
        td_error = reward + gamma * max_next_q - current_q
        new_q = current_q + alpha * td_error
        q_table[discretized_state][action] = new_q

        # 4.4 切换到下一个状态
        discretized_state = discretized_next_state

    # 记录每轮奖励
    reward_history.append(total_reward)
    # 每100轮打印进度
    if (episode + 1) % 100 == 0:
        print(f"Episode {episode+1}, Average Reward: {np.mean(reward_history[-100:]):.2f}")

# 5. 可视化训练过程
plt.plot(reward_history)
plt.xlabel("Episode")
plt.ylabel("Total Reward")
plt.title("Q-Learning Training Progress (CartPole)")
plt.show()

# 6. 测试训练好的智能体（不用探索，全选最优动作）
test_episodes = 10
test_rewards = []
for _ in range(test_episodes):
    state, _ = env.reset()
    discretized_state = discretize_state(state)
    total_reward = 0
    done = False
    while not done:
        if discretized_state not in q_table:
            action = 0
        else:
            action = np.argmax(q_table[discretized_state])
        next_state, reward, done, _, _ = env.step(action)
        total_reward += reward
        discretized_state = discretize_state(next_state)
    test_rewards.append(total_reward)
print(f"Test Average Reward: {np.mean(test_rewards):.2f}")

env.close()