强化学习入门：MDP 与 Bellman 方程

# 监督学习 vs 强化学习的训练循环对比

# 监督学习：有明确标签，直接计算损失
def supervised_train(model, X, y):
    predictions = model.predict(X)
    loss = mean_squared_error(y, predictions)
    model.backward(loss)  # 有标准答案直接回传

# 强化学习：只有延迟奖励，需要策略评估
def rl_train(agent, env, episodes=100):
    for ep in range(episodes):
        state = env.reset()
        done = False
        total_reward = 0
        while not done:
            action = agent.choose_action(state)
            next_state, reward, done, _ = env.step(action)
            agent.update(state, action, reward, next_state)
            state = next_state
            total_reward += reward

# 信用分配问题示例：智能体走了10步才吃到奖励
# 哪一步是关键动作？需要 Bellman 方程来解答

rewards = [0, 0, 0, 0, 0, 0, 0, 0, 0, +10]  # 第10步才有奖励

# 朴素方法：把所有奖励平均分配给每个动作（不推荐）
def naive_credit_assign(rewards):
    avg = sum(rewards) / len(rewards)
    return [avg] * len(rewards)

# 折扣回报：越近的动作权重越大
def discounted_return(rewards, gamma=0.9):
    returns = []
    G = 0
    for r in reversed(rewards):
        G = r + gamma * G
        returns.insert(0, G)
    return returns

print(naive_credit_assign(rewards))     # [1.0, 1.0, ..., 1.0]
print(discounted_return(rewards, 0.9))  # 越后面的值越大

# 简单的 MDP 示例：学生日常生活决策
# 状态: 上课/刷手机/学习/通过考试
# 动作: 认真听/走神/复习/放弃

class SimpleMDP:
    def __init__(self):
        self.states = ["上课", "刷手机", "学习", "通过考试", "挂科"]
        self.actions = ["认真听", "走神", "复习", "放弃"]
        # 状态转移概率 P[s][a] = [(prob, next_s, reward), ...]
        self.P = {
            "上课": {
                "认真听": [(0.8, "学习", 2), (0.2, "刷手机", -1)],
                "走神": [(0.7, "刷手机", -1), (0.3, "学习", 1)],
            },
            "刷手机": {
                "复习": [(0.6, "学习", 3), (0.4, "刷手机", -1)],
                "放弃": [(1.0, "挂科", -10)],
            },
            "学习": {
                "复习": [(0.9, "通过考试", 10), (0.1, "学习", 0)],
                "放弃": [(0.5, "刷手机", -1), (0.5, "学习", 0)],
            },
        }
        self.gamma = 0.9  # 折扣因子

    def step(self, state, action):
        transitions = self.P[state][action]
        probs, next_states, rewards = zip(*transitions)
        idx = np.random.choice(len(probs), p=probs)
        return next_states[idx], rewards[idx]

# 马尔可夫性验证：下一状态是否只依赖当前状态和动作？

import numpy as np

def check_markov_property(history_data):
    """检查历史数据是否满足马尔可夫性
    比较 P(s'|s,a) 和 P(s'|s,a,history) 是否一致
    """
    # 构建条件概率分布
    p_given_current = {}   # P(s'|s,a)
    p_given_history = {}   # P(s'|s,a,s_prev)

    for h in history_data:
        s_prev, s_curr, action, s_next = h
        key_current = (s_curr, action)
        key_history = (s_prev, s_curr, action)

        p_given_current.setdefault(key_current, {})
        p_given_current[key_current][s_next] =             p_given_current[key_current].get(s_next, 0) + 1

        p_given_history.setdefault(key_history, {})
        p_given_history[key_history][s_next] =             p_given_history[key_history].get(s_next, 0) + 1

    # 如果两种条件下概率分布差异大，说明不满足马尔可夫性
    return p_given_current, p_given_history

# 策略与值函数的实现

import numpy as np

class Policy:
    """随机策略：每个动作有固定概率"""
    def __init__(self, n_actions):
        self.probs = np.ones(n_actions) / n_actions  # 均匀初始

    def choose_action(self, state):
        return np.random.choice(len(self.probs), p=self.probs)

    def update(self, state, action, advantage, lr=0.1):
        """REINFORCE 风格的策略更新"""
        self.probs[action] += lr * advantage
        self.probs = np.clip(self.probs, 1e-6, None)
        self.probs /= self.probs.sum()  # 重新归一化


class ValueFunction:
    """状态值函数 V(s)"""
    def __init__(self, n_states):
        self.V = np.zeros(n_states)

    def update(self, state, td_error, lr=0.1):
        """时序差分更新：V(s) += alpha * (r + gamma*V(s') - V(s))"""
        self.V[state] += lr * td_error

    def get_value(self, state):
        return self.V[state]

# Q 函数：评估状态-动作对的长期价值

class QFunction:
    """动作值函数 Q(s, a)"""
    def __init__(self, n_states, n_actions):
        self.Q = np.zeros((n_states, n_actions))

    def choose_action(self, state, epsilon=0.1):
        """epsilon-greedy 策略：大部分时间选最优，偶尔探索"""
        if np.random.random() < epsilon:
            return np.random.randint(self.Q.shape[1])
        return np.argmax(self.Q[state])

    def update(self, state, action, reward, next_state, gamma=0.9, lr=0.1):
        """Q-learning 更新规则"""
        best_next = np.max(self.Q[next_state])
        td_target = reward + gamma * best_next
        td_error = td_target - self.Q[state, action]
        self.Q[state, action] += lr * td_error
        return td_error

    def get_q(self, state, action):
        return self.Q[state, action]

# Bellman 期望方程：计算给定策略下的状态值

def bellman_expectation(V, state, policy, P, R, gamma=0.9):
    """
    V(s) = sum_a pi(a|s) * sum_{s'} P(s'|s,a) * [R(s,a,s') + gamma * V(s')]
    """
    v = 0.0
    for action, prob_a in enumerate(policy[state]):
        for next_state, prob_s in P[state][action].items():
            reward = R[state][action].get(next_state, 0)
            v += prob_a * prob_s * (reward + gamma * V[next_state])
    return v

# 迭代计算值函数直到收敛
def compute_value_function(policy, P, R, gamma=0.9, theta=1e-6):
    n_states = len(policy)
    V = np.zeros(n_states)
    while True:
        delta = 0
        for s in range(n_states):
            v_old = V[s]
            V[s] = bellman_expectation(V, s, policy, P, R, gamma)
            delta = max(delta, abs(v_old - V[s]))
        if delta < theta:
            break
    return V

# Bellman 最优方程：寻找最优策略下的值函数

def bellman_optimality(V, state, P, R, gamma=0.9):
    """
    V*(s) = max_a sum_{s'} P(s'|s,a) * [R(s,a,s') + gamma * V*(s')]
    """
    q_values = {}
    for action in P[state]:
        q = 0.0
        for next_state, prob_s in P[state][action].items():
            reward = R[state][action].get(next_state, 0)
            q += prob_s * (reward + gamma * V[next_state])
        q_values[action] = q
    return max(q_values.values())  # 取最优动作的 Q 值

# 策略提取：给定最优值函数，提取最优策略
def extract_optimal_policy(V_star, P, R, gamma=0.9):
    n_states = len(V_star)
    n_actions = len(P[0])
    policy = np.zeros((n_states, n_actions))
    for s in range(n_states):
        q_values = []
        for a in range(n_actions):
            q = 0.0
            for sp, prob in P[s][a].items():
                r = R[s][a].get(sp, 0)
                q += prob * (r + gamma * V_star[sp])
            q_values.append(q)
        best_action = np.argmax(q_values)
        policy[s, best_action] = 1.0
    return policy

# 策略迭代算法：策略评估 + 策略改进

def policy_iteration(P, R, gamma=0.9, theta=1e-6):
    n_states = len(P)
    n_actions = len(P[0])
    V = np.zeros(n_states)
    policy = np.ones((n_states, n_actions)) / n_actions  # 随机初始策略

    while True:
        # === 策略评估 ===
        while True:
            delta = 0
            for s in range(n_states):
                v_old = V[s]
                V[s] = sum(
                    policy[s][a] * sum(
                        prob * (R[s][a].get(sp, 0) + gamma * V[sp])
                        for sp, prob in P[s][a].items()
                    )
                    for a in range(n_actions)
                )
                delta = max(delta, abs(v_old - V[s]))
            if delta < theta:
                break

        # === 策略改进 ===
        policy_stable = True
        for s in range(n_states):
            old_action = np.argmax(policy[s])
            q_values = np.zeros(n_actions)
            for a in range(n_actions):
                q_values[a] = sum(
                    prob * (R[s][a].get(sp, 0) + gamma * V[sp])
                    for sp, prob in P[s][a].items()
                )
            best_action = np.argmax(q_values)
            policy[s] = np.eye(n_actions)[best_action]
            if old_action != best_action:
                policy_stable = False

        if policy_stable:
            break
    return V, policy

# 值迭代算法：更简洁的替代方案

def value_iteration(P, R, gamma=0.9, theta=1e-6):
    """值迭代：每步直接用 Bellman 最优方程更新"""
    n_states = len(P)
    n_actions = len(P[0])
    V = np.zeros(n_states)

    while True:
        delta = 0
        for s in range(n_states):
            v_old = V[s]
            # 对每个动作计算 Q 值，取最大
            q_values = np.zeros(n_actions)
            for a in range(n_actions):
                for sp, prob in P[s][a].items():
                    reward = R[s][a].get(sp, 0)
                    q_values[a] += prob * (reward + gamma * V[sp])
            V[s] = np.max(q_values)
            delta = max(delta, abs(v_old - V[s]))
        if delta < theta:
            break

    # 从最优值函数提取最优策略
    policy = np.zeros((n_states, n_actions))
    for s in range(n_states):
        q_values = np.zeros(n_actions)
        for a in range(n_actions):
            for sp, prob in P[s][a].items():
                reward = R[s][a].get(sp, 0)
                q_values[a] += prob * (reward + gamma * V[sp])
        best = np.argmax(q_values)
        policy[s, best] = 1.0
    return V, policy

import gymnasium as gym

# 创建 Frozen Lake 环境
env = gym.make("FrozenLake-v1", desc=None, map_name="4x4", is_slippery=True)

# 查看环境信息
print(f"状态空间: {env.observation_space}")  # Discrete(16)
print(f"动作空间: {env.action_space}")        # Discrete(4): 上下左右

# 渲染地图
# SFFF       (S: 起点, 安全)
# FHFH       (F: 冻结冰面, 安全)
# FFFH       (H: 冰洞, 危险)
# HFFG       (G: 终点, 目标)

# 随机策略试玩
state, info = env.reset()
total_reward = 0
done = False
while not done:
    action = env.action_space.sample()  # 随机动作
    state, reward, terminated, truncated, info = env.step(action)
    done = terminated or truncated
    total_reward += reward
    env.render()

print(f"总奖励: {total_reward}")  # 到达终点=1.0，掉洞=0.0

# 用值迭代求解 Frozen Lake

def solve_frozen_lake(env, gamma=0.99, theta=1e-8):
    n_states = env.observation_space.n
    n_actions = env.action_space.n

    # 提取转移概率 P[s][a] = [(prob, next_s, reward, done), ...]
    P = {}
    R = {}
    for s in range(n_states):
        P[s] = {}
        R[s] = {}
        for a in range(n_actions):
            P[s][a] = {}
            R[s][a] = {}
            for prob, next_s, reward, done in env.unwrapped.P[s][a]:
                P[s][a][next_s] = prob
                R[s][a][next_s] = reward

    # 值迭代
    V = np.zeros(n_states)
    while True:
        delta = 0
        for s in range(n_states):
            v_old = V[s]
            q_vals = np.zeros(n_actions)
            for a in range(n_actions):
                for sp, prob in P[s][a].items():
                    r = R[s][a].get(sp, 0)
                    q_vals[a] += prob * (r + gamma * V[sp])
            V[s] = np.max(q_vals)
            delta = max(delta, abs(v_old - V[s]))
        if delta < theta:
            break

    # 提取策略
    policy = np.zeros(n_states, dtype=int)
    for s in range(n_states):
        q_vals = np.zeros(n_actions)
        for a in range(n_actions):
            for sp, prob in P[s][a].items():
                r = R[s][a].get(sp, 0)
                q_vals[a] += prob * (r + gamma * V[sp])
        policy[s] = np.argmax(q_vals)
    return V, policy

V_star, pi_star = solve_frozen_lake(env)
print(f"最优值函数 (起点): {V_star[0]:.4f}")
print(f"最优策略 (起点): {pi_star[0]} (0=左,1=下,2=右,3=上)")

import gymnasium as gym

# 标准交互循环（推荐写法）
def run_episode(env, policy, render=False):
    """运行一集，返回总奖励和轨迹"""
    state, info = env.reset()
    trajectory = []  # 记录 (s, a, r, s') 用于后续学习
    total_reward = 0
    terminated = False
    truncated = False

    while not (terminated or truncated):
        action = policy(state)  # 策略决定动作
        next_state, reward, terminated, truncated, info = env.step(action)
        trajectory.append((state, action, reward, next_state))
        total_reward += reward
        state = next_state
        if render:
            env.render()

    env.close()
    return total_reward, trajectory


# epsilon-greedy 策略：平衡探索与利用
def epsilon_greedy_policy(Q, state, epsilon=0.1):
    n_actions = Q.shape[1] if hasattr(Q, 'shape') else 4
    if np.random.random() < epsilon:
        return np.random.randint(n_actions)
    return np.argmax(Q[state])

# 完整的 Q-learning 训练循环

def q_learning_train(env, episodes=1000, alpha=0.1, gamma=0.99, epsilon=0.1):
    """Q-learning 训练：无需模型，直接通过与环境交互学习"""
    n_states = env.observation_space.n
    n_actions = env.action_space.n
    Q = np.zeros((n_states, n_actions))
    rewards_per_episode = []

    for ep in range(episodes):
        state, _ = env.reset()
        total_reward = 0
        terminated = False
        truncated = False

        while not (terminated or truncated):
            # epsilon-greedy 选择动作
            if np.random.random() < epsilon:
                action = np.random.randint(n_actions)
            else:
                action = np.argmax(Q[state])

            next_state, reward, terminated, truncated, _ = env.step(action)

            # Q-learning 核心更新
            best_next_q = np.max(Q[next_state]) if not terminated else 0
            td_target = reward + gamma * best_next_q
            Q[state, action] += alpha * (td_target - Q[state, action])

            state = next_state
            total_reward += reward

        rewards_per_episode.append(total_reward)

        # 逐渐减少探索
        epsilon = max(0.01, epsilon * 0.999)

    return Q, rewards_per_episode

# 训练并验证
env = gym.make("FrozenLake-v1", is_slippery=True)
Q_star, history = q_learning_train(env, episodes=2000)
print(f"最终 Q 表形状: {Q_star.shape}")
print(f"最近 100 集平均奖励: {np.mean(history[-100:]):.3f}")