方策勾配法の理論と導出を解説して実装する

import numpy as np
import matplotlib.pyplot as plt

# --- CartPole環境の簡易実装 ---
class CartPole:
    """CartPole環境の簡易実装（OpenAI Gym互換のインターフェース）"""
    def __init__(self):
        self.gravity = 9.8
        self.masscart = 1.0
        self.masspole = 0.1
        self.total_mass = self.masscart + self.masspole
        self.length = 0.5  # ポールの半分の長さ
        self.polemass_length = self.masspole * self.length
        self.force_mag = 10.0
        self.tau = 0.02  # 時間刻み
        self.theta_threshold = 12 * np.pi / 180  # 角度閾値（12度）
        self.x_threshold = 2.4  # 位置閾値

    def reset(self):
        """環境リセット: 状態 = [x, x_dot, theta, theta_dot]"""
        self.state = np.random.uniform(-0.05, 0.05, size=4)
        return self.state.copy()

    def step(self, action):
        """行動(0: 左, 1: 右)を実行"""
        x, x_dot, theta, theta_dot = self.state
        force = self.force_mag if action == 1 else -self.force_mag

        cos_theta = np.cos(theta)
        sin_theta = np.sin(theta)

        # 運動方程式
        temp = (force + self.polemass_length * theta_dot**2 * sin_theta) / self.total_mass
        theta_acc = (self.gravity * sin_theta - cos_theta * temp) / (
            self.length * (4.0/3.0 - self.masspole * cos_theta**2 / self.total_mass))
        x_acc = temp - self.polemass_length * theta_acc * cos_theta / self.total_mass

        # オイラー法による状態更新
        x += self.tau * x_dot
        x_dot += self.tau * x_acc
        theta += self.tau * theta_dot
        theta_dot += self.tau * theta_acc

        self.state = np.array([x, x_dot, theta, theta_dot])

        # 終了判定
        done = bool(
            x < -self.x_threshold or x > self.x_threshold
            or theta < -self.theta_threshold or theta > self.theta_threshold
        )
        reward = 1.0 if not done else 0.0

        return self.state.copy(), reward, done

# --- 方策ネットワーク（ソフトマックス線形方策） ---
class SoftmaxPolicy:
    """ソフトマックス線形方策"""
    def __init__(self, n_features, n_actions):
        self.n_features = n_features
        self.n_actions = n_actions
        # 重みの初期化（ゼロ初期化）
        self.theta = np.zeros((n_features, n_actions))

    def _softmax(self, x):
        """数値安定なソフトマックス"""
        e_x = np.exp(x - np.max(x))
        return e_x / e_x.sum()

    def get_action_probs(self, state):
        """状態から行動確率を計算"""
        logits = state @ self.theta  # (n_actions,)
        probs = self._softmax(logits)
        return probs

    def select_action(self, state):
        """行動を確率的に選択"""
        probs = self.get_action_probs(state)
        action = np.random.choice(self.n_actions, p=probs)
        return action, probs

    def compute_grad_log_pi(self, state, action):
        """∇_θ log π_θ(a|s) を計算（ソフトマックス方策の場合）"""
        probs = self.get_action_probs(state)
        # ∇_θ log π = φ(s) (e_a - π) （ソフトマックスの勾配）
        # e_a: 行動aの位置が1のone-hotベクトル
        one_hot = np.zeros(self.n_actions)
        one_hot[action] = 1.0
        grad = np.outer(state, one_hot - probs)  # (n_features, n_actions)
        return grad

# --- REINFORCE アルゴリズム ---
def reinforce(env, n_episodes=2000, gamma=0.99, alpha=0.01):
    """REINFORCEアルゴリズム（ベースラインなし）"""
    policy = SoftmaxPolicy(n_features=4, n_actions=2)
    rewards_history = []

    for ep in range(n_episodes):
        state = env.reset()
        episode_states = []
        episode_actions = []
        episode_rewards = []
        done = False

        # エピソードの生成
        while not done:
            action, _ = policy.select_action(state)
            next_state, reward, done = env.step(action)

            episode_states.append(state)
            episode_actions.append(action)
            episode_rewards.append(reward)

            state = next_state

        # リターンの計算（後ろから）
        T = len(episode_rewards)
        G = np.zeros(T)
        G[-1] = episode_rewards[-1]
        for t in range(T - 2, -1, -1):
            G[t] = episode_rewards[t] + gamma * G[t + 1]

        # パラメータ更新
        for t in range(T):
            grad = policy.compute_grad_log_pi(episode_states[t], episode_actions[t])
            policy.theta += alpha * (gamma ** t) * G[t] * grad

        rewards_history.append(sum(episode_rewards))

    return policy, rewards_history

# --- REINFORCE with Baseline ---
def reinforce_baseline(env, n_episodes=2000, gamma=0.99, alpha_theta=0.01, alpha_w=0.05):
    """REINFORCEアルゴリズム（ベースライン付き）"""
    policy = SoftmaxPolicy(n_features=4, n_actions=2)
    # ベースライン（状態価値関数）の重み
    w = np.zeros(4)  # V_w(s) = w^T s（線形近似）
    rewards_history = []

    for ep in range(n_episodes):
        state = env.reset()
        episode_states = []
        episode_actions = []
        episode_rewards = []
        done = False

        while not done:
            action, _ = policy.select_action(state)
            next_state, reward, done = env.step(action)

            episode_states.append(state)
            episode_actions.append(action)
            episode_rewards.append(reward)

            state = next_state

        # リターンの計算
        T = len(episode_rewards)
        G = np.zeros(T)
        G[-1] = episode_rewards[-1]
        for t in range(T - 2, -1, -1):
            G[t] = episode_rewards[t] + gamma * G[t + 1]

        # パラメータ更新
        for t in range(T):
            s_t = episode_states[t]
            baseline = w @ s_t  # V_w(s_t)
            advantage = G[t] - baseline

            # Actor（方策）の更新
            grad = policy.compute_grad_log_pi(s_t, episode_actions[t])
            policy.theta += alpha_theta * (gamma ** t) * advantage * grad

            # Critic（ベースライン）の更新
            w += alpha_w * (G[t] - baseline) * s_t

        rewards_history.append(sum(episode_rewards))

    return policy, rewards_history

# --- 実験の実行 ---
np.random.seed(42)
env = CartPole()

# REINFORCEの実行
policy_no_bl, rewards_no_bl = reinforce(env, n_episodes=3000, alpha=0.005)

# REINFORCE with Baselineの実行
np.random.seed(42)
env2 = CartPole()
policy_bl, rewards_bl = reinforce_baseline(env2, n_episodes=3000, alpha_theta=0.005, alpha_w=0.01)

# 結果の可視化
fig, axes = plt.subplots(1, 2, figsize=(16, 6))

# 学習曲線の比較
window = 100
def moving_avg(data, w):
    return np.convolve(data, np.ones(w)/w, mode='valid')

axes[0].plot(moving_avg(rewards_no_bl, window), label='REINFORCE', alpha=0.8)
axes[0].plot(moving_avg(rewards_bl, window), label='REINFORCE + Baseline', alpha=0.8)
axes[0].set_xlabel('Episode', fontsize=12)
axes[0].set_ylabel(f'Average Reward (window={window})', fontsize=12)
axes[0].set_title('REINFORCE: Learning Curves on CartPole', fontsize=14)
axes[0].legend(fontsize=12)
axes[0].grid(True, alpha=0.3)

# エピソード長の推移（生の値）
axes[1].scatter(range(len(rewards_bl)), rewards_bl, s=1, alpha=0.3, label='Episode Length')
axes[1].plot(moving_avg(rewards_bl, window), 'r-', linewidth=2, label=f'Moving Avg (w={window})')
axes[1].set_xlabel('Episode', fontsize=12)
axes[1].set_ylabel('Episode Length', fontsize=12)
axes[1].set_title('REINFORCE + Baseline: Episode Lengths', fontsize=14)
axes[1].legend(fontsize=12)
axes[1].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# 学習後の方策のテスト
def test_policy(env, policy, n_tests=20):
    """学習済み方策のテスト"""
    total_rewards = []
    for _ in range(n_tests):
        state = env.reset()
        total_reward = 0
        done = False
        steps = 0
        while not done and steps < 500:
            action, _ = policy.select_action(state)
            state, reward, done = env.step(action)
            total_reward += reward
            steps += 1
        total_rewards.append(total_reward)
    return total_rewards

test_results = test_policy(CartPole(), policy_bl, n_tests=50)
print(f"学習後のテスト結果（50回平均）: {np.mean(test_results):.1f} ± {np.std(test_results):.1f} ステップ")

import numpy as np
import matplotlib.pyplot as plt

class CartPole:
    """CartPole環境（再掲）"""
    def __init__(self):
        self.gravity = 9.8
        self.masscart = 1.0
        self.masspole = 0.1
        self.total_mass = self.masscart + self.masspole
        self.length = 0.5
        self.polemass_length = self.masspole * self.length
        self.force_mag = 10.0
        self.tau = 0.02
        self.theta_threshold = 12 * np.pi / 180
        self.x_threshold = 2.4

    def reset(self):
        self.state = np.random.uniform(-0.05, 0.05, size=4)
        return self.state.copy()

    def step(self, action):
        x, x_dot, theta, theta_dot = self.state
        force = self.force_mag if action == 1 else -self.force_mag
        cos_theta = np.cos(theta)
        sin_theta = np.sin(theta)
        temp = (force + self.polemass_length * theta_dot**2 * sin_theta) / self.total_mass
        theta_acc = (self.gravity * sin_theta - cos_theta * temp) / (
            self.length * (4.0/3.0 - self.masspole * cos_theta**2 / self.total_mass))
        x_acc = temp - self.polemass_length * theta_acc * cos_theta / self.total_mass
        x += self.tau * x_dot
        x_dot += self.tau * x_acc
        theta += self.tau * theta_dot
        theta_dot += self.tau * theta_acc
        self.state = np.array([x, x_dot, theta, theta_dot])
        done = bool(
            x < -self.x_threshold or x > self.x_threshold
            or theta < -self.theta_threshold or theta > self.theta_threshold
        )
        reward = 1.0 if not done else 0.0
        return self.state.copy(), reward, done

def actor_critic(env, n_episodes=3000, gamma=0.99,
                 alpha_actor=0.005, alpha_critic=0.01):
    """Actor-Criticアルゴリズム（線形近似）"""
    n_features = 4
    n_actions = 2

    # Actor: ソフトマックス方策
    theta = np.zeros((n_features, n_actions))
    # Critic: 線形価値関数 V(s) = w^T s
    w = np.zeros(n_features)

    rewards_history = []

    def softmax(x):
        e_x = np.exp(x - np.max(x))
        return e_x / e_x.sum()

    for ep in range(n_episodes):
        state = env.reset()
        total_reward = 0
        done = False
        I = 1.0  # 割引率の累積

        while not done:
            # 行動の選択
            logits = state @ theta
            probs = softmax(logits)
            action = np.random.choice(n_actions, p=probs)

            # 環境との相互作用
            next_state, reward, done = env.step(action)
            total_reward += reward

            # TD誤差の計算
            v_current = w @ state
            v_next = 0 if done else w @ next_state
            delta = reward + gamma * v_next - v_current

            # Criticの更新
            w += alpha_critic * delta * state

            # Actorの更新
            one_hot = np.zeros(n_actions)
            one_hot[action] = 1.0
            grad_log_pi = np.outer(state, one_hot - probs)
            theta += alpha_actor * I * delta * grad_log_pi

            I *= gamma
            state = next_state

        rewards_history.append(total_reward)

    return theta, w, rewards_history

# Actor-Criticの実行
np.random.seed(42)
env_ac = CartPole()
theta_ac, w_ac, rewards_ac = actor_critic(env_ac, n_episodes=3000)

# 可視化
fig, ax = plt.subplots(figsize=(10, 6))

window = 100
avg_rewards = np.convolve(rewards_ac, np.ones(window)/window, mode='valid')
ax.plot(avg_rewards, linewidth=2, label='Actor-Critic')
ax.set_xlabel('Episode', fontsize=12)
ax.set_ylabel(f'Average Reward (window={window})', fontsize=12)
ax.set_title('Actor-Critic on CartPole', fontsize=14)
ax.legend(fontsize=12)
ax.grid(True, alpha=0.3)
ax.axhline(y=200, color='r', linestyle='--', alpha=0.5, label='Solved threshold')
ax.legend(fontsize=12)
plt.tight_layout()
plt.show()

print(f"最後100エピソードの平均報酬: {np.mean(rewards_ac[-100:]):.1f}")