05 - 蒙特卡洛方法¶

import numpy as np
from collections import defaultdict
from typing import Callable, List, Tuple

def generate_episode(env, policy):
    """
    生成一个episode

    参数:
        env: 环境，需要有reset()和step(action)方法
        policy: 策略函数，输入状态返回动作概率分布或动作
    返回:
        episode: [(state, action, reward), ...]
    """
    episode = []
    state, _ = env.reset()
    done = False

    while not done:
        # 根据策略选择动作
        if callable(policy):
            if hasattr(policy, 'shape'):  # 如果是数组形式的策略
                action = np.random.choice(len(policy[state]), p=policy[state])
            else:
                action = policy(state)
        else:
            action = policy

        next_state, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated
        episode.append((state, action, reward))
        state = next_state

        # 防止无限循环
        if len(episode) > 10000:
            break

    return episode

class VarianceReducedMC:
    """带方差缩减技术的蒙特卡洛方法"""

    def __init__(self, n_actions: int, gamma: float = 0.9):
        self.n_actions = n_actions
        self.gamma = gamma
        self.V = defaultdict(float)  # defaultdict访问不存在的键时返回默认值
        self.N = defaultdict(int)  # 访问计数
        self.returns_sum = defaultdict(float)
        self.returns_sq_sum = defaultdict(float)  # 用于方差估计

    def compute_return(self, rewards: List[float], start_idx: int = 0) -> float:
        """计算从start_idx开始的折扣回报"""
        G = 0.0
        for i in range(len(rewards) - 1, start_idx - 1, -1):
            G = self.gamma * G + rewards[i]
        return G

    def update_with_baseline(self, state: int, G: float,
                            baseline: float = None) -> float:
        """使用基线控制的增量更新"""
        if baseline is None:
            baseline = self.V[state]

        self.N[state] += 1
        n = self.N[state]

        # 增量均值更新（标准MC）
        self.V[state] += (G - self.V[state]) / n

        # 记录用于方差估计
        self.returns_sum[state] += G
        self.returns_sq_sum[state] += G * G

        return self.V[state]

    def get_variance_estimate(self, state: int) -> float:
        """估计回报的方差"""
        n = self.N[state]
        if n < 2:
            return float('inf')

        mean_sq = self.returns_sq_sum[state] / n
        mean = self.returns_sum[state] / n
        return mean_sq - mean ** 2

    def sample_complexity_for_accuracy(self, state: int,
                                       epsilon: float,
                                       delta: float = 0.05) -> int:
        """
        计算达到(epsilon, delta)精度所需的样本数

        使用切比雪夫不等式：P(|V_hat - V| >= epsilon) <= Var/(n*epsilon^2)

        参数:
            epsilon: 精度要求
            delta: 失败概率
        返回:
            所需样本数
        """
        var_estimate = self.get_variance_estimate(state)
        # 切比雪夫: n >= Var / (delta * epsilon^2)
        return int(np.ceil(var_estimate / (delta * epsilon ** 2)))

def weighted_importance_sampling_update(Q: dict, C: dict,
                                        state: int, action: int,
                                        G: float, rho: float) -> float:
    """
    加权重要性采样的增量更新

    参数:
        Q: 动作值函数
        C: 累积权重
        state, action: 状态动作对
        G: 回报
        rho: 重要性采样比率
    返回:
        更新后的Q值
    """
    sa = (state, action)
    C[sa] = C.get(sa, 0.0) + rho

    if C[sa] > 0:
        Q[sa] = Q.get(sa, 0.0) + (rho / C[sa]) * (G - Q.get(sa, 0.0))

    return Q[sa]

def stratified_mc_prediction(env, policy, strata_fn: Callable,
                             gamma: float = 0.9,
                             samples_per_stratum: int = 1000) -> dict:
    """
    分层蒙特卡洛预测

    参数:
        env: 环境
        policy: 策略
        strata_fn: 分层函数，将状态映射到层编号
        gamma: 折扣因子
        samples_per_stratum: 每层样本数
    返回:
        V: 状态值函数估计
    """
    # 收集各层数据
    stratum_data = defaultdict(lambda: {'returns': [], 'weights': []})

    # 估计各层权重（通过采样）
    exploration_episodes = 10000
    state_counts = defaultdict(int)

    for _ in range(exploration_episodes):
        episode = generate_episode(env, policy)
        visited_states = set()
        for state, _, _ in episode:
            if state not in visited_states:
                visited_states.add(state)
                stratum_id = strata_fn(state)
                state_counts[stratum_id] += 1

    total_states = sum(state_counts.values())
    stratum_weights = {k: v / total_states for k, v in state_counts.items()}

    # 分层采样
    for stratum_id in stratum_weights.keys():
        returns = []
        for _ in range(samples_per_stratum):
            episode = generate_episode(env, policy)

            # 反向计算回报
            G = 0
            for t in range(len(episode) - 1, -1, -1):
                state, action, reward = episode[t]
                G = gamma * G + reward

                if strata_fn(state) == stratum_id:
                    returns.append((state, G))

        stratum_data[stratum_id]['returns'] = returns

    # 合并各层估计
    V = defaultdict(float)
    for stratum_id, data in stratum_data.items():
        weight = stratum_weights[stratum_id]
        for state, G in data['returns']:
            V[state] += weight * G / len(data['returns'])

    return V

def off_policy_mc_control(env, behavior_policy, gamma=0.9, num_episodes=10000):
    """
    Off-Policy MC控制（使用加权重要性采样）

    参数:
        env: 环境
        behavior_policy: 行为策略（用于生成数据）
        gamma: 折扣因子
        num_episodes: episode数量
    返回:
        Q: 最优动作值函数估计
        target_policy: 目标策略（确定性贪心）
    """
    # 初始化
    Q = defaultdict(lambda: np.zeros(env.n_actions))
    C = defaultdict(lambda: np.zeros(env.n_actions))  # 累积权重
    target_policy = {}

    for i_episode in range(num_episodes):
        # 使用行为策略生成episode
        episode = generate_episode(env, behavior_policy)

        G = 0.0
        W = 1.0  # 重要性采样权重

        # 反向遍历
        for t in range(len(episode) - 1, -1, -1):
            state, action, reward = episode[t]

            G = gamma * G + reward

            # 更新累积权重
            C[state][action] += W

            # 加权重要性采样更新
            if C[state][action] > 0:
                Q[state][action] += (W / C[state][action]) * (G - Q[state][action])

            # 更新目标策略（确定性贪心）
            target_policy[state] = np.argmax(Q[state])

            # 如果行为策略选择的动作不是目标策略会选择的，终止
            if action != target_policy[state]:
                break

            # 更新权重
            # 注意：这里假设目标策略是确定性的
            # 如果是随机策略，需要乘以 pi(action|state) / b(action|state)
            W *= 1.0 / behavior_policy(state)[action]

            # 防止权重爆炸
            if W > 1e6:
                break

        if (i_episode + 1) % 1000 == 0:
            print(f"Episode {i_episode + 1}/{num_episodes} 完成")

    return Q, target_policy