策略梯度：REINFORCE 算法

# Softmax 策略：离散动作空间的经典表示

import numpy as np

class SoftmaxPolicy:
    def __init__(self, n_states, n_actions, n_features=4):
        self.theta = np.random.randn(n_actions, n_features) * 0.01

    def feature_extractor(self, state):
        return np.array([1.0, state, state**2, np.sin(state)])

    def get_preferences(self, state):
        features = self.feature_extractor(state)
        return self.theta @ features

    def get_probs(self, state):
        h = self.get_preferences(state)
        h_shifted = h - np.max(h)  # 数值稳定
        exp_h = np.exp(h_shifted)
        return exp_h / np.sum(exp_h)

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

    def log_prob(self, state, action):
        probs = self.get_probs(state)
        return np.log(probs[action] + 1e-10)

# 神经网络策略：更强大的函数近似器

class NeuralPolicy:
    def __init__(self, n_features, n_actions, hidden_size=64):
        scale1 = np.sqrt(2.0 / (n_features + hidden_size))
        scale2 = np.sqrt(2.0 / (hidden_size + n_actions))
        self.W1 = np.random.randn(n_features, hidden_size) * scale1
        self.b1 = np.zeros(hidden_size)
        self.W2 = np.random.randn(hidden_size, n_actions) * scale2
        self.b2 = np.zeros(n_actions)

    def forward(self, state):
        hidden = np.maximum(0, state @ self.W1 + self.b1)
        logits = hidden @ self.W2 + self.b2
        logits_shifted = logits - np.max(logits)
        probs = np.exp(logits_shifted) / np.sum(np.exp(logits_shifted))
        return probs, hidden, logits

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

    def log_prob(self, state, action):
        probs, _, _ = self.forward(state)
        return np.log(probs[action] + 1e-10)

# 策略梯度定理的直观理解

def policy_gradient_intuition():
    """
    策略梯度定理：
    J(theta) = E[sum_t log(pi(a_t|s_t)) * G_t]

    推导思路：
    1. J(theta) = E_pi[G_0] = sum_tau p(tau|theta) * R(tau)
    2. grad J = sum_tau grad p(tau|theta) * R(tau)
    3. 利用 grad p = p * grad log p
    4. 得到 grad J = E_tau[grad log p(tau|theta) * R(tau)]
    5. grad log p(tau) = sum_t grad log pi(a_t|s_t)
    """
    pass

def estimate_policy_gradient(states, actions, returns, agent):
    grad = np.zeros_like(agent.theta)
    for traj_s, traj_a, traj_r in zip(states, actions, returns):
        for s, a, G in zip(traj_s, traj_a, traj_r):
            probs = agent.get_probs(s)
            grad_log_pi = -probs.copy()
            grad_log_pi[a] += 1.0
            grad += np.outer(np.ones_like(agent.theta[a]), grad_log_pi * G)
    return grad / len(states)

# 策略梯度定理的数值验证

def verify_policy_gradient_theorem():
    np.random.seed(42)
    P = np.array([
        [[0.8, 0.2], [0.3, 0.7]],
        [[0.6, 0.4], [0.5, 0.5]],
    ])
    R = np.array([[1.0, 0.0], [0.5, 2.0]])
    gamma = 0.9
    theta = np.array([[0.1, -0.1], [0.0, 0.2]])

    def softmax(theta_row):
        e = np.exp(theta_row - np.max(theta_row))
        return e / e.sum()

    def expected_return(theta):
        pi = np.array([softmax(theta[0]), softmax(theta[1])])
        P_pi = np.zeros((2, 2))
        for s in range(2):
            for a in range(2):
                P_pi[s] += pi[s, a] * P[s, a]
        eigvals, eigvecs = np.linalg.eig(P_pi.T)
        stationary = np.real(eigvecs[:, np.argmax(eigvals)])
        stationary /= stationary.sum()
        R_pi = np.array([sum(pi[s,a]*R[s,a] for a in range(2)) for s in range(2)])
        V = np.linalg.solve(np.eye(2) - gamma * P_pi, R_pi)
        return stationary @ V

    eps = 1e-5
    numerical_grad = np.zeros_like(theta)
    J0 = expected_return(theta)
    for i in range(2):
        for j in range(2):
            tp = theta.copy()
            tp[i, j] += eps
            numerical_grad[i, j] = (expected_return(tp) - J0) / eps
    print(f"数值梯度:\n{numerical_grad}")

verify_policy_gradient_theorem()

# REINFORCE 算法完整实现

import numpy as np
import gymnasium as gym

class REINFORCE:
    def __init__(self, n_features, n_actions, lr=0.01, gamma=0.99):
        self.lr = lr
        self.gamma = gamma
        self.W1 = np.random.randn(n_features, 32) * 0.01
        self.b1 = np.zeros(32)
        self.W2 = np.random.randn(32, n_actions) * 0.01
        self.b2 = np.zeros(n_actions)

    def forward(self, state):
        h = np.maximum(0, state @ self.W1 + self.b1)
        logits = h @ self.W2 + self.b2
        logits -= np.max(logits)
        probs = np.exp(logits) / np.sum(np.exp(logits))
        return probs, h

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

    def compute_returns(self, rewards):
        returns = []
        G = 0
        for r in reversed(rewards):
            G = r + self.gamma * G
            returns.insert(0, G)
        returns = np.array(returns)
        return (returns - returns.mean()) / (returns.std() + 1e-8)

    def update(self, states, actions, rewards):
        returns = self.compute_returns(rewards)
        dW1, db1, dW2, db2 = 0, 0, 0, 0
        for s, a, G in zip(states, actions, returns):
            probs, h = self.forward(s)
            d_logits = -np.eye(len(probs))[a] + probs
            d_logits *= G
            dW2 += np.outer(h, d_logits)
            db2 += d_logits
            d_h = d_logits @ self.W2.T
            d_h[h <= 0] = 0
            dW1 += np.outer(s, d_h)
            db1 += d_h
        n = len(states)
        self.W1 += self.lr * dW1 / n
        self.b1 += self.lr * db1 / n
        self.W2 += self.lr * dW2 / n
        self.b2 += self.lr * db2 / n

# REINFORCE 训练循环

def train_reinforce(env_name="CartPole-v1", episodes=1000, lr=0.01, gamma=0.99):
    env = gym.make(env_name)
    n_features = env.observation_space.shape[0]
    n_actions = env.action_space.n
    agent = REINFORCE(n_features, n_actions, lr, gamma)
    history = []

    for ep in range(episodes):
        state, _ = env.reset()
        states, actions, rewards = [], [], []
        done = False
        total_reward = 0

        while not done:
            action = agent.choose_action(state)
            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            states.append(state)
            actions.append(action)
            rewards.append(reward)
            total_reward += reward
            state = next_state

        agent.update(states, actions, rewards)
        history.append(total_reward)

        if ep % 100 == 0:
            avg = np.mean(history[-100:])
            print(f"Episode {ep}: avg_reward = {avg:.1f}")

    env.close()
    return agent, history

# 基线函数：为什么减去基线不影响期望

def prove_baseline_property():
    """数学证明：E[grad log pi(a|s) * b(s)] = 0"""
    theta = np.array([0.5, -0.3, 0.8])
    probs = np.exp(theta) / np.sum(np.exp(theta))
    grad_log_pi = np.eye(3) - probs
    expected_grad = probs @ grad_log_pi
    print(f"E[grad log pi] = {expected_grad}")  # 接近零

class BaselineREINFORCE:
    def __init__(self, n_features, n_actions, lr=0.01, gamma=0.99):
        self.lr = lr
        self.gamma = gamma
        self.W1 = np.random.randn(n_features, 32) * 0.01
        self.b1 = np.zeros(32)
        self.W2 = np.random.randn(32, n_actions) * 0.01
        self.b2 = np.zeros(n_actions)
        self.Vw1 = np.random.randn(n_features, 16) * 0.01
        self.Vb1 = np.zeros(16)
        self.Vw2 = np.random.randn(16, 1) * 0.01
        self.Vb2 = np.zeros(1)

    def get_value(self, state):
        h = np.maximum(0, state @ self.Vw1 + self.Vb1)
        return float(h @ self.Vw2 + self.Vb2)

    def update_baseline(self, states, returns, lr_v=0.05):
        for s, G in zip(states, returns):
            V = self.get_value(s)
            td_error = G - V
            h = np.maximum(0, s @ self.Vw1 + self.Vb1)
            self.Vw2 += lr_v * td_error * h
            self.Vb2 += lr_v * td_error

# 基线对方差的降低效果模拟

def compare_variance_with_baseline(n_samples=10000):
    np.random.seed(42)
    Q_values = np.array([1.0, 2.0, 5.0, 0.5])
    probs = np.array([0.1, 0.2, 0.5, 0.2])

    grads_no_baseline = []
    for _ in range(n_samples):
        a = np.random.choice(4, p=probs)
        grad = (np.eye(4)[a] - probs) * Q_values[a]
        grads_no_baseline.append(grad)

    V = np.sum(probs * Q_values)
    grads_with_baseline = []
    for _ in range(n_samples):
        a = np.random.choice(4, p=probs)
        advantage = Q_values[a] - V
        grad = (np.eye(4)[a] - probs) * advantage
        grads_with_baseline.append(grad)

    var_no = np.var(grads_no_baseline, axis=0).sum()
    var_with = np.var(grads_with_baseline, axis=0).sum()
    print(f"无基线方差: {var_no:.4f}")
    print(f"有基线方差: {var_with:.4f}")
    print(f"方差降低: {(1 - var_with/var_no)*100:.1f}%")

compare_variance_with_baseline()

# 折扣回报 vs 优势函数

class DiscountedReturns:
    def __init__(self, gamma=0.99):
        self.gamma = gamma

    def compute(self, rewards):
        returns = []
        G = 0
        for r in reversed(rewards):
            G = r + self.gamma * G
            returns.insert(0, G)
        return np.array(returns)

class AdvantageEstimator:
    def __init__(self, gamma=0.99, lambda_=0.95):
        self.gamma = gamma
        self.lambda_ = lambda_

    def compute_gae(self, rewards, values, next_values, dones):
        advantages = []
        gae = 0
        for t in reversed(range(len(rewards))):
            delta = rewards[t] + self.gamma * next_values[t] * (1 - dones[t]) - values[t]
            gae = delta + self.gamma * self.lambda_ * gae * (1 - dones[t])
            advantages.insert(0, gae)
        return np.array(advantages)

    def compute_simple(self, rewards, values):
        dr = DiscountedReturns(self.gamma)
        returns = dr.compute(rewards)
        return returns - np.array(values)

# 折扣因子 gamma 对策略学习的影响

def analyze_gamma_impact():
    rewards = [0, 0, 0, 0, 0, 0, 0, 0, 0, 10]
    for gamma in [0.5, 0.9, 0.99, 0.999, 1.0]:
        G = 0
        for r in reversed(rewards):
            G = r + gamma * G
        print(f"gamma={gamma:.3f}: G_0 = {G:.2f}")

analyze_gamma_impact()

def analyze_gae_lambda():
    rewards = [1, 2, 3, 4, 5]
    values = [0.5, 1.5, 2.5, 3.5, 4.5]
    next_values = [1.5, 2.5, 3.5, 4.5, 0.0]
    dones = [0, 0, 0, 0, 1]
    for lam in [0.0, 0.5, 0.95, 1.0]:
        gae = AdvantageEstimator(gamma=0.99, lambda_=lam)
        advantages = gae.compute_gae(rewards, values, next_values, dones)
        print(f"lambda={lam:.2f}: advantages = {advantages}")

analyze_gae_lambda()

import numpy as np
import gymnasium as gym

class CartPoleREINFORCE:
    def __init__(self, lr=0.005, gamma=0.99):
        self.gamma = gamma
        self.lr = lr
        self.p_w1 = np.random.randn(4, 64) * 0.01
        self.p_b1 = np.zeros(64)
        self.p_w2 = np.random.randn(64, 2) * 0.01
        self.p_b2 = np.zeros(2)
        self.v_w1 = np.random.randn(4, 32) * 0.01
        self.v_b1 = np.zeros(32)
        self.v_w2 = np.random.randn(32, 1) * 0.01
        self.v_b2 = np.zeros(1)

    def policy_forward(self, state):
        h = np.tanh(state @ self.p_w1 + self.p_b1)
        logits = h @ self.p_w2 + self.p_b2
        logits -= np.max(logits)
        exp_l = np.exp(logits)
        probs = exp_l / exp_l.sum()
        return probs, h

    def value_forward(self, state):
        h = np.tanh(state @ self.v_w1 + self.v_b1)
        return float(h @ self.v_w2 + self.v_b2)

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

    def compute_returns(self, rewards):
        returns = []
        G = 0
        for r in reversed(rewards):
            G = r + self.gamma * G
            returns.insert(0, G)
        returns = np.array(returns)
        return (returns - returns.mean()) / (returns.std() + 1e-8)

    def update(self, states, actions, rewards):
        returns = self.compute_returns(rewards)
        for s, a, G in zip(states, actions, returns):
            probs, h = self.policy_forward(s)
            d_logits = -np.eye(2)[a] + probs
            d_logits *= G
            self.p_w2 += self.lr * np.outer(h, d_logits)
            self.p_b2 += self.lr * d_logits
            d_h = d_logits @ self.p_w2.T * (1 - h**2)
            self.p_w1 += self.lr * np.outer(s, d_h)
            self.p_b1 += self.lr * d_h
            V = self.value_forward(s)
            v_error = G - V
            vh = np.tanh(s @ self.v_w1 + self.v_b1)
            self.v_w2 += 0.01 * v_error * vh
            self.v_b2 += 0.01 * v_error

# 训练与评估

def train_cartpole(episodes=500):
    env = gym.make("CartPole-v1")
    agent = CartPoleREINFORCE()
    history = []
    for ep in range(episodes):
        state, _ = env.reset()
        states, actions, rewards = [], [], []
        total_reward = 0
        while True:
            action = agent.choose_action(state)
            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            states.append(state)
            actions.append(action)
            rewards.append(reward)
            total_reward += reward
            state = next_state
            if done:
                break
        agent.update(states, actions, rewards)
        history.append(total_reward)
        if ep % 50 == 0 and ep > 0:
            avg = np.mean(history[-50:])
            print(f"Ep {ep:4d} | avg_50 = {avg:6.1f}")
        if np.mean(history[-100:]) >= 195:
            print(f"Solved at episode {ep}!")
            break
    env.close()
    return agent, history

def evaluate(agent, n_episodes=10):
    env = gym.make("CartPole-v1")
    scores = []
    for _ in range(n_episodes):
        state, _ = env.reset()
        total = 0
        while True:
            probs, _ = agent.policy_forward(state)
            action = np.argmax(probs)
            state, r, term, trunc, _ = env.step(action)
            total += r
            if term or trunc:
                break
        scores.append(total)
    env.close()
    print(f"avg = {np.mean(scores):.1f}, min = {min(scores)}, max = {max(scores)}")
    return scores

# Actor-Critic：结合值方法和策略梯度的优势

import numpy as np

class ActorCritic:
    def __init__(self, n_features, n_actions, lr_actor=0.001, lr_critic=0.005):
        self.lr_actor = lr_actor
        self.lr_critic = lr_critic
        self.gamma = 0.99
        self.actor_w1 = np.random.randn(n_features, 64) * 0.01
        self.actor_w2 = np.random.randn(64, n_actions) * 0.01
        self.critic_w1 = np.random.randn(n_features, 64) * 0.01
        self.critic_w2 = np.random.randn(64, 1) * 0.01

    def select_action(self, state):
        h = np.tanh(state @ self.actor_w1)
        logits = h @ self.actor_w2
        logits -= np.max(logits)
        probs = np.exp(logits) / np.sum(np.exp(logits))
        return np.random.choice(len(probs), p=probs)

    def learn(self, state, action, reward, next_state, done):
        V_s = self._value(state)
        V_sp = 0.0 if done else self._value(next_state)
        td_error = reward + self.gamma * V_sp - V_s
        self._update_actor(state, action, td_error)
        self._update_critic(state, td_error)
        return td_error

    def _value(self, state):
        h = np.tanh(state @ self.critic_w1)
        return float(h @ self.critic_w2)

    def _update_actor(self, state, action, td_error):
        h = np.tanh(state @ self.actor_w1)
        probs = self._probs(state)
        d_logits = -np.eye(2)[action] + probs
        self.actor_w2 += self.lr_actor * td_error * np.outer(h, d_logits)

    def _update_critic(self, state, td_error):
        h = np.tanh(state @ self.critic_w1)
        self.critic_w2 += self.lr_critic * td_error * h

    def _probs(self, state):
        h = np.tanh(state @ self.actor_w1)
        logits = h @ self.actor_w2
        logits -= np.max(logits)
        exp_l = np.exp(logits)
        return exp_l / exp_l.sum()

# 算法选择决策树

def choose_algorithm(action_space, state_dim, need_efficiency):
    if action_space == "continuous":
        if need_efficiency:
            return ["SAC（最大熵，样本效率好）", "TD3（确定性策略梯度）"]
        return ["PPO（稳定，工业界首选）", "DDPG（基础连续控制）"]
    else:
        if state_dim > 100:
            return ["DQN/Rainbow（高维离散）", "C51（分布型 DQN）"]
        else:
            if need_efficiency:
                return ["Q-learning（表格）", "SARSA（on-policy）"]
            return ["REINFORCE（策略梯度入门）", "Actor-Critic（混合架构）"]

# 使用示例
scenarios = [
    ("continuous", 24, False, "机器人控制"),
    ("discrete", 84*84*3, True, "Atari 游戏"),
    ("discrete", 4, False, "CartPole"),
]
for action_s, state_d, eff, name in scenarios:
    recs = choose_algorithm(action_s, state_d, eff)
    print(f"\n{name}: {recs}")