Policy-based approach — REINFORCE algorithm

English translation

[Policy-based method]

A Q-function is not explicitly used; the policy itself is directly parameterized and optimized
- Policy $\pi_\theta(a|s)$ : probability of choosing action a in state s
Value-based methods directly approximate and learn the state–action value function $Q^*(s,a)$ , selecting the action with the highest Q-value in each state
Policy-based methods directly optimize the policy itself without going through a value function.
Benefits
- Handles continuous and large action spaces easily
  - Direct policy approximation is convenient when approximating a Q-table or Q-function for all actions is difficult.
  - A stochastic policy has built-in exploration.
    - Probabilistic action sampling naturally produces exploration without a separate exploration strategy.
  - Intuitive policy improvement: The goal is to maximize expected policy return, which can be optimized directly.
Disadvantages
- Training may be unstable compared with value-based methods

[REINFORCE algorithm]

Monte Carlo policy gradient method: Update the policy after experiencing an entire episode.
Policy Gradient Theorem
- Estimate the gradient of policy performance and update policy parameters by gradient ascent.
  - Adjust policy parameters in the direction that improves policy performance.
- Formula: $\nabla_\theta J(\theta) = E_\pi \big[\, \nabla_\theta \log \pi_\theta(A_t \mid S_t)\, G_t \,\big]$
  - $G_t$ : From time $t$ , the return
  - Intuition: Update the policy to increase the probability of actions that received high rewards.
    - If an action produces a good outcome, learn to increase the probability of taking that action.
Update equation
- $\theta \leftarrow \theta + \alpha\, G_t \, \nabla_\theta \log \pi_\theta(A_t \mid S_t)$
  - $\alpha$ : Learning rate
  - High rewards lead to an increased probability of the action taken; poor rewards decrease it.
    - For example, if an action earns a high game score, increase its probability in similar situations.
REINFORCE with Advantage
- Advantage function $A_t = G_t - V^\pi(S_t)$ : Replace $G_t$ with $A_t$ .
  - $V^\pi(S_t)$ : State value function
  - The advantage function represents the additional benefit of taking an action in a particular state.
  - Subtract the state value as a baseline, and update the policy based on how much better the return was than expected for that state.

[REINFORCE implementation]

Design the policy network
A neural network that takes a state and outputs a probability distribution over actions
$\pi(a|s)$ : Probability of choosing action $a$ in state $s$
$\psi_a$ : Output of Layer 2 corresponding to action $a$

$\begin{array}{c} \boxed{\text{Input: State }(s)} \\ \downarrow \\ \boxed{\text{Layer 1: Fully Connected (Linear), } \phi = W_1 s + b_1} \\ \downarrow \\ \boxed{\text{Activation Function: ReLU (Rectified Linear Unit), } a = \max(0, \phi)} \\ \downarrow \\ \boxed{\text{Layer 2: Fully Connected (Linear), } \psi = W_2 a + b_2} \\ \downarrow \\ \boxed{\text{Output Activation: Softmax}, \pi(a|s) = \frac{e^{\psi_a}}{\sum_{a'} e^{\psi_{a'}}}} \\ \downarrow \\ \boxed{\text{Output: Action Probabilities } \pi(a|s)\;} \end{array}$

Collect an episode
- Interact with the environment to collect states, actions, rewards, and other data.
Use collected episodes to update the policy
- Update policy network parameters using the policy gradient theorem and update rule. The REINFORCE policy update is expressed as follows. The key is to update policy parameters using the policy gradient.

1. Basic notation

$\theta$ : Policy neural network parameters (weights and biases)
$\pi_\theta(a|s)$ : Policy with parameters $\theta$ (probability of taking action $a$ in state $s$ )
$R_t = \sum_{k=t}^{T}\gamma^{k-t} r_k$ : Discounted total reward from time $t$ to the end of the episode
- $r_k$ : Reward at time $k$
- $\gamma$ : Discount factor ( $0 \le \gamma \le 1$ )
- $T$ : Episode end time
$\nabla_\theta$ : Gradient with respect to $\theta$

2. Policy Gradient (Policy Gradient) calculation

$\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[ \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t|s_t) R_t \right]$

$J(\theta)$ : policy performance metric (expectation)
- In REINFORCE, the expected total reward obtained during an episode
$\mathbb{E}_{\tau \sim \pi_\theta}[\cdot]$ : Expectation over trajectory $\tau = (s_0, a_0, r_1, s_1, a_1, ... , s_T, a_T)$ generated by following policy $\pi_\theta$
$\nabla_\theta \log \pi_\theta(a_t|s_t)$ : Gradient of the log policy\
Indicates the direction that increases the probability of policy $\pi_\theta$ choosing action $a_t$ in state $s_t$
$R_t$ : Return from time t

3. Compute the policy update

After calculating the policy gradient, update the policy parameters using gradient ascent
$\theta \leftarrow \theta + \alpha \nabla_\theta J(\theta)$
- $\alpha$ : Learning rate

4. Monte Carlo approximation

Since the expectation cannot be calculated exactly in practice, approximate it using Monte Carlo sampling\
Sample one episode $\tau$ to estimate the policy gradient as follows
$\nabla_\theta J(\theta) \approx \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t|s_t) R_t$

5. Final update rule

$\theta \leftarrow \theta + \alpha \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t|s_t) R_t$
- Log-policy gradient $\nabla_\theta \log \pi_\theta(a_t|s_t)$
  - Direction that makes the policy select the current action more or less often
- Total reward (return) $R_t$ obtained after $t$
  - Positive for a good action and negative for a bad action
- Learning rate $\alpha$
  - Controls the update size

[Limitations]

High variance: episode-level return calculations have greater variance in long episodes or complex problems, reducing update efficiency
- REINFORCE updates the policy only after an episode ends, making it difficult to use information obtained mid-episode
Slow response: Updates occur only after an episode ends, making real-time responses difficult.
- Immediate adaptation to environmental changes is difficult.
Solution: Actor–critic methods (Next Week)
- Actor–critic methods address REINFORCE's limitations by jointly training an actor that determines the policy and a critic that estimates values.

[Exercise: REINFORCE]

1. Install and import required libraries

import gymnasium as gym
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical
import numpy as np
import matplotlib.pyplot as plt
import imageio

2. Set hyperparameters

learning_rate: Learning rate
gamma: Discount factor
n_episodes: maximum number of episodes
print_interval: progress reporting interval
early_stopping_patience: Early-stopping patience (how long to wait)
target_score: Target score (training succeeds when this score is exceeded)

learning_rate = 0.0003
gamma = 0.99
n_episodes = 3000  # Maximum number of episodes
print_interval = 50 # Intermediate result reporting interval
early_stopping_patience = 10 # Number of consecutive episodes exceeding target_score
target_score = 400  # CartPole-v1 success criterion (typically a score of 200 or higher)

3. Define the neural network model (policy network)

A simple 2-layer neural network (input–hidden–output)
forward(): return a probability distribution with softmax

class Policy(nn.Module):
    def __init__(self, state_size, action_size, hidden_size=32):
        super(Policy, self).__init__()
        self.fc1 = nn.Linear(state_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, action_size)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.softmax(x, dim=0)  # Probabilities for a single state

4. Implement REINFORCE (train function)

optimizer: Use the Adam optimizer.
For each episode:
- rewards, log_probs: lists storing rewards and log probabilities
- Until done is True:
  - Obtain a probability distribution with policy(state), sample an action with Categorical, and calculate its log probability
  - Use env.step() to obtain the action, reward, termination status, etc.
  - Store in rewards and log_probs
  - Update state
- At the end of the episode:
  - R: calculate return (discounted total reward)
  - policy_loss: sum of -log_prob * R (basic REINFORCE loss)
  - Update the policy with optimizer (gradient ascent)
- Print the mean score every print_interval and check the early stopping condition

def train(env, policy, optimizer, n_episodes, gamma, print_interval, early_stopping_patience, target_score):
    scores = []

    for i_episode in range(1, n_episodes + 1):
        state, _ = env.reset()
        if isinstance(state, tuple):
            state = state[0]  # gym version compatibility
        state = torch.from_numpy(state).float()

        rewards = []
        log_probs = []
        done = False

        while not done:
            probs = policy(state)
            m = Categorical(probs)
            action = m.sample()
            log_prob = m.log_prob(action)
            next_state, reward, terminated, truncated, _ = env.step(action.item())
            rewards.append(reward)
            log_probs.append(log_prob)
            state = torch.from_numpy(next_state).float()
            done = terminated or truncated

        # Compute returns
        R = 0
        returns = []
        for r in reversed(rewards):
            R = r + gamma * R
            returns.insert(0, R)

        # Compute the Policy Loss (basic form)
        returns = torch.tensor(returns) # Convert a list to a tensor
        log_probs = torch.stack(log_probs) # Convert a list of tensor objects to a tensor
        policy_loss = torch.sum(-log_probs * returns)

        optimizer.zero_grad()
        policy_loss.backward()
        optimizer.step()

        total_reward = sum(rewards)
        scores.append(total_reward)

        if (i_episode+1) % print_interval == 0:
            avg_score = np.mean(scores[-print_interval:])
            print(f"Episode {(i_episode+1)}\tAverage Score: {avg_score:.2f}")

        if np.min(scores[-early_stopping_patience:])>=target_score:
            print(f"Early stopping after {i_episode+1} episodes.  Target score reached.")
            break

    return scores

5. Run training and check results

env = gym.make('CartPole-v1', render_mode = 'rgb_array')
state_size = env.observation_space.shape[0]
action_size = env.action_space.n

policy = Policy(state_size, action_size)
optimizer = optim.Adam(policy.parameters(), lr=learning_rate)

scores = train(env, policy, optimizer, n_episodes, gamma, print_interval, early_stopping_patience, target_score)

plt.plot(scores)
plt.xlabel("Episode")
plt.ylabel("Score")
plt.show()

Episode 50	Average Score: 29.59
Episode 100	Average Score: 29.46
Episode 150	Average Score: 22.94
Episode 200	Average Score: 28.10
Episode 250	Average Score: 33.60
Episode 300	Average Score: 31.22
Episode 350	Average Score: 35.42
Episode 400	Average Score: 38.36
Episode 450	Average Score: 33.10
Episode 500	Average Score: 42.88
Episode 550	Average Score: 46.48
Episode 600	Average Score: 55.06
Episode 650	Average Score: 52.90
Episode 700	Average Score: 56.98
Episode 750	Average Score: 60.50
Episode 800	Average Score: 49.84
Episode 850	Average Score: 61.66
Episode 900	Average Score: 68.46
Episode 950	Average Score: 70.76
Episode 1000	Average Score: 87.88
Episode 1050	Average Score: 95.34
Episode 1100	Average Score: 104.16
Episode 1150	Average Score: 112.28
Episode 1200	Average Score: 161.80
Episode 1250	Average Score: 148.26
Episode 1300	Average Score: 151.28
Episode 1350	Average Score: 194.38
Episode 1400	Average Score: 209.04
Episode 1450	Average Score: 223.46
Episode 1500	Average Score: 213.40
Episode 1550	Average Score: 205.94
Episode 1600	Average Score: 242.70
Episode 1650	Average Score: 268.14
Episode 1700	Average Score: 300.26
Episode 1750	Average Score: 250.06
Episode 1800	Average Score: 279.60
Episode 1850	Average Score: 279.02
Episode 1900	Average Score: 287.54
Episode 1950	Average Score: 301.08
Episode 2000	Average Score: 335.34
Episode 2050	Average Score: 351.32
Episode 2100	Average Score: 341.72
Episode 2150	Average Score: 298.76
Episode 2200	Average Score: 313.38
Episode 2250	Average Score: 331.06
Episode 2300	Average Score: 368.56
Episode 2350	Average Score: 353.54
Episode 2400	Average Score: 340.40
Episode 2450	Average Score: 376.02
Early stopping after 2482 episodes.  Target score reached.

<Figure size 640x480 with 1 Axes>

state, _ = env.reset()
if isinstance(state, tuple):
    state = state[0]
state = torch.from_numpy(state).float()

frames = []
total_reward = 0
done = False

while not done:
    probs = policy(state)
    m = Categorical(probs)
    action = m.sample()

    next_state, reward, terminated, truncated, _ = env.step(action.item())
    total_reward += reward
    state = torch.from_numpy(next_state).float()
    done = terminated or truncated
    frames.append(env.render())

imageio.mimsave("cartpole_simulation_reinforce.gif", frames, duration=33)  # Adjust duration
print(f"Total reward: {total_reward}")
env.close()

Total reward: 500.0