01 - 时序差分学习基础¶

import numpy as np
from collections import defaultdict

def td_zero_prediction(env, policy, alpha=0.1, gamma=0.9, num_episodes=1000):
    """
    TD(0)策略评估

    参数:
        env: 环境
        policy: 策略，state -> action probabilities
        alpha: 学习率
        gamma: 折扣因子
        num_episodes: episode数量

    返回:
        V: 状态值函数估计
    """
    V = defaultdict(float)  # defaultdict访问不存在的键时返回默认值

    for episode in range(num_episodes):
        state, _ = env.reset()
        done = False

        while not done:
            # 根据策略选择动作
            action = np.random.choice(len(policy[state]), p=policy[state])

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

            # TD更新
            td_target = reward + gamma * V[next_state] * (not done)
            td_error = td_target - V[state]
            V[state] += alpha * td_error

            state = next_state

        if (episode + 1) % 100 == 0:
            # 计算当前策略的平均回报
            avg_reward = evaluate_policy(env, policy, V)
            print(f"Episode {episode + 1}/{num_episodes}, Avg Reward: {avg_reward:.2f}")

    return V

def evaluate_policy(env, policy, V, num_episodes=10):
    """评估策略的平均回报"""
    total_rewards = []

    for _ in range(num_episodes):
        state, _ = env.reset()
        done = False
        episode_reward = 0

        while not done:
            action = np.argmax(policy[state])
            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            episode_reward += reward
            state = next_state

        total_rewards.append(episode_reward)

    return np.mean(total_rewards)

# 1. 常数学习率
alpha = 0.1

# 2. 线性衰减
alpha = max(0.01, 0.1 - episode * 0.0001)

# 3. 指数衰减
alpha = 0.1 * (0.995 ** episode)

# 4. 理论上保证收敛
alpha = 1 / (episode + 1)

def batch_td_learning(env, policy, alpha=0.1, gamma=0.9,
                      batch_size=100, num_batches=100):
    """
    批量TD学习

    收集batch_size个transition，然后进行多次更新
    """
    V = defaultdict(float)

    for batch_idx in range(num_batches):
        # 收集数据
        transitions = []
        for _ in range(batch_size):
            state, _ = env.reset()
            done = False

            while not done and len(transitions) < batch_size:
                action = np.random.choice(len(policy[state]), p=policy[state])
                next_state, reward, terminated, truncated, _ = env.step(action)
                done = terminated or truncated
                transitions.append((state, action, reward, next_state, done))
                state = next_state

        # 批量更新（多次遍历数据）
        for _ in range(10):  # 每个batch更新10次
            for state, action, reward, next_state, done in transitions:
                td_target = reward + gamma * V[next_state] * (not done)
                td_error = td_target - V[state]
                V[state] += alpha * td_error

        if (batch_idx + 1) % 10 == 0:
            print(f"Batch {batch_idx + 1}/{num_batches}")

    return V

# 使用附录中的GridWorld环境
from gridworld import GridWorld

env = GridWorld(size=5)

# 定义随机策略
n_states = env.get_state_space()
n_actions = env.get_action_space()
random_policy = np.ones((n_states, n_actions)) / n_actions

# 运行TD(0)
V = td_zero_prediction(env, random_policy, alpha=0.1, num_episodes=1000)

# 可视化值函数
import matplotlib.pyplot as plt

V_matrix = np.array([V[s] for s in range(n_states)]).reshape(5, 5)  # 重塑张量形状  # np.array创建NumPy数组
plt.imshow(V_matrix, cmap='viridis')
plt.colorbar()
plt.title('Value Function from TD(0)')
plt.show()

alphas = [0.01, 0.05, 0.1, 0.3, 0.5]
results = {}

for alpha in alphas:
    V = td_zero_prediction(env, random_policy, alpha=alpha, num_episodes=500)
    results[alpha] = V

# 比较最终值函数的差异
for alpha, V in results.items():
    print(f"α={alpha}: Mean V = {np.mean(list(V.values())):.3f}")

# 实现MC预测（首次访问）
def mc_prediction(env, policy, gamma=0.9, num_episodes=1000):
    V = defaultdict(float)
    returns = defaultdict(list)

    for episode in range(num_episodes):
        episode_data = []
        state, _ = env.reset()
        done = False

        while not done:
            action = np.random.choice(len(policy[state]), p=policy[state])
            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            episode_data.append((state, reward))
            state = next_state

        # 计算回报
        G = 0
        visited = set()
        for t in range(len(episode_data) - 1, -1, -1):
            state, reward = episode_data[t]
            G = gamma * G + reward

            if state not in visited:  # 首次访问
                visited.add(state)
                returns[state].append(G)
                V[state] = np.mean(returns[state])

    return V

# 对比TD和MC
td_V = td_zero_prediction(env, random_policy, alpha=0.1, num_episodes=1000)
mc_V = mc_prediction(env, random_policy, gamma=0.9, num_episodes=1000)

# 计算差异
diffs = [abs(td_V[s] - mc_V[s]) for s in range(n_states)]
print(f"Mean absolute difference: {np.mean(diffs):.4f}")